Esempio n. 1
0
    def generate_mmodes(self, ts_data=None):
        """Calculate the m-modes corresponding to the Timestream.

        Perform an MPI transpose for efficiency.
        """

        completed_file = self._mdir + 'COMPLETED_M'
        if os.path.exists(completed_file):
            if mpiutil.rank0:
                print "******* m-files already generated ********"
            mpiutil.barrier()
            return

        # Make directory if required
        # if mpiutil.rank0 and not os.path.exists(self._mdir):
        #     os.makedirs(self._mdir)

        try:
            os.makedirs(self._mdir)
        except OSError:
            # directory exists
            pass

        tel = self.telescope
        mmax = tel.mmax
        ntime = ts_data.shape[0] if ts_data is not None else self.ntime
        nbl = tel.nbase
        nfreq = tel.nfreq

        indices = list(itertools.product(np.arange(nfreq), np.arange(nbl)))
        lind, sind, eind = mpiutil.split_local(nfreq * nbl)

        # load the local section of the time stream
        tstream = np.zeros((ntime, lind), dtype=np.complex128)
        for ind, (f_ind, bl_ind) in enumerate(indices[sind:eind]):
            if ts_data is not None:
                tstream[:, ind] = ts_data[:, f_ind, bl_ind]
            else:
                with h5py.File(self._tsfile, 'r') as f:
                    tstream[:, ind] = f['/timestream'][:, f_ind, bl_ind]

        # FFT to get m-mode
        mmodes = np.fft.fft(tstream, axis=0) / ntime # m = 0 is at left
        mmodes = MPIArray.wrap(mmodes, axis=1)
        # redistribute along different m
        mmodes = mmodes.redistribute(axis=0)

        # save m-modes to file
        ms = np.concatenate([np.arange(0, mmax+1), np.arange(-mmax, 0)])
        for ind, mi in enumerate(mpiutil.mpilist(ms, method='con')):
            with h5py.File(self._mfile(mi), 'w') as f:
              f.create_dataset('/mmode', data=mmodes[ind].view(np.ndarray).reshape(nfreq, nbl))
              f.attrs['m'] = mi

        mpiutil.barrier()

        if mpiutil.rank0:

            # Make file marker that the m's have been correctly generated:
            open(completed_file, 'a').close()
Esempio n. 2
0
    def _generate_invbeam(self, regen=False):

        completed_file = self._inv_mdir + 'COMPLETED_IBEAM'
        if os.path.exists(completed_file) and not regen:
            if mpiutil.rank0:
                print
                print '=' * 80
                print "******* Inverse beam transfer m-files already generated ********"
            mpiutil.barrier()
            return

        if mpiutil.rank0:
            print
            print '=' * 80
            print 'Create inverse beam transfer m-files...'

        st = time.time()

        mmax = self.telescope.mmax
        for mi in mpiutil.mpilist(range(mmax+1)):
            inv_beam, W1 = self.compute_invbeam_m(mi)
            # save to file
            with h5py.File(self._inv_mfile(mi), 'w') as f:
                f.create_dataset('ibeam_m', data=inv_beam)
                f.create_dataset('W1_m', data=W1)
                f.attrs['m'] = mi

        mpiutil.barrier()

        et = time.time()

        if mpiutil.rank0:

            # Make file marker that the m's have been correctly generated:
            open(completed_file, 'a').close()

            # Print out timing
            print "=== Create inverse beam transfer m-files took %f s ===" % (et - st)
Esempio n. 3
0
    def process(self, ts):

        assert isinstance(
            ts, Timestream
        ), '%s only works for Timestream object' % self.__class__.__name__

        calibrator = self.params['calibrator']
        catalog = self.params['catalog']
        vis_conj = self.params['vis_conj']
        zero_diag = self.params['zero_diag']
        span = self.params['span']
        reserve_high_gain = self.params['reserve_high_gain']
        plot_figs = self.params['plot_figs']
        fig_prefix = self.params['fig_name']
        tag_output_iter = self.params['tag_output_iter']
        save_src_vis = self.params['save_src_vis']
        src_vis_file = self.params['src_vis_file']
        subtract_src = self.params['subtract_src']
        replace_with_src = self.params['replace_with_src']
        apply_gain = self.params['apply_gain']
        save_gain = self.params['save_gain']
        save_phs_change = self.params['save_phs_change']
        gain_file = self.params['gain_file']
        # temperature_convert = self.params['temperature_convert']
        show_progress = self.params['show_progress']
        progress_step = self.params['progress_step']

        if save_src_vis or subtract_src or apply_gain or save_gain:
            pol_type = ts['pol'].attrs['pol_type']
            if pol_type != 'linear':
                raise RuntimeError('Can not do ps_cal for pol_type: %s' %
                                   pol_type)

            ts.redistribute('baseline')

            feedno = ts['feedno'][:].tolist()
            pol = [ts.pol_dict[p] for p in ts['pol'][:]]  # as string
            gain_pd = {
                'xx': 0,
                'yy': 1,
                0: 'xx',
                1: 'yy'
            }  # for gain related op
            bls = mpiutil.gather_array(ts.local_bl[:], root=None, comm=ts.comm)
            # # antpointing = np.radians(ts['antpointing'][-1, :, :]) # radians
            # transitsource = ts['transitsource'][:]
            # transit_time = transitsource[-1, 0] # second, sec1970
            # int_time = ts.attrs['inttime'] # second

            # get the calibrator
            try:
                s = calibrators.get_src(calibrator)
            except KeyError:
                if mpiutil.rank0:
                    print 'Calibrator %s is unavailable, available calibrators are:'
                    for key, d in calibrators.src_data.items():
                        print '%8s  ->  %12s' % (key, d[0])
                raise RuntimeError('Calibrator %s is unavailable')
            if mpiutil.rank0:
                print 'Try to calibrate with %s...' % s.src_name

            # get transit time of calibrator
            # array
            aa = ts.array
            aa.set_jultime(ts['jul_date'][0])  # the first obs time point
            next_transit = aa.next_transit(s)
            transit_time = a.phs.ephem2juldate(next_transit)  # Julian date
            # get time zone
            pattern = '[-+]?\d+'
            tz = re.search(pattern, ts.attrs['timezone']).group()
            tz = int(tz)
            local_next_transit = ephem.Date(
                next_transit + tz * ephem.hour)  # plus 8h to get Beijing time
            # if transit_time > ts['jul_date'][-1]:
            if transit_time > max(ts['jul_date'][-1], ts['jul_date'][:].max()):
                raise RuntimeError(
                    'Data does not contain local transit time %s of source %s'
                    % (local_next_transit, calibrator))

            # the first transit index
            transit_inds = [np.searchsorted(ts['jul_date'][:], transit_time)]
            # find all other transit indices
            aa.set_jultime(ts['jul_date'][0] + 1.0)
            transit_time = a.phs.ephem2juldate(
                aa.next_transit(s))  # Julian date
            cnt = 2
            while (transit_time <= ts['jul_date'][-1]):
                transit_inds.append(
                    np.searchsorted(ts['jul_date'][:], transit_time))
                aa.set_jultime(ts['jul_date'][0] + 1.0 * cnt)
                transit_time = a.phs.ephem2juldate(
                    aa.next_transit(s))  # Julian date
                cnt += 1

            if mpiutil.rank0:
                print 'transit ind of %s: %s, time: %s' % (
                    s.src_name, transit_inds, local_next_transit)

            ### now only use the first transit point to do the cal
            ### may need to improve in the future
            transit_ind = transit_inds[0]
            int_time = ts.attrs['inttime']  # second
            start_ind = transit_ind - np.int(span / int_time)
            end_ind = transit_ind + np.int(
                span /
                int_time) + 1  # plus 1 to make transit_ind is at the center

            start_ind = max(0, start_ind)
            end_ind = min(end_ind, ts.vis.shape[0])

            if vis_conj:
                ts.local_vis[:] = ts.local_vis.conj()

            nt = end_ind - start_ind
            t_inds = range(start_ind, end_ind)
            freq = ts.freq[:]  # MHz
            nf = len(freq)
            nlb = len(ts.local_bl[:])
            nfeed = len(feedno)
            tfp_inds = list(
                itertools.product(
                    t_inds, range(nf),
                    [pol.index('xx'), pol.index('yy')]))  # only for xx and yy
            ns, ss, es = mpiutil.split_all(len(tfp_inds), comm=ts.comm)
            # gather data to make each process to have its own data which has all bls
            for ri, (ni, si, ei) in enumerate(zip(ns, ss, es)):
                lvis = np.zeros((ni, nlb), dtype=ts.vis.dtype)
                lvis_mask = np.zeros((ni, nlb), dtype=ts.vis_mask.dtype)
                for ii, (ti, fi, pi) in enumerate(tfp_inds[si:ei]):
                    lvis[ii] = ts.local_vis[ti, fi, pi]
                    lvis_mask[ii] = ts.local_vis_mask[ti, fi, pi]
                # gather vis from all process for separate bls
                gvis = mpiutil.gather_array(lvis,
                                            axis=1,
                                            root=ri,
                                            comm=ts.comm)
                gvis_mask = mpiutil.gather_array(lvis_mask,
                                                 axis=1,
                                                 root=ri,
                                                 comm=ts.comm)
                if ri == mpiutil.rank:
                    tfp_linds = tfp_inds[si:ei]  # inds for this process
                    this_vis = gvis
                    this_vis_mask = gvis_mask
            del tfp_inds
            del lvis
            del lvis_mask
            tfp_len = len(tfp_linds)

            # lotl_mask = np.zeros((tfp_len, nfeed, nfeed), dtype=bool)
            cnan = complex(np.nan, np.nan)  # complex nan
            if save_src_vis or subtract_src:
                # save calibrator src vis
                lsrc_vis = np.full((tfp_len, nfeed, nfeed),
                                   cnan,
                                   dtype=ts.vis.dtype)
                if save_src_vis:
                    # save sky vis
                    lsky_vis = np.full((tfp_len, nfeed, nfeed),
                                       cnan,
                                       dtype=ts.vis.dtype)
                    # save outlier vis
                    lotl_vis = np.full((tfp_len, nfeed, nfeed),
                                       cnan,
                                       dtype=ts.vis.dtype)

            if apply_gain or save_gain:
                lGain = np.full((tfp_len, nfeed), cnan, dtype=ts.vis.dtype)

            # find indices mapping between Vmat and vis
            # bis = range(nbl)
            bis_conj = []  # indices that shold be conj
            mis = [
            ]  # indices in the nfeed x nfeed matrix by flatten it to a vector
            mis_conj = [
            ]  # indices (of conj vis) in the nfeed x nfeed matrix by flatten it to a vector
            for bi, (fdi, fdj) in enumerate(bls):
                ai, aj = feedno.index(fdi), feedno.index(fdj)
                mis.append(ai * nfeed + aj)
                if ai != aj:
                    bis_conj.append(bi)
                    mis_conj.append(aj * nfeed + ai)

            # construct visibility matrix for a single time, freq, pol
            Vmat = np.full((nfeed, nfeed), cnan, dtype=ts.vis.dtype)
            # get flus of the calibrator in the observing frequencies
            if show_progress and mpiutil.rank0:
                pg = progress.Progress(tfp_len, step=progress_step)
            for ii, (ti, fi, pi) in enumerate(tfp_linds):
                if show_progress and mpiutil.rank0:
                    pg.show(ii)
                # when noise on, just pass
                if 'ns_on' in ts.iterkeys() and ts['ns_on'][ti]:
                    continue
                # aa.set_jultime(ts['jul_date'][ti])
                # s.compute(aa)
                # get the topocentric coordinate of the calibrator at the current time
                # s_top = s.get_crds('top', ncrd=3)
                # aa.sim_cache(cat.get_crds('eq', ncrd=3)) # for compute bm_response and sim
                Vmat.flat[mis] = np.ma.array(
                    this_vis[ii], mask=this_vis_mask[ii]).filled(cnan)
                Vmat.flat[mis_conj] = np.ma.array(
                    this_vis[ii, bis_conj],
                    mask=this_vis_mask[ii, bis_conj]).conj().filled(cnan)

                if save_src_vis:
                    lsky_vis[ii] = Vmat

                # set invalid val to 0
                invalid = ~np.isfinite(Vmat)  # a bool array
                # if too many masks
                if np.where(invalid)[0].shape[0] > 0.3 * nfeed**2:
                    continue
                Vmat[invalid] = 0
                # if all are zeros
                if np.allclose(Vmat, 0.0):
                    continue

                # fill diagonal of Vmat to 0
                if zero_diag:
                    np.fill_diagonal(Vmat, 0)

                # initialize the outliers
                med = np.median(Vmat.real) + 1.0J * np.median(Vmat.imag)
                diff = Vmat - med
                S0 = np.where(
                    np.abs(diff) > 3.0 * rpca_decomp.MAD(Vmat), diff, 0)
                # stable PCA decomposition
                V0, S = rpca_decomp.decompose(Vmat,
                                              rank=1,
                                              S=S0,
                                              max_iter=200,
                                              threshold='hard',
                                              tol=1.0e-6,
                                              debug=False)

                # # find abnormal values in S
                # # first check diagonal elements
                # import pdb; pdb.set_trace()
                # svals = np.diag(S)
                # smed = np.median(svals.real) + 1.0J * np.median(svals.imag)
                # smad = rpca_decomp.MAD(svals)
                # # abnormal indices
                # abis =  np.where(np.abs(svals - smed) > 3.0 * smad)[0]
                # for abi in abis:
                #     lotl_mask[ii, abi, abi] = True
                # # then check non-diagonal elements
                # for rii in range(nfeed):
                #     for cii in range(nfeed):
                #         if rii == cii:
                #             continue
                #         rli = max(0, rii-2)
                #         rhi = min(nfeed, rii+3)
                #         cli = max(0, cii-2)
                #         chi = min(nfeed, cii+3)
                #         svals = np.array([ S[xi, yi] for xi in range(rli, rhi) for yi in range(cli, chi) if xi != yi ])
                #         smed = np.median(svals.real) + 1.0J * np.median(svals.imag)
                #         smad = rpca_decomp.MAD(svals)
                #         if np.abs(S[rii, cii] - smed) > 3.0 * smad:
                #             lotl_mask[ii, rii, cii] = True

                if save_src_vis or subtract_src:
                    lsrc_vis[ii] = V0
                    if save_src_vis:
                        lotl_vis[ii] = S

                # plot
                if plot_figs:
                    ind = ti - start_ind
                    # plot Vmat
                    plt.figure(figsize=(13, 5))
                    plt.subplot(121)
                    plt.imshow(Vmat.real,
                               aspect='equal',
                               origin='lower',
                               interpolation='nearest')
                    plt.colorbar(shrink=1.0)
                    plt.subplot(122)
                    plt.imshow(Vmat.imag,
                               aspect='equal',
                               origin='lower',
                               interpolation='nearest')
                    plt.colorbar(shrink=1.0)
                    fig_name = '%s_V_%d_%d_%s.png' % (fig_prefix, ind, fi,
                                                      pol[pi])
                    if tag_output_iter:
                        fig_name = output_path(fig_name,
                                               iteration=self.iteration)
                    else:
                        fig_name = output_path(fig_name)
                    plt.savefig(fig_name)
                    plt.close()
                    # plot V0
                    plt.figure(figsize=(13, 5))
                    plt.subplot(121)
                    plt.imshow(V0.real,
                               aspect='equal',
                               origin='lower',
                               interpolation='nearest')
                    plt.colorbar(shrink=1.0)
                    plt.subplot(122)
                    plt.imshow(V0.imag,
                               aspect='equal',
                               origin='lower',
                               interpolation='nearest')
                    plt.colorbar(shrink=1.0)
                    fig_name = '%s_V0_%d_%d_%s.png' % (fig_prefix, ind, fi,
                                                       pol[pi])
                    if tag_output_iter:
                        fig_name = output_path(fig_name,
                                               iteration=self.iteration)
                    else:
                        fig_name = output_path(fig_name)
                    plt.savefig(fig_name)
                    plt.close()
                    # plot S
                    plt.figure(figsize=(13, 5))
                    plt.subplot(121)
                    plt.imshow(S.real,
                               aspect='equal',
                               origin='lower',
                               interpolation='nearest')
                    plt.colorbar(shrink=1.0)
                    plt.subplot(122)
                    plt.imshow(S.imag,
                               aspect='equal',
                               origin='lower',
                               interpolation='nearest')
                    plt.colorbar(shrink=1.0)
                    fig_name = '%s_S_%d_%d_%s.png' % (fig_prefix, ind, fi,
                                                      pol[pi])
                    if tag_output_iter:
                        fig_name = output_path(fig_name,
                                               iteration=self.iteration)
                    else:
                        fig_name = output_path(fig_name)
                    plt.savefig(fig_name)
                    plt.close()
                    # plot N
                    N = Vmat - V0 - S
                    plt.figure(figsize=(13, 5))
                    plt.subplot(121)
                    plt.imshow(N.real,
                               aspect='equal',
                               origin='lower',
                               interpolation='nearest')
                    plt.colorbar(shrink=1.0)
                    plt.subplot(122)
                    plt.imshow(N.imag,
                               aspect='equal',
                               origin='lower',
                               interpolation='nearest')
                    plt.colorbar(shrink=1.0)
                    fig_name = '%s_N_%d_%d_%s.png' % (fig_prefix, ind, fi,
                                                      pol[pi])
                    if tag_output_iter:
                        fig_name = output_path(fig_name,
                                               iteration=self.iteration)
                    else:
                        fig_name = output_path(fig_name)
                    plt.savefig(fig_name)
                    plt.close()

                if apply_gain or save_gain:
                    # use v_ij = gi gj^* \int Ai Aj^* e^(2\pi i n \cdot uij) T(x) d^2n
                    # precisely, we shold have
                    # V0 = (lambda^2 * Sc / (2 k_B)) * gi gj^* Ai Aj^* e^(2\pi i n0 \cdot uij)
                    e, U = la.eigh(V0, eigvals=(nfeed - 1, nfeed - 1))
                    g = U[:, -1] * e[
                        -1]**0.5  # = \sqrt(lambda^2 * Sc / (2 k_B)) * gi Ai * e^(2\pi i n0 \cdot ui)
                    if g[0].real < 0:
                        g *= -1.0  # make all g[0] phase 0, instead of pi
                    lGain[ii] = g

                    # plot Gain
                    if plot_figs:
                        plt.figure()
                        plt.plot(feedno, g.real, 'b-', label='real')
                        plt.plot(feedno, g.real, 'bo')
                        plt.plot(feedno, g.imag, 'g-', label='imag')
                        plt.plot(feedno, g.imag, 'go')
                        plt.plot(feedno, np.abs(g), 'r-', label='abs')
                        plt.plot(feedno, np.abs(g), 'ro')
                        plt.xlim(feedno[0] - 1, feedno[-1] + 1)
                        yl, yh = plt.ylim()
                        plt.ylim(yl, yh + (yh - yl) / 5)
                        plt.xlabel('Feed number')
                        plt.legend()
                        fig_name = '%s_ants_%d_%d_%s.png' % (fig_prefix, ind,
                                                             fi, pol[pi])
                        if tag_output_iter:
                            fig_name = output_path(fig_name,
                                                   iteration=self.iteration)
                        else:
                            fig_name = output_path(fig_name)
                        plt.savefig(fig_name)
                        plt.close()

            # # apply outlier mask
            # nbl = len(bls)
            # lom = np.zeros((lotl_mask.shape[0], nbl), dtype=lotl_mask.dtype)
            # for bi, (fd1, fd2) in enumerate(bls):
            #     b1, b2 = feedno.index(fd1), feedno.index(fd2)
            #     lom[:, bi] = lotl_mask[:, b1, b2]
            # lom = mpiarray.MPIArray.wrap(lom, axis=0, comm=ts.comm)
            # lom = lom.redistribute(axis=1).local_array.reshape(nt, nf, 2, -1)
            # ts.local_vis_mask[start_ind:end_ind, :, pol.index('xx')] |= lom[:, :, 0]
            # ts.local_vis_mask[start_ind:end_ind, :, pol.index('yy')] |= lom[:, :, 1]

            # subtract the vis of calibrator from self.vis
            if subtract_src:
                nbl = len(bls)
                lv = np.zeros((lsrc_vis.shape[0], nbl), dtype=lsrc_vis.dtype)
                for bi, (fd1, fd2) in enumerate(bls):
                    b1, b2 = feedno.index(fd1), feedno.index(fd2)
                    lv[:, bi] = lsrc_vis[:, b1, b2]
                lv = mpiarray.MPIArray.wrap(lv, axis=0, comm=ts.comm)
                lv = lv.redistribute(axis=1).local_array.reshape(nt, nf, 2, -1)
                if replace_with_src:
                    ts.local_vis[start_ind:end_ind, :,
                                 pol.index('xx')] = lv[:, :, 0]
                    ts.local_vis[start_ind:end_ind, :,
                                 pol.index('yy')] = lv[:, :, 1]
                else:
                    if 'ns_on' in ts.iterkeys():
                        lv[ts['ns_on']
                           [start_ind:
                            end_ind]] = 0  # avoid ns_on signal to become nan
                    ts.local_vis[start_ind:end_ind, :,
                                 pol.index('xx')] -= lv[:, :, 0]
                    ts.local_vis[start_ind:end_ind, :,
                                 pol.index('yy')] -= lv[:, :, 1]

                del lv

            if not save_src_vis:
                if subtract_src:
                    del lsrc_vis
            else:
                if tag_output_iter:
                    src_vis_file = output_path(src_vis_file,
                                               iteration=self.iteration)
                else:
                    src_vis_file = output_path(src_vis_file)
                # create file and allocate space first by rank0
                if mpiutil.rank0:
                    with h5py.File(src_vis_file, 'w') as f:
                        # allocate space
                        shp = (nt, nf, 2, nfeed, nfeed)
                        f.create_dataset('sky_vis', shp, dtype=lsky_vis.dtype)
                        f.create_dataset('src_vis', shp, dtype=lsrc_vis.dtype)
                        f.create_dataset('outlier_vis',
                                         shp,
                                         dtype=lotl_vis.dtype)
                        # f.create_dataset('outlier_mask', shp, dtype=lotl_mask.dtype)
                        f.attrs['calibrator'] = calibrator
                        f.attrs['dim'] = 'time, freq, pol, feed, feed'
                        try:
                            f.attrs['time'] = ts.time[start_ind:end_ind]
                        except RuntimeError:
                            f.create_dataset('time',
                                             data=ts.time[start_ind:end_ind])
                            f.attrs['time'] = '/time'
                        f.attrs['freq'] = freq
                        f.attrs['pol'] = np.array(['xx', 'yy'])
                        f.attrs['feed'] = np.array(feedno)

                mpiutil.barrier()

                # write data to file
                for i in range(10):
                    try:
                        # NOTE: if write simultaneously, will loss data with processes distributed in several nodes
                        for ri in xrange(mpiutil.size):
                            if ri == mpiutil.rank:
                                with h5py.File(src_vis_file, 'r+') as f:
                                    for ii, (ti, fi,
                                             pi) in enumerate(tfp_linds):
                                        ti_ = ti - start_ind
                                        pi_ = gain_pd[pol[pi]]
                                        f['sky_vis'][ti_, fi,
                                                     pi_] = lsky_vis[ii]
                                        f['src_vis'][ti_, fi,
                                                     pi_] = lsrc_vis[ii]
                                        f['outlier_vis'][ti_, fi,
                                                         pi_] = lotl_vis[ii]
                                        # f['outlier_mask'][ti_, fi, pi_] = lotl_mask[ii]
                            mpiutil.barrier()
                        break
                    except IOError:
                        time.sleep(0.5)
                        continue
                else:
                    raise RuntimeError('Could not open file: %s...' %
                                       src_vis_file)

                del lsrc_vis
                del lsky_vis
                del lotl_vis
                # del lotl_mask

                mpiutil.barrier()

            if apply_gain or save_gain:
                # flag outliers in lGain along each feed
                lG_abs = np.full_like(lGain, np.nan, dtype=lGain.real.dtype)
                for i in range(lGain.shape[0]):
                    valid_inds = np.where(np.isfinite(lGain[i]))[0]
                    if len(valid_inds) > 3:
                        vabs = np.abs(lGain[i, valid_inds])
                        vmed = np.median(vabs)
                        vabs_diff = np.abs(vabs - vmed)
                        vmad = np.median(vabs_diff) / 0.6745
                        if reserve_high_gain:
                            # reserve significantly higher ones, flag only significantly lower ones
                            lG_abs[i, valid_inds] = np.where(
                                vmed - vabs > 3.0 * vmad, np.nan, vabs)
                        else:
                            # flag both significantly higher and lower ones
                            lG_abs[i, valid_inds] = np.where(
                                vabs_diff > 3.0 * vmad, np.nan, vabs)

                # choose data slice near the transit time
                li = max(start_ind, transit_ind - 10) - start_ind
                hi = min(end_ind, transit_ind + 10 + 1) - start_ind
                ci = transit_ind - start_ind  # center index for transit_ind
                # compute s_top for this time range
                n0 = np.zeros(((hi - li), 3))
                for ti, jt in enumerate(ts.time[start_ind:end_ind][li:hi]):
                    aa.set_jultime(jt)
                    s.compute(aa)
                    n0[ti] = s.get_crds('top', ncrd=3)
                if save_phs_change:
                    n0t = np.zeros((nt, 3))
                    for ti, jt in enumerate(ts.time[start_ind:end_ind]):
                        aa.set_jultime(jt)
                        s.compute(aa)
                        n0t[ti] = s.get_crds('top', ncrd=3)

                # get the positions of feeds
                feedpos = ts['feedpos'][:]

                # wrap and redistribute Gain and flagged G_abs
                Gain = mpiarray.MPIArray.wrap(lGain, axis=0, comm=ts.comm)
                Gain = Gain.redistribute(axis=1).reshape(
                    nt, nf, 2, None).redistribute(axis=0).reshape(
                        None, nf * 2 * nfeed).redistribute(axis=1)
                G_abs = mpiarray.MPIArray.wrap(lG_abs, axis=0, comm=ts.comm)
                G_abs = G_abs.redistribute(axis=1).reshape(
                    nt, nf, 2, None).redistribute(axis=0).reshape(
                        None, nf * 2 * nfeed).redistribute(axis=1)

                fpd_inds = list(
                    itertools.product(range(nf), range(2),
                                      range(nfeed)))  # only for xx and yy
                fpd_linds = mpiutil.mpilist(fpd_inds,
                                            method='con',
                                            comm=ts.comm)
                del fpd_inds
                # create data to save the solved gain for each feed
                lgain = np.full((len(fpd_linds), ), cnan,
                                dtype=Gain.dtype)  # gain for each feed
                if save_phs_change:
                    lphs = np.full((nt, len(fpd_linds)),
                                   np.nan,
                                   dtype=Gain.real.dtype
                                   )  # phase change with time for each feed

                # check for conj
                num_conj = 0
                for ii, (fi, pi, di) in enumerate(fpd_linds):
                    y = G_abs.local_array[li:hi, ii]
                    inds = np.where(np.isfinite(y))[0]
                    if len(inds) >= max(4, 0.5 * len(y)):
                        # get the approximate magnitude by averaging the central G_abs
                        # solve phase by least square fit
                        ui = (feedpos[di] - feedpos[0]) * (
                            1.0e6 * freq[fi]
                        ) / const.c  # position of this feed (relative to the first feed) in unit of wavelength
                        exp_factor = np.exp(2.0J * np.pi * np.dot(n0, ui))
                        ef = exp_factor
                        Gi = Gain.local_array[li:hi, ii]
                        e_phs = np.dot(ef[inds].conj(),
                                       Gi[inds] / y[inds]) / len(inds)
                        ea = np.abs(e_phs)
                        e_phs_conj = np.dot(ef[inds],
                                            Gi[inds] / y[inds]) / len(inds)
                        eac = np.abs(e_phs_conj)
                        if eac > ea:
                            num_conj += 1
                # reduce num_conj from all processes
                num_conj = mpiutil.allreduce(num_conj, comm=ts.comm)
                if num_conj > 0.5 * (nf * 2 * nfeed):  # 2 for 2 pols
                    if mpiutil.rank0:
                        print '!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!'
                        print '!!!   Detect data should be their conjugate...   !!!'
                        print '!!!   Correct it automatically...                !!!'
                        print '!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!'
                    mpiutil.barrier()
                    # correct vis
                    ts.local_vis[:] = ts.local_vis.conj()
                    # correct G
                    Gain.local_array[:] = Gain.local_array.conj()

                # solve for gain
                for ii, (fi, pi, di) in enumerate(fpd_linds):
                    y = G_abs.local_array[li:hi, ii]
                    inds = np.where(np.isfinite(y))[0]
                    if len(inds) >= max(4, 0.5 * len(y)):
                        # get the approximate magnitude by averaging the central G_abs
                        mag = np.mean(
                            y[inds]
                        )  # = \sqrt(lambda^2 * Sc / (2 k_B)) * |gi| Ai
                        # solve phase by least square fit
                        ui = (feedpos[di] - feedpos[0]) * (
                            1.0e6 * freq[fi]
                        ) / const.c  # position of this feed (relative to the first feed) in unit of wavelength
                        exp_factor = np.exp(2.0J * np.pi * np.dot(n0, ui))
                        ef = exp_factor
                        Gi = Gain.local_array[li:hi, ii]
                        e_phs = np.dot(ef[inds].conj(), Gi[inds] /
                                       y[inds]) / len(inds)  # the phase of gi
                        ea = np.abs(e_phs)
                        if np.abs(ea - 1.0) < 0.1:
                            # compute gain for this feed
                            lgain[
                                ii] = mag * e_phs  # \sqrt(lambda^2 * Sc / (2 k_B)) * gi Ai
                            if save_phs_change:
                                lphs[:, ii] = np.angle(
                                    np.exp(-2.0J * np.pi * np.dot(n0t, ui)) *
                                    Gain.local_array[:, ii])
                        else:
                            e_phs_conj = np.dot(ef[inds],
                                                Gi[inds] / y[inds]) / len(inds)
                            eac = np.abs(e_phs_conj)
                            if eac > ea:
                                if np.abs(eac - 1.0) < 0.01:
                                    print 'feedno = %d, fi = %d, pol = %s: may need to be conjugated' % (
                                        feedno[di], fi, gain_pd[pi])
                            else:
                                print 'feedno = %d, fi = %d, pol = %s: maybe wrong abs(e_phs): %s' % (
                                    feedno[di], fi, gain_pd[pi], ea)

                # gather local gain
                gain = mpiutil.gather_array(lgain,
                                            axis=0,
                                            root=None,
                                            comm=ts.comm)
                del lgain
                gain = gain.reshape(nf, 2, nfeed)
                if save_phs_change:
                    phs = mpiutil.gather_array(lphs,
                                               axis=1,
                                               root=0,
                                               comm=ts.comm)
                    del lphs
                    if mpiutil.rank0:
                        phs = phs.reshape(nt, nf, 2, nfeed)

                # normalize to get the exact gain
                Sc = s.get_jys(1.0e-3 * freq)
                # Omega = aa.ants[0].beam.Omega ### TODO: implement Omega for dish
                Ai = aa.ants[0].beam.response(n0[ci - li])
                lmd = const.c / (1.0e6 * freq)
                factor = np.sqrt(
                    (lmd**2 * 1.0e-26 * Sc) /
                    (2 * const.k_B)) * Ai  # NOTE: 1Jy = 1.0e-26 W m^-2 Hz^-1
                gain /= factor[:, np.newaxis, np.newaxis]

                # apply gain to vis
                if apply_gain:
                    for fi in range(nf):
                        for pi in [pol.index('xx'), pol.index('yy')]:
                            pi_ = gain_pd[pol[pi]]
                            for bi, (fd1, fd2) in enumerate(
                                    ts['blorder'].local_data):
                                g1 = gain[fi, pi_, feedno.index(fd1)]
                                g2 = gain[fi, pi_, feedno.index(fd2)]
                                if np.isfinite(g1) and np.isfinite(g2):
                                    if fd1 == fd2:
                                        # auto-correlation should be real
                                        ts.local_vis[:, fi, pi,
                                                     bi] /= (g1 *
                                                             np.conj(g2)).real
                                    else:
                                        ts.local_vis[:, fi, pi,
                                                     bi] /= (g1 * np.conj(g2))
                                else:
                                    # mask the un-calibrated vis
                                    ts.local_vis_mask[:, fi, pi, bi] = True

                    # in unit K after the calibration
                    ts.vis.attrs['unit'] = 'K'

                # save gain to file
                if save_gain:
                    if tag_output_iter:
                        gain_file = output_path(gain_file,
                                                iteration=self.iteration)
                    else:
                        gain_file = output_path(gain_file)
                    if mpiutil.rank0:
                        with h5py.File(gain_file, 'w') as f:
                            # allocate space for Gain
                            dset = f.create_dataset('Gain', (nt, nf, 2, nfeed),
                                                    dtype=Gain.dtype)
                            dset.attrs['calibrator'] = calibrator
                            dset.attrs['dim'] = 'time, freq, pol, feed'
                            try:
                                dset.attrs['time'] = ts.time[start_ind:end_ind]
                            except RuntimeError:
                                f.create_dataset(
                                    'time', data=ts.time[start_ind:end_ind])
                                dset.attrs['time'] = '/time'
                            dset.attrs['freq'] = freq
                            dset.attrs['pol'] = np.array(['xx', 'yy'])
                            dset.attrs['feed'] = np.array(feedno)
                            dset.attrs['transit_ind'] = transit_ind
                            # save gain
                            dset = f.create_dataset('gain', data=gain)
                            dset.attrs['calibrator'] = calibrator
                            dset.attrs['dim'] = 'freq, pol, feed'
                            dset.attrs['freq'] = freq
                            dset.attrs['pol'] = np.array(['xx', 'yy'])
                            dset.attrs['feed'] = np.array(feedno)
                            # save phs
                            if save_phs_change:
                                f.create_dataset('phs', data=phs)

                    mpiutil.barrier()

                    # save Gain
                    for i in range(10):
                        try:
                            # NOTE: if write simultaneously, will loss data with processes distributed in several nodes
                            for ri in xrange(mpiutil.size):
                                if ri == mpiutil.rank:
                                    with h5py.File(gain_file, 'r+') as f:
                                        for ii, (ti, fi,
                                                 pi) in enumerate(tfp_linds):
                                            ti_ = ti - start_ind
                                            pi_ = gain_pd[pol[pi]]
                                            f['Gain'][ti_, fi, pi_] = lGain[ii]
                                mpiutil.barrier()
                            break
                        except IOError:
                            time.sleep(0.5)
                            continue
                    else:
                        raise RuntimeError('Could not open file: %s...' %
                                           gain_file)

                    mpiutil.barrier()

        return super(PsCal, self).process(ts)
Esempio n. 4
0
    def _generate_mfiles(self, regen=False):

        completed_file = self._mdir + 'COMPLETED_BEAM'
        if os.path.exists(completed_file) and not regen:
            if mpiutil.rank0:
                print
                print '=' * 80
                print "******* Beam transfer m-files already generated ********"
            mpiutil.barrier()
            return

        if mpiutil.rank0:
            print
            print '=' * 80
            print 'Create beam transfer m-files...'

        st = time.time()

        # Calculate the Beam Transfer Matrices
        nfreq = self.telescope.nfreq
        nbl = self.telescope.nbase
        npol = self.telescope.num_pol_sky
        ntheta = self.telescope.theta_size
        nphi = self.telescope.phi_size

        # create file to save beam transfer matrices
        dsize = (nfreq, nbl, npol, ntheta)
        csize = (nfreq, 1, npol, ntheta)
        mmax = self.telescope.mmax
        ms = np.concatenate([np.arange(0, mmax+1), np.arange(-mmax, 0)])
        # get local section of m'th
        for ind, mi in enumerate(mpiutil.mpilist(ms, method='con')):
            with h5py.File(self._mfile(mi), 'w') as f:
                f.create_dataset('beam_m', dsize, chunks=csize, compression='lzf', dtype=np.complex128)
                f.attrs['m'] = mi

        # calculate the total memory needed for the transfer matrix
        total_memory = nfreq * nbl * npol * ntheta * nphi * 16.0 # Bytes, 16 for complex128
        limit = 1.0 # GB, memory limit for each process
        # make each process have maximum `limit` GB
        sigle_memory = limit * 2**30 # Bytes
        # how many chunks
        num_chunks = np.int(np.ceil(total_memory / (mpiutil.size * sigle_memory)))

        # split bls to num_chunks sections
        if nbl < num_chunks:
            warnings.warn('Could not split to %d chunks for %d baselines' % (num_chunks, nbl))
        num_chunks = min(num_chunks, nbl)
        num, start, end = mpiutil.split_m(nbl, num_chunks)
        for ci in range(num_chunks):
            if mpiutil.rank0:
                print "Starting chunk %i of %i" % (ci+1, num_chunks)

            tarray = self.telescope.transfer_matrices(np.arange(start[ci], end[ci]), np.arange(nfreq))
            tarray = MPIArray.wrap(tarray, axis=0)
            # redistribute along different m
            tarray = tarray.redistribute(axis=3)

            # save beam transfer matrices to file
            for ind, mi in enumerate(mpiutil.mpilist(ms, method='con')):
                with h5py.File(self._mfile(mi), 'r+') as f:
                    f['beam_m'][:, start[ci]:end[ci]] = tarray[..., ind].view(np.ndarray).reshape(nfreq, num[ci], npol, ntheta)

        mpiutil.barrier()

        et = time.time()

        if mpiutil.rank0:

            # Make file marker that the m's have been correctly generated:
            open(completed_file, 'a').close()

            # Print out timing
            print "=== Create beam transfer m-files took %f s ===" % (et - st)
Esempio n. 5
0
def hough_transform(I,
                    time,
                    freq,
                    dl,
                    dh,
                    nd=None,
                    t0l=None,
                    t0h=None,
                    nt0=None,
                    threshold=3.0,
                    comm=None):
    """Hough transform algorithm.

    Parameters
    ----------
    I : 2D np.ndarray
        Input data image, row time, column freq, both in ascending. order
    time : 1D np.ndarray
        The corresponding time of `I`, in ms.
    freq : 1D np.ndarray
        The corresponding frequency of `I`, in GHz.
    dl : float
        Lowest d = 4.15 DM.
    dh : float
        High d = 4.15 DM.
    nd : integer
        Number of d.
    t0l : float
        Lowest time offset t0, in ms.
    t0h : float
        Highest time offset t0, in ms.
    nt0 : integer
        Number of time offset t0.
    threshold : float
        How many sigmas to truncate the data.
    comm : MPI communicator
        MPI communicator. Required if executed parallelly by using MPI.
    """

    try:
        from caput import mpiutil
    except ImportError:
        # no mpiutil, can not use MPI to speed up
        comm = None

    fl, fh = freq[0], freq[-1]
    Nf = freq.shape[0]
    tl, th = time[0], time[-1]
    Nt = time.shape[0]
    df = (fh - fl) / Nf
    dt = (th - tl) / Nt

    # mask data based on a given threshold
    med = np.median(I)
    if threshold > 0.0:
        abs_diff = np.abs(I - med)
        mad = np.median(abs_diff) / 0.6745
        Im = np.where(abs_diff > threshold * mad, I - med,
                      np.nan)  # subtract median
    else:
        Im = I - med

    inds = np.where(np.isfinite(Im))  # inds of non-masked vals

    if nd is None:
        dd = fl**2 * dt
        nd = np.int(np.ceil((dh - dl) / dd))
    d = np.linspace(dl, dh, nd)
    # compute the range of t0
    if t0l is None:
        t0l = -fl**-2 * dh + tl
    if t0h is None:
        t0h = -fh**-2 * dl + th
    assert t0h > t0l, "Must have t0h (= %g) > t0l (= %g)" % (t0h, t0l)
    if nt0 is None:
        nt0 = nd
    dt0 = (t0h - t0l) / nt0

    # initialize the accumulator
    A = np.zeros((nt0, nd))  # the accumulator

    # accumulat the accumulator
    if comm is None or comm.size == 1:
        linds = zip(inds[0], inds[1])
    else:
        linds = mpiutil.mpilist(zip(inds[0], inds[1]), comm=comm)
    for (ti, fi) in linds:
        fv = fl + fi * df
        tv = tl + ti * dt
        t0v = -fv**-2 * d + tv
        for di in xrange(nd):
            t0i = np.int(np.around((t0v[di] - t0l) / dt0))
            t0i = max(0, t0i)
            t0i = min(nt0 - 1, t0i)
            A[t0i, di] += Im[ti, fi]

    # gather and accumulat A
    if not (comm is None or comm.size == 1):
        A = A.reshape((1, nt0, nd))
        Ac = mpiutil.gather_array(A, comm=comm)
        if mpiutil.rank0:
            A = Ac.sum(axis=0)
        else:
            A = None

    return A, dl, dh, t0l, t0h
Esempio n. 6
0
def separator(sec, tag):
    # sleep, sync, and flush to avoid output of different parts being mixed
    time.sleep(sec)
    mpiutil.barrier()
    sys.stdout.flush()

    if rank == 0:
        print
        print '-' * 35 + ' ' + tag + ' ' + '-' * 35


# mpilist
separator(sec, 'mpilist')
full_list = [1, 2.5, 'a', True, (3, 4), {'x': 1}]
local_list = mpiutil.mpilist(full_list)
print "rank %d has %s with method = 'con'" % (rank, local_list)
local_list = mpiutil.mpilist(full_list, method='alt')
print "rank %d has %s with method = 'alt'" % (rank, local_list)
local_list = mpiutil.mpilist(full_list, method='rand')
print "rank %d has %s with method = 'rand'" % (rank, local_list)

# mpirange
separator(sec, 'mpirange')
local_ary = mpiutil.mpirange(1, 7)
print "rank %d has %s with method = 'con'" % (rank, local_ary)
local_ary = mpiutil.mpirange(1, 7, method='alt')
print "rank %d has %s with method = 'alt'" % (rank, local_ary)
local_ary = mpiutil.mpirange(1, 7, method='rand')
print "rank %d has %s with method = 'rand'" % (rank, local_ary)
Esempio n. 7
0
    def generate_mmodes(self, ts_data=None):
        """Calculate the m-modes corresponding to the Timestream.

        Perform an MPI transpose for efficiency.
        """

        completed_file = self._mdir + 'COMPLETED_M'
        if os.path.exists(completed_file):
            if mpiutil.rank0:
                print "******* m-files already generated ********"
            mpiutil.barrier()
            return

        # Make directory if required
        # if mpiutil.rank0 and not os.path.exists(self._mdir):
        #     os.makedirs(self._mdir)

        try:
            os.makedirs(self._mdir)
        except OSError:
            # directory exists
            pass

        tel = self.telescope
        mmax = tel.mmax
        ntime = ts_data.shape[0] if ts_data is not None else self.ntime
        nbl = tel.nbase
        nfreq = tel.nfreq

        indices = list(itertools.product(np.arange(nfreq), np.arange(nbl)))
        lind, sind, eind = mpiutil.split_local(nfreq * nbl)

        # load the local section of the time stream
        tstream = np.zeros((ntime, lind), dtype=np.complex128)
        for ind, (f_ind, bl_ind) in enumerate(indices[sind:eind]):
            if ts_data is not None:
                tstream[:, ind] = ts_data[:, f_ind, bl_ind]
            else:
                with h5py.File(self._tsfile, 'r') as f:
                    tstream[:, ind] = f['/timestream'][:, f_ind, bl_ind]

        # FFT to get m-mode
        mmodes = np.fft.fft(tstream, axis=0) / ntime  # m = 0 is at left
        mmodes = MPIArray.wrap(mmodes, axis=1)
        # redistribute along different m
        mmodes = mmodes.redistribute(axis=0)

        # save m-modes to file
        ms = np.concatenate([np.arange(0, mmax + 1), np.arange(-mmax, 0)])
        for ind, mi in enumerate(mpiutil.mpilist(ms, method='con')):
            with h5py.File(self._mfile(mi), 'w') as f:
                f.create_dataset('/mmode',
                                 data=mmodes[ind].view(np.ndarray).reshape(
                                     nfreq, nbl))
                f.attrs['m'] = mi

        mpiutil.barrier()

        if mpiutil.rank0:

            # Make file marker that the m's have been correctly generated:
            open(completed_file, 'a').close()
Esempio n. 8
0
    def process(self, ts):

        assert isinstance(
            ts, Timestream
        ), '%s only works for Timestream object' % self.__class__.__name__

        calibrator = self.params['calibrator']
        catalog = self.params['catalog']
        span = self.params['span']
        plot_figs = self.params['plot_figs']
        fig_prefix = self.params['fig_name']
        tag_output_iter = self.params['tag_output_iter']
        save_gain = self.params['save_gain']
        gain_file = self.params['gain_file']
        temperature_convert = self.params['temperature_convert']

        ts.redistribute('baseline')

        feedno = ts['feedno'][:].tolist()
        pol = [ts.pol_dict[p] for p in ts['pol'][:]]
        bl = mpiutil.gather_array(ts.local_bl[:], root=None, comm=ts.comm)
        bls = [tuple(b) for b in bl]
        # # antpointing = np.radians(ts['antpointing'][-1, :, :]) # radians
        # transitsource = ts['transitsource'][:]
        # transit_time = transitsource[-1, 0] # second, sec1970
        # int_time = ts.attrs['inttime'] # second

        # calibrator
        srclist, cutoff, catalogs = a.scripting.parse_srcs(calibrator, catalog)
        cat = a.src.get_catalog(srclist, cutoff, catalogs)
        assert (len(cat) == 1), 'Allow only one calibrator'
        s = cat.values()[0]
        if mpiutil.rank0:
            print 'Calibrating for source %s with' % calibrator,
            print 'strength', s._jys, 'Jy',
            print 'measured at', s.mfreq, 'GHz',
            print 'with index', s.index

        # get transit time of calibrator
        # array
        aa = ts.array
        aa.set_jultime(ts['jul_date'][0])  # the first obs time point
        next_transit = aa.next_transit(s)
        transit_time = a.phs.ephem2juldate(next_transit)  # Julian date
        # get time zone
        pattern = '[-+]?\d+'
        tz = re.search(pattern, ts.attrs['timezone']).group()
        tz = int(tz)
        local_next_transit = ephem.Date(
            next_transit + tz * ephem.hour)  # plus 8h to get Beijing time
        if transit_time > ts['jul_date'][-1]:
            raise RuntimeError(
                'Data does not contain local transit time %s of source %s' %
                (local_next_transit, calibrator))

        # the first transit index
        transit_inds = [np.searchsorted(ts['jul_date'][:], transit_time)]
        # find all other transit indices
        aa.set_jultime(ts['jul_date'][0] + 1.0)
        transit_time = a.phs.ephem2juldate(aa.next_transit(s))  # Julian date
        cnt = 2
        while (transit_time <= ts['jul_date'][-1]):
            transit_inds.append(
                np.searchsorted(ts['jul_date'][:], transit_time))
            aa.set_jultime(ts['jul_date'][0] + 1.0 * cnt)
            transit_time = a.phs.ephem2juldate(
                aa.next_transit(s))  # Julian date
            cnt += 1

        if mpiutil.rank0:
            print 'transit ind of %s: %s, time: %s' % (
                calibrator, transit_inds, local_next_transit)

        ### now only use the first transit point to do the cal
        ### may need to improve in the future
        transit_ind = transit_inds[0]
        int_time = ts.attrs['inttime']  # second
        start_ind = transit_ind - np.int(span / int_time)
        end_ind = transit_ind + np.int(
            span / int_time) + 1  # plus 1 to make transit_ind is at the center

        # check if data contain this range
        if start_ind < 0:
            raise RuntimeError('start_ind: %d < 0' % start_ind)
        if end_ind > ts.vis.shape[0]:
            raise RuntimeError('end_ind: %d > %d' % (end_ind, ts.vis.shape[0]))

        ############################################
        # if ts.is_cylinder:
        #     ts.local_vis[:] = ts.local_vis.conj() # now for cylinder array
        ############################################

        nt = end_ind - start_ind
        t_inds = range(start_ind, end_ind)
        freq = ts.freq[:]
        nf = len(freq)
        nlb = len(ts.local_bl[:])
        nfeed = len(feedno)
        tfp_inds = list(
            itertools.product(
                t_inds, range(nf),
                [pol.index('xx'), pol.index('yy')]))  # only for xx and yy
        ns, ss, es = mpiutil.split_all(len(tfp_inds), comm=ts.comm)
        # gather data to make each process to have its own data which has all bls
        for ri, (ni, si, ei) in enumerate(zip(ns, ss, es)):
            lvis = np.zeros((ni, nlb), dtype=ts.vis.dtype)
            lvis_mask = np.zeros((ni, nlb), dtype=ts.vis_mask.dtype)
            for ii, (ti, fi, pi) in enumerate(tfp_inds[si:ei]):
                lvis[ii] = ts.local_vis[ti, fi, pi]
                lvis_mask[ii] = ts.local_vis_mask[ti, fi, pi]
            # gather vis from all process for separate bls
            gvis = mpiutil.gather_array(lvis, axis=1, root=ri, comm=ts.comm)
            gvis_mask = mpiutil.gather_array(lvis_mask,
                                             axis=1,
                                             root=ri,
                                             comm=ts.comm)
            if ri == mpiutil.rank:
                tfp_linds = tfp_inds[si:ei]  # inds for this process
                this_vis = gvis
                this_vis_mask = gvis_mask
        del tfp_inds
        del lvis
        del lvis_mask
        lGain = np.empty((len(tfp_linds), nfeed), dtype=np.complex128)
        lGain[:] = complex(np.nan, np.nan)

        # construct visibility matrix for a single time, freq, pol
        Vmat = np.zeros((nfeed, nfeed), dtype=ts.vis.dtype)
        Sc = s.get_jys()
        for ii, (ti, fi, pi) in enumerate(tfp_linds):
            # when noise on, just pass
            if 'ns_on' in ts.iterkeys() and ts['ns_on'][ti]:
                continue
            # aa.set_jultime(ts['jul_date'][ti])
            # s.compute(aa)
            # get fluxes vs. freq of the calibrator
            # Sc = s.get_jys()
            # get the topocentric coordinate of the calibrator at the current time
            # s_top = s.get_crds('top', ncrd=3)
            # aa.sim_cache(cat.get_crds('eq', ncrd=3)) # for compute bm_response and sim
            mask_cnt = 0
            for i, ai in enumerate(feedno):
                for j, aj in enumerate(feedno):
                    try:
                        bi = bls.index((ai, aj))
                        if this_vis_mask[ii, bi] and not np.isfinite(
                                this_vis[ii, bi]):
                            mask_cnt += 1
                            Vmat[i, j] = 0
                        else:
                            Vmat[i, j] = this_vis[ii, bi] / Sc[fi]  # xx, yy
                    except ValueError:
                        bi = bls.index((aj, ai))
                        if this_vis_mask[ii, bi] and not np.isfinite(
                                this_vis[ii, bi]):
                            mask_cnt += 1
                            Vmat[i, j] = 0
                        else:
                            Vmat[i, j] = np.conj(this_vis[ii, bi] /
                                                 Sc[fi])  # xx, yy

            # if too many masks
            if mask_cnt > 0.3 * nfeed**2:
                continue

            # set invalid val to 0
            # Vmat = np.where(np.isfinite(Vmat), Vmat, 0)

            # initialize the outliers
            med = np.median(Vmat.real) + 1.0J * np.median(Vmat.imag)
            diff = Vmat - med
            S0 = np.where(np.abs(diff) > 3.0 * rpca_decomp.MAD(Vmat), diff, 0)
            # stable PCA decomposition
            V0, S = rpca_decomp.decompose(Vmat,
                                          rank=1,
                                          S=S0,
                                          max_iter=100,
                                          threshold='hard',
                                          tol=1.0e-6,
                                          debug=False)
            # V0, S = rpca_decomp.decompose(Vmat, rank=1, S=S0, max_iter=100, threshold='soft', tol=1.0e-6, debug=False)

            # plot
            if plot_figs:
                ind = ti - start_ind
                # plot Vmat
                plt.figure(figsize=(13, 5))
                plt.subplot(121)
                plt.imshow(Vmat.real,
                           aspect='equal',
                           origin='lower',
                           interpolation='nearest')
                plt.colorbar(shrink=1.0)
                plt.subplot(122)
                plt.imshow(Vmat.imag,
                           aspect='equal',
                           origin='lower',
                           interpolation='nearest')
                plt.colorbar(shrink=1.0)
                fig_name = '%s_V_%d_%d_%s.png' % (fig_prefix, ind, fi,
                                                  ts.pol_dict[pi])
                if tag_output_iter:
                    fig_name = output_path(fig_name, iteration=self.iteration)
                else:
                    fig_name = output_path(fig_name)
                plt.savefig(fig_name)
                plt.close()
                # plot V0
                plt.figure(figsize=(13, 5))
                plt.subplot(121)
                plt.imshow(V0.real,
                           aspect='equal',
                           origin='lower',
                           interpolation='nearest')
                plt.colorbar(shrink=1.0)
                plt.subplot(122)
                plt.imshow(V0.imag,
                           aspect='equal',
                           origin='lower',
                           interpolation='nearest')
                plt.colorbar(shrink=1.0)
                fig_name = '%s_V0_%d_%d_%s.png' % (fig_prefix, ind, fi,
                                                   ts.pol_dict[pi])
                if tag_output_iter:
                    fig_name = output_path(fig_name, iteration=self.iteration)
                else:
                    fig_name = output_path(fig_name)
                plt.savefig(fig_name)
                plt.close()
                # plot S
                plt.figure(figsize=(13, 5))
                plt.subplot(121)
                plt.imshow(S.real,
                           aspect='equal',
                           origin='lower',
                           interpolation='nearest')
                plt.colorbar(shrink=1.0)
                plt.subplot(122)
                plt.imshow(S.imag,
                           aspect='equal',
                           origin='lower',
                           interpolation='nearest')
                plt.colorbar(shrink=1.0)
                fig_name = '%s_S_%d_%d_%s.png' % (fig_prefix, ind, fi,
                                                  ts.pol_dict[pi])
                if tag_output_iter:
                    fig_name = output_path(fig_name, iteration=self.iteration)
                else:
                    fig_name = output_path(fig_name)
                plt.savefig(fig_name)
                plt.close()
                # plot N
                N = Vmat - V0 - S
                plt.figure(figsize=(13, 5))
                plt.subplot(121)
                plt.imshow(N.real,
                           aspect='equal',
                           origin='lower',
                           interpolation='nearest')
                plt.colorbar(shrink=1.0)
                plt.subplot(122)
                plt.imshow(N.imag,
                           aspect='equal',
                           origin='lower',
                           interpolation='nearest')
                plt.colorbar(shrink=1.0)
                fig_name = '%s_N_%d_%d_%s.png' % (fig_prefix, ind, fi,
                                                  ts.pol_dict[pi])
                if tag_output_iter:
                    fig_name = output_path(fig_name, iteration=self.iteration)
                else:
                    fig_name = output_path(fig_name)
                plt.savefig(fig_name)
                plt.close()

            e, U = la.eigh(V0, eigvals=(nfeed - 1, nfeed - 1))
            g = U[:, -1] * e[-1]**0.5
            lGain[ii] = g

            # plot Gain
            if plot_figs:
                plt.figure()
                plt.plot(feedno, g.real, 'b-', label='real')
                plt.plot(feedno, g.real, 'bo')
                plt.plot(feedno, g.imag, 'g-', label='imag')
                plt.plot(feedno, g.imag, 'go')
                plt.plot(feedno, np.abs(g), 'r-', label='abs')
                plt.plot(feedno, np.abs(g), 'ro')
                plt.xlim(feedno[0] - 1, feedno[-1] + 1)
                yl, yh = plt.ylim()
                plt.ylim(yl, yh + (yh - yl) / 5)
                plt.xlabel('Feed number')
                plt.legend()
                fig_name = '%s_ants_%d_%d_%s.png' % (fig_prefix, ind, fi,
                                                     ts.pol_dict[pi])
                if tag_output_iter:
                    fig_name = output_path(fig_name, iteration=self.iteration)
                else:
                    fig_name = output_path(fig_name)
                plt.savefig(fig_name)
                plt.close()

        # gather Gain from each processes
        Gain = mpiutil.gather_array(lGain, axis=0, root=None, comm=ts.comm)
        Gain = Gain.reshape(nt, nf, 2, nfeed)

        # choose data slice near the transit time
        c = nt / 2  # center ind
        li = max(0, c - 100)
        hi = min(nt, c + 100 + 1)
        x = np.arange(li, hi)
        # compute s_top for this time range
        n0 = np.zeros(((hi - li), 3))
        for ti, jt in enumerate(ts.time[start_ind:end_ind][li:hi]):
            aa.set_jultime(jt)
            s.compute(aa)
            n0[ti] = s.get_crds('top', ncrd=3)

        # get the positions of feeds
        feedpos = ts['feedpos'][:]

        # create data to save the solved gain for each feed
        local_fp_inds = mpiutil.mpilist(
            list(itertools.product(range(nf), range(2))))
        lgain = np.zeros((len(local_fp_inds), nfeed),
                         dtype=Gain.dtype)  # gain for each feed
        lgain[:] = complex(np.nan, np.nan)

        for ii, (fi, pi) in enumerate(local_fp_inds):
            data = np.abs(Gain[li:hi, fi, pi, :]).T
            # flag outliers
            median = np.ma.median(data, axis=0)
            abs_diff = np.ma.abs(data - median[np.newaxis, :])
            mad = np.ma.median(abs_diff, axis=0) / 0.6745
            with warnings.catch_warnings():
                warnings.filterwarnings(
                    'ignore', 'invalid value encountered in greater')
                warnings.filterwarnings(
                    'ignore', 'invalid value encountered in greater_equal')
                warnings.filterwarnings(
                    'ignore', 'invalid value encountered in absolute')
                data = np.where(abs_diff > 3.0 * mad[np.newaxis, :], np.nan,
                                data)
            # gaussian/sinc fit
            for idx in range(nfeed):
                y = data[idx]
                inds = np.where(np.isfinite(y))[0]
                if len(inds) > 0.75 * len(y):
                    # get the best estimate of the central val
                    cval = y[inds[np.argmin(np.abs(inds - c))]]
                    try:
                        # gaussian fit
                        # popt, pcov = optimize.curve_fit(fg, x[inds], y[inds], p0=(cval, c, 90, 0))
                        # sinc function seems fit better
                        popt, pcov = optimize.curve_fit(fc,
                                                        x[inds],
                                                        y[inds],
                                                        p0=(cval, c, 1.0e-2,
                                                            0))
                        # print 'popt:', popt
                    except RuntimeError:
                        print 'curve_fit failed for fi = %d, pol = %s, feed = %d' % (
                            fi, ['xx', 'yy'][pi], feedno[idx])
                        continue

                    An = y / fc(popt[1], *popt)  # the beam profile
                    ui = (feedpos[idx] - feedpos[0]) * (
                        1.0e6 * freq[fi]
                    ) / const.c  # position of this feed (relative to the first feed) in unit of wavelength
                    exp_factor = np.exp(2.0J * np.pi * np.dot(n0, ui))
                    Ae = An * exp_factor
                    Gi = Gain[li:hi, fi, pi, idx]
                    # compute gain for this feed
                    lgain[ii,
                          idx] = np.dot(Ae[inds].conj(), Gi[inds]) / np.dot(
                              Ae[inds].conj(), Ae[inds])

        # gather local gain
        gain = mpiutil.gather_array(lgain, axis=0, root=None, comm=ts.comm)
        gain = gain.reshape(nf, 2, nfeed)

        # apply gain to vis
        for fi in range(nf):
            for pi in [pol.index('xx'), pol.index('yy')]:
                for bi, (fd1, fd2) in enumerate(ts['blorder'].local_data):
                    g1 = gain[fi, pi, feedno.index(fd1)]
                    g2 = gain[fi, pi, feedno.index(fd2)]
                    if np.isfinite(g1) and np.isfinite(g2):
                        ts.local_vis[:, fi, pi, bi] /= (g1 * np.conj(g2))
                    else:
                        # mask the un-calibrated vis
                        ts.local_vis_mask[:, fi, pi, bi] = True

        # convert vis from intensity unit to temperature unit in K
        if temperature_convert:
            factor = 1.0e-26 * (const.c**2 / (2 * const.k_B *
                                              (1.0e6 * freq)**2)
                                )  # NOTE: 1Jy = 1.0e-26 W m^-2 Hz^-1
            ts.local_vis[:] *= factor[np.newaxis, :, np.newaxis, np.newaxis]
            ts.vis.attrs['unit'] = 'K'

        # save gain to file
        if mpiutil.rank0 and save_gain:
            if tag_output_iter:
                gain_file = output_path(gain_file, iteration=self.iteration)
            else:
                gain_file = output_path(gain_file)
            with h5py.File(gain_file, 'w') as f:
                # save Gain
                Gain = f.create_dataset('Gain', data=Gain)
                Gain.attrs['dim'] = 'time, freq, pol, feed'
                Gain.attrs['time'] = ts.time[start_ind:end_ind]
                Gain.attrs['freq'] = freq
                Gain.attrs['pol'] = np.array(['xx', 'yy'])
                Gain.attrs['feed'] = np.array(feedno)
                # save gain
                gain = f.create_dataset('gain', data=gain)
                gain.attrs['dim'] = 'freq, pol, feed'
                gain.attrs['freq'] = freq
                gain.attrs['pol'] = np.array(['xx', 'yy'])
                gain.attrs['feed'] = np.array(feedno)

        return super(PsCal, self).process(ts)
Esempio n. 9
0
    def process(self, ts):

        assert isinstance(
            ts, Timestream
        ), '%s only works for Timestream object' % self.__class__.__name__

        #        if mpiutil.rank0:
        #            saveVmat = []

        calibrator = self.params['calibrator']
        catalog = self.params['catalog']
        vis_conj = self.params['vis_conj']
        zero_diag = self.params['zero_diag']
        span = self.params['span']
        reserve_high_gain = self.params['reserve_high_gain']
        plot_figs = self.params['plot_figs']
        fig_prefix = self.params['fig_name']
        tag_output_iter = self.params['tag_output_iter']
        save_src_vis = self.params['save_src_vis']
        src_vis_file = self.params['src_vis_file']
        subtract_src = self.params['subtract_src']
        apply_gain = self.params['apply_gain']
        save_gain = self.params['save_gain']
        save_phs_change = self.params['save_phs_change']
        gain_file = self.params['gain_file']
        temperature_convert = self.params['temperature_convert']
        show_progress = self.params['show_progress']
        progress_step = self.params['progress_step']
        srcdict = self.params['srcdict']
        max_iter = self.params['max_iter']

        #        if mpiutil.rank0:
        #            print(ts.keys())
        #            print(ts.attrs.keys())
        #
        #        import sys
        #        sys.exit()

        if save_src_vis or subtract_src or apply_gain or save_gain:
            pol_type = ts['pol'].attrs['pol_type']
            if pol_type != 'linear':
                raise RuntimeError('Can not do ps_cal for pol_type: %s' %
                                   pol_type)

            ts.redistribute('baseline')

            feedno = ts['feedno'][:].tolist()
            pol = [ts.pol_dict[p] for p in ts['pol'][:]]  # as string
            gain_pd = {
                'xx': 0,
                'yy': 1,
                0: 'xx',
                1: 'yy'
            }  # for gain related op
            bls = mpiutil.gather_array(ts.local_bl[:], root=None, comm=ts.comm)
            # # antpointing = np.radians(ts['antpointing'][-1, :, :]) # radians
            # transitsource = ts['transitsource'][:]
            # transit_time = transitsource[-1, 0] # second, sec1970
            # int_time = ts.attrs['inttime'] # second

            # calibrator
            #see whether the source should be observed
            #only for single calibrator
            if ts.is_dish:
                obsd_sources = ts['transitsource'].attrs['srcname']
                obsd_sources = obsd_sources.split(',')
                #                srcdict = {'cas':'CassiopeiaA'}
                src_index = 0.1
                for src_short, src_long in srcdict.items():
                    #                    print(src_short,src_long,obsd_sources,calibrator)
                    if (src_short == calibrator) and (src_long
                                                      in obsd_sources):
                        src_index = obsd_sources.index(src_long)
                        break
                else:
                    raise Exception(
                        'the transit source is not included in the observation plan!'
                    )
                obs_transit_time = ts['transitsource'][src_index][0]

#            srclist, cutoff, catalogs = a.scripting.parse_srcs(calibrator, catalog)
#            cat = a.src.get_catalog(srclist, cutoff, catalogs)
#            assert(len(cat) == 1), 'Allow only one calibrator'
#            s = cat.values()[0]

# get the calibrator
            try:
                s = calibrators.get_src(calibrator)
            except KeyError:
                if mpiutil.rank0:
                    print 'Calibrator %s is unavailable, available calibrators are:'
                    for key, d in calibrators.src_data.items():
                        print '%8s  ->  %12s' % (key, d[0])
                raise RuntimeError('Calibrator %s is unavailable')
            if mpiutil.rank0:
                print 'Try to calibrate with %s...' % s.src_name
#            if mpiutil.rank0:
#                print 'Calibrating for source %s with' % calibrator,
#                print 'strength', s._jys, 'Jy',
#                print 'measured at', s.mfreq, 'GHz',
#                print 'with index', s.index

# get transit time of calibrator
# array
            aa = ts.array
            aa.set_jultime(ts['jul_date'][0])  # the first obs time point
            next_transit = aa.next_transit(s)
            next_transit = aa.next_transit(
                s) if not ts.is_dish else ephem.date(
                    datetime.utcfromtimestamp(obs_transit_time))
            transit_time = a.phs.ephem2juldate(next_transit)  # Julian date
            # get time zone
            pattern = '[-+]?\d+'
            tz = re.search(pattern, ts.attrs['timezone']).group()
            tz = int(tz)

            local_next_transit = ephem.Date(
                next_transit + tz * ephem.hour)  # plus 8h to get Beijing time
            # if transit_time > ts['jul_date'][-1]:
            #========================================================
            if ts.is_dish:
                if (transit_time >
                        max(ts['jul_date'][-1],
                            ts['jul_date'][:].max())) or (transit_time < min(
                                ts['jul_date'][0], ts['jul_date'][:].min())):
                    #                    raise RuntimeError('dish data does not contain local transit time %s of source %s' % (local_next_transit, calibrator))
                    raise NoTransit(
                        'Dish data does not contain local transit time %s of source %s'
                        % (local_next_transit, calibrator))
                transit_inds = [
                    np.searchsorted(ts['jul_date'][:], transit_time)
                ]
                peak_span = int(np.around(10. / ts.attrs['inttime']))
                peak_start = ts['jul_date'][transit_inds[0] - peak_span]
                peak_end = ts['jul_date'][transit_inds[0] + peak_span]
                tspan = (peak_end - peak_start) * 86400.
                if mpiutil.rank0:
                    print('transit peak: %s' % local_next_transit)
                    print('%d point (should be about 10s) previous: %s' %
                          (peak_span,
                           ephem.date(
                               a.phs.juldate2ephem(peak_start) +
                               tz * ephem.hour)))
                    print(
                        '%d point (should be about 10s) later: %s' %
                        (peak_span,
                         ephem.date(
                             a.phs.juldate2ephem(peak_end) + tz * ephem.hour)))
                if tspan > 30.:
                    warnings.warn(
                        'Peak of transit is not continuous in the data! May lead to poor performance!'
                    )
            else:
                if transit_time > max(ts['jul_date'][-1],
                                      ts['jul_date'][:].max()):
                    #                    raise RuntimeError('Cylinder data does not contain local transit time %s of source %s' % (local_next_transit, calibrator))
                    raise NoTransit(
                        'Cylinder data does not contain local transit time %s of source %s'
                        % (local_next_transit, calibrator))

                # the first transit index
                transit_inds = [
                    np.searchsorted(ts['jul_date'][:], transit_time)
                ]
                # find all other transit indices
                aa.set_jultime(ts['jul_date'][0] + 1.0)
                transit_time = a.phs.ephem2juldate(
                    aa.next_transit(s))  # Julian date
                cnt = 2
                while (transit_time <= ts['jul_date'][-1]):
                    transit_inds.append(
                        np.searchsorted(ts['jul_date'][:], transit_time))
                    aa.set_jultime(ts['jul_date'][0] + 1.0 * cnt)
                    # Julian date
                    #                transit_time = a.phs.ephem2juldate(aa.next_transit(s) if not ts.is_dish else ephem.date(datetime.utcfromtimestamp(obs_transit_time)))
                    transit_time = a.phs.ephem2juldate(aa.next_transit(s))
                    cnt += 1
#========================================================

            if mpiutil.rank0:
                print 'transit ind of %s: %s, time: %s' % (
                    s.src_name, transit_inds, local_next_transit)

            if (not ts.ps_first) and ts.interp_all_masked:
                raise NotEnoughPointToInterpolateError(
                    'More than 80% of the data was masked due to shortage of noise points for interpolation(need at least 4 to perform cubic spline)! The pointsource calibration may not be done due to too many masked points!'
                )
            ### now only use the first transit point to do the cal
            ### may need to improve in the future
            transit_ind = transit_inds[0]
            int_time = ts.attrs['inttime']  # second
            start_ind = transit_ind - np.int(span / int_time)
            end_ind = transit_ind + np.int(
                span /
                int_time) + 1  # plus 1 to make transit_ind is at the center

            start_ind = max(0, start_ind)
            end_ind = min(end_ind, ts.vis.shape[0])

            if vis_conj:
                ts.local_vis[:] = ts.local_vis.conj()

            nt = end_ind - start_ind
            t_inds = range(start_ind, end_ind)
            freq = ts.freq[:]
            nf = len(freq)
            nlb = len(ts.local_bl[:])
            nfeed = len(feedno)
            tfp_inds = list(
                itertools.product(
                    t_inds, range(nf),
                    [pol.index('xx'), pol.index('yy')]))  # only for xx and yy
            ns, ss, es = mpiutil.split_all(len(tfp_inds), comm=ts.comm)
            # gather data to make each process to have its own data which has all bls
            for ri, (ni, si, ei) in enumerate(zip(ns, ss, es)):
                lvis = np.zeros((ni, nlb), dtype=ts.vis.dtype)
                lvis_mask = np.zeros((ni, nlb), dtype=ts.vis_mask.dtype)
                for ii, (ti, fi, pi) in enumerate(tfp_inds[si:ei]):
                    lvis[ii] = ts.local_vis[ti, fi, pi]
                    lvis_mask[ii] = ts.local_vis_mask[ti, fi, pi]
                # gather vis from all process for separate bls
                gvis = mpiutil.gather_array(lvis,
                                            axis=1,
                                            root=ri,
                                            comm=ts.comm)
                gvis_mask = mpiutil.gather_array(lvis_mask,
                                                 axis=1,
                                                 root=ri,
                                                 comm=ts.comm)
                if ri == mpiutil.rank:
                    tfp_linds = tfp_inds[si:ei]  # inds for this process
                    this_vis = gvis
                    this_vis_mask = gvis_mask
            del tfp_inds
            del lvis
            del lvis_mask
            tfp_len = len(tfp_linds)

            cnan = complex(np.nan, np.nan)  # complex nan
            if save_src_vis or subtract_src:
                # save calibrator src vis
                lsrc_vis = np.full((tfp_len, nfeed, nfeed),
                                   cnan,
                                   dtype=ts.vis.dtype)
                if save_src_vis:
                    # save sky vis
                    lsky_vis = np.full((tfp_len, nfeed, nfeed),
                                       cnan,
                                       dtype=ts.vis.dtype)
                    # save outlier vis
                    lotl_vis = np.full((tfp_len, nfeed, nfeed),
                                       cnan,
                                       dtype=ts.vis.dtype)

            if apply_gain or save_gain:
                lGain = np.full((tfp_len, nfeed), cnan, dtype=ts.vis.dtype)

            # find indices mapping between Vmat and vis
            # bis = range(nbl)
            bis_conj = []  # indices that shold be conj
            mis = [
            ]  # indices in the nfeed x nfeed matrix by flatten it to a vector
            mis_conj = [
            ]  # indices (of conj vis) in the nfeed x nfeed matrix by flatten it to a vector
            for bi, (fdi, fdj) in enumerate(bls):
                ai, aj = feedno.index(fdi), feedno.index(fdj)
                mis.append(ai * nfeed + aj)
                if ai != aj:
                    bis_conj.append(bi)
                    mis_conj.append(aj * nfeed + ai)

            # construct visibility matrix for a single time, freq, pol
            Vmat = np.full((nfeed, nfeed), cnan, dtype=ts.vis.dtype)
            # get flus of the calibrator in the observing frequencies
            Sc = s.get_jys(freq)
            if show_progress and mpiutil.rank0:
                pg = progress.Progress(tfp_len, step=progress_step)
            for ii, (ti, fi, pi) in enumerate(tfp_linds):
                if show_progress and mpiutil.rank0:
                    pg.show(ii)
                # when noise on, just pass
                if 'ns_on' in ts.iterkeys() and ts['ns_on'][ti]:
                    continue
                # aa.set_jultime(ts['jul_date'][ti])
                # s.compute(aa)
                # get the topocentric coordinate of the calibrator at the current time
                # s_top = s.get_crds('top', ncrd=3)
                # aa.sim_cache(cat.get_crds('eq', ncrd=3)) # for compute bm_response and sim
                Vmat.flat[mis] = np.ma.array(
                    this_vis[ii], mask=this_vis_mask[ii]).filled(cnan)
                Vmat.flat[mis_conj] = np.ma.array(
                    this_vis[ii, bis_conj],
                    mask=this_vis_mask[ii, bis_conj]).conj().filled(cnan)

                #                if mpiutil.rank0:
                #                    saveVmat += [Vmat.copy()]

                if save_src_vis:
                    lsky_vis[ii] = Vmat

                # set invalid val to 0
                invalid = ~np.isfinite(Vmat)  # a bool array
                # if too many masks

                if np.where(invalid)[0].shape[0] > 0.3 * nfeed**2:
                    continue
                Vmat[invalid] = 0
                #                # if all are zeros
                if np.allclose(Vmat, 0.0):
                    continue

                # fill diagonal of Vmat to 0
                if zero_diag:
                    np.fill_diagonal(Vmat, 0)

                # initialize the outliers
                med = np.median(Vmat.real) + 1.0J * np.median(Vmat.imag)
                diff = Vmat - med
                S0 = np.where(
                    np.abs(diff) > 3.0 * rpca_decomp.MAD(Vmat), diff, 0)
                # stable PCA decomposition
                #                V0, S = rpca_decomp.decompose(Vmat, rank=1, S=S0, max_iter=100, threshold='hard', tol=1.0e-6, debug=False)
                V0, S = rpca_decomp.decompose(Vmat,
                                              rank=1,
                                              S=S0,
                                              max_iter=max_iter,
                                              threshold='hard',
                                              tol=1.0e-6,
                                              debug=False)
                #                V0, S = rpca_decomp.decompose(Vmat, rank=1, S=S0, max_iter=1, threshold='hard', tol=1., debug=False)
                if save_src_vis or subtract_src:
                    lsrc_vis[ii] = V0
                    if save_src_vis:
                        lotl_vis[ii] = S

                # plot
                if plot_figs:
                    ind = ti - start_ind
                    # plot Vmat
                    plt.figure(figsize=(13, 5))
                    plt.subplot(121)
                    plt.imshow(Vmat.real,
                               aspect='equal',
                               origin='lower',
                               interpolation='nearest')
                    plt.colorbar(shrink=1.0)
                    plt.subplot(122)
                    plt.imshow(Vmat.imag,
                               aspect='equal',
                               origin='lower',
                               interpolation='nearest')
                    plt.colorbar(shrink=1.0)
                    fig_name = '%s_V_%d_%d_%s.png' % (fig_prefix, ind, fi,
                                                      pol[pi])
                    if tag_output_iter:
                        fig_name = output_path(fig_name,
                                               iteration=self.iteration)
                    else:
                        fig_name = output_path(fig_name)
                    plt.savefig(fig_name)
                    plt.close()
                    # plot V0
                    plt.figure(figsize=(13, 5))
                    plt.subplot(121)
                    plt.imshow(V0.real,
                               aspect='equal',
                               origin='lower',
                               interpolation='nearest')
                    plt.colorbar(shrink=1.0)
                    plt.subplot(122)
                    plt.imshow(V0.imag,
                               aspect='equal',
                               origin='lower',
                               interpolation='nearest')
                    plt.colorbar(shrink=1.0)
                    fig_name = '%s_V0_%d_%d_%s.png' % (fig_prefix, ind, fi,
                                                       pol[pi])
                    if tag_output_iter:
                        fig_name = output_path(fig_name,
                                               iteration=self.iteration)
                    else:
                        fig_name = output_path(fig_name)
                    plt.savefig(fig_name)
                    plt.close()
                    # plot S
                    plt.figure(figsize=(13, 5))
                    plt.subplot(121)
                    plt.imshow(S.real,
                               aspect='equal',
                               origin='lower',
                               interpolation='nearest')
                    plt.colorbar(shrink=1.0)
                    plt.subplot(122)
                    plt.imshow(S.imag,
                               aspect='equal',
                               origin='lower',
                               interpolation='nearest')
                    plt.colorbar(shrink=1.0)
                    fig_name = '%s_S_%d_%d_%s.png' % (fig_prefix, ind, fi,
                                                      pol[pi])
                    if tag_output_iter:
                        fig_name = output_path(fig_name,
                                               iteration=self.iteration)
                    else:
                        fig_name = output_path(fig_name)
                    plt.savefig(fig_name)
                    plt.close()
                    # plot N
                    N = Vmat - V0 - S
                    plt.figure(figsize=(13, 5))
                    plt.subplot(121)
                    plt.imshow(N.real,
                               aspect='equal',
                               origin='lower',
                               interpolation='nearest')
                    plt.colorbar(shrink=1.0)
                    plt.subplot(122)
                    plt.imshow(N.imag,
                               aspect='equal',
                               origin='lower',
                               interpolation='nearest')
                    plt.colorbar(shrink=1.0)
                    fig_name = '%s_N_%d_%d_%s.png' % (fig_prefix, ind, fi,
                                                      pol[pi])
                    if tag_output_iter:
                        fig_name = output_path(fig_name,
                                               iteration=self.iteration)
                    else:
                        fig_name = output_path(fig_name)
                    plt.savefig(fig_name)
                    plt.close()

                if apply_gain or save_gain:
                    #                    if mpiutil.rank0:
                    #                        np.save('Vmat',saveVmat)
                    e, U = la.eigh(V0 / Sc[fi], eigvals=(nfeed - 1, nfeed - 1))
                    g = U[:, -1] * e[-1]**0.5
                    if g[0].real < 0:
                        g *= -1.0  # make all g[0] phase 0, instead of pi
                    lGain[ii] = g

                    # plot Gain
                    liplot = max(start_ind, transit_ind - 1)
                    hiplot = min(end_ind, transit_ind + 1 + 1)
                    if plot_figs and ti >= liplot and ti <= hiplot:
                        #                    if ti >= liplot and ti <= hiplot:
                        ind = ti - start_ind
                        plt.figure()
                        plt.plot(feedno, g.real, 'b-', label='real')
                        plt.plot(feedno, g.real, 'bo')
                        plt.plot(feedno, g.imag, 'g-', label='imag')
                        plt.plot(feedno, g.imag, 'go')
                        plt.plot(feedno, np.abs(g), 'r-', label='abs')
                        plt.plot(feedno, np.abs(g), 'ro')
                        plt.xlim(feedno[0] - 1, feedno[-1] + 1)
                        yl, yh = plt.ylim()
                        plt.ylim(yl, yh + (yh - yl) / 5)
                        plt.xlabel('Feed number')
                        plt.legend()
                        fig_name = '%s_ants_%d_%d_%s.png' % (fig_prefix, ind,
                                                             fi, pol[pi])
                        print('plot %s' % fig_name)
                        if tag_output_iter:
                            fig_name = output_path(fig_name,
                                                   iteration=self.iteration)
                        else:
                            fig_name = output_path(fig_name)
                        plt.savefig(fig_name)
                        plt.close()

            # subtract the vis of calibrator from self.vis
            if subtract_src:
                nbl = len(bls)
                lv = np.zeros((lsrc_vis.shape[0], nbl), dtype=lsrc_vis.dtype)
                for bi, (fd1, fd2) in enumerate(bls):
                    b1, b2 = feedno.index(fd1), feedno.index(fd2)
                    lv[:, bi] = lsrc_vis[:, b1, b2]
                lv = mpiarray.MPIArray.wrap(lv, axis=0, comm=ts.comm)
                lv = lv.redistribute(axis=1).local_array.reshape(nt, nf, 2, -1)
                if 'ns_on' in ts.iterkeys():
                    lv[ts['ns_on']
                       [start_ind:
                        end_ind]] = 0  # avoid ns_on signal to become nan
                ts.local_vis[start_ind:end_ind, :,
                             pol.index('xx')] -= lv[:, :, 0]
                ts.local_vis[start_ind:end_ind, :,
                             pol.index('yy')] -= lv[:, :, 1]

                del lv

            if not save_src_vis:
                if subtract_src:
                    del lsrc_vis
            else:
                if tag_output_iter:
                    src_vis_file = output_path(src_vis_file,
                                               iteration=self.iteration)
                else:
                    src_vis_file = output_path(src_vis_file)
                # create file and allocate space first by rank0
                if mpiutil.rank0:
                    with h5py.File(src_vis_file, 'w') as f:
                        # allocate space
                        shp = (nt, nf, 2, nfeed, nfeed)
                        f.create_dataset('sky_vis', shp, dtype=lsky_vis.dtype)
                        f.create_dataset('src_vis', shp, dtype=lsrc_vis.dtype)
                        f.create_dataset('outlier_vis',
                                         shp,
                                         dtype=lotl_vis.dtype)
                        f.attrs['calibrator'] = calibrator
                        f.attrs['dim'] = 'time, freq, pol, feed, feed'
                        try:
                            f.attrs['time'] = ts.time[start_ind:end_ind]
                        except RuntimeError:
                            f.create_dataset('time',
                                             data=ts.time[start_ind:end_ind])
                            f.attrs['time'] = '/time'
                        f.attrs['freq'] = freq
                        f.attrs['pol'] = np.array(['xx', 'yy'])
                        f.attrs['feed'] = np.array(feedno)

                mpiutil.barrier()

                # write data to file
                for i in range(10):
                    try:
                        # NOTE: if write simultaneously, will loss data with processes distributed in several nodes
                        for ri in xrange(mpiutil.size):
                            if ri == mpiutil.rank:
                                with h5py.File(src_vis_file, 'r+') as f:
                                    for ii, (ti, fi,
                                             pi) in enumerate(tfp_linds):
                                        ti_ = ti - start_ind
                                        pi_ = gain_pd[pol[pi]]
                                        f['sky_vis'][ti_, fi,
                                                     pi_] = lsky_vis[ii]
                                        f['src_vis'][ti_, fi,
                                                     pi_] = lsrc_vis[ii]
                                        f['outlier_vis'][ti_, fi,
                                                         pi_] = lotl_vis[ii]
                            mpiutil.barrier()
                        break
                    except IOError:
                        time.sleep(0.5)
                        continue
                else:
                    raise RuntimeError('Could not open file: %s...' %
                                       src_vis_file)

                del lsrc_vis
                del lsky_vis
                del lotl_vis

                mpiutil.barrier()

            if apply_gain or save_gain:
                # flag outliers in lGain along each feed
                lG_abs = np.full_like(lGain, np.nan, dtype=lGain.real.dtype)
                for i in range(lGain.shape[0]):
                    valid_inds = np.where(np.isfinite(lGain[i]))[0]
                    if len(valid_inds) > 3:
                        vabs = np.abs(lGain[i, valid_inds])
                        vmed = np.median(vabs)
                        vabs_diff = np.abs(vabs - vmed)
                        vmad = np.median(vabs_diff) / 0.6745
                        if reserve_high_gain:
                            # reserve significantly higher ones, flag only significantly lower ones
                            lG_abs[i, valid_inds] = np.where(
                                vmed - vabs > 3.0 * vmad, np.nan, vabs)
                        else:
                            # flag both significantly higher and lower ones
                            lG_abs[i, valid_inds] = np.where(
                                vabs_diff > 3.0 * vmad, np.nan, vabs)

                # choose data slice near the transit time
                li = max(start_ind, transit_ind - 10) - start_ind
                hi = min(end_ind, transit_ind + 10 + 1) - start_ind
                # compute s_top for this time range
                n0 = np.zeros(((hi - li), 3))
                for ti, jt in enumerate(ts.time[start_ind:end_ind][li:hi]):
                    aa.set_jultime(jt)
                    s.compute(aa)
                    n0[ti] = s.get_crds('top', ncrd=3)
                if save_phs_change:
                    n0t = np.zeros((nt, 3))
                    for ti, jt in enumerate(ts.time[start_ind:end_ind]):
                        aa.set_jultime(jt)
                        s.compute(aa)
                        n0t[ti] = s.get_crds('top', ncrd=3)

                # get the positions of feeds
                feedpos = ts['feedpos'][:]

                # wrap and redistribute Gain and flagged G_abs
                Gain = mpiarray.MPIArray.wrap(lGain, axis=0, comm=ts.comm)
                Gain = Gain.redistribute(axis=1).reshape(
                    nt, nf, 2, None).redistribute(axis=0).reshape(
                        None, nf * 2 * nfeed).redistribute(axis=1)
                G_abs = mpiarray.MPIArray.wrap(lG_abs, axis=0, comm=ts.comm)
                G_abs = G_abs.redistribute(axis=1).reshape(
                    nt, nf, 2, None).redistribute(axis=0).reshape(
                        None, nf * 2 * nfeed).redistribute(axis=1)

                fpd_inds = list(
                    itertools.product(range(nf), range(2),
                                      range(nfeed)))  # only for xx and yy
                fpd_linds = mpiutil.mpilist(fpd_inds,
                                            method='con',
                                            comm=ts.comm)
                del fpd_inds
                # create data to save the solved gain for each feed
                lgain = np.full((len(fpd_linds), ), cnan,
                                dtype=Gain.dtype)  # gain for each feed
                if save_phs_change:
                    lphs = np.full((nt, len(fpd_linds)),
                                   np.nan,
                                   dtype=Gain.real.dtype
                                   )  # phase change with time for each feed

                # check for conj
                num_conj = 0
                for ii, (fi, pi, di) in enumerate(fpd_linds):
                    y = G_abs.local_array[li:hi, ii]
                    inds = np.where(np.isfinite(y))[0]
                    if len(inds) >= max(4, 0.5 * len(y)):
                        # get the approximate magnitude by averaging the central G_abs
                        # solve phase by least square fit
                        ui = (feedpos[di] - feedpos[0]) * (
                            1.0e6 * freq[fi]
                        ) / const.c  # position of this feed (relative to the first feed) in unit of wavelength
                        exp_factor = np.exp(2.0J * np.pi * np.dot(n0, ui))
                        ef = exp_factor
                        Gi = Gain.local_array[li:hi, ii]
                        e_phs = np.dot(ef[inds].conj(),
                                       Gi[inds] / y[inds]) / len(inds)
                        ea = np.abs(e_phs)
                        e_phs_conj = np.dot(ef[inds],
                                            Gi[inds] / y[inds]) / len(inds)
                        eac = np.abs(e_phs_conj)
                        if eac > ea:
                            num_conj += 1
                # reduce num_conj from all processes
                num_conj = mpiutil.allreduce(num_conj, comm=ts.comm)
                if num_conj > 0.5 * (nf * 2 * nfeed):  # 2 for 2 pols
                    if mpiutil.rank0:
                        print '!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!'
                        print '!!!   Detect data should be their conjugate...   !!!'
                        print '!!!   Correct it automatically...                !!!'
                        print '!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!'
                    mpiutil.barrier()
                    # correct vis
                    ts.local_vis[:] = ts.local_vis.conj()
                    # correct G
                    Gain.local_array[:] = Gain.local_array.conj()

                # solve for gain
#                feedplotG = []
#                savenum = 0
                for ii, (fi, pi, di) in enumerate(fpd_linds):
                    y = G_abs.local_array[li:hi, ii]
                    inds = np.where(np.isfinite(y))[0]

                    if len(inds) >= max(4, 0.5 * len(y)):
                        # get the approximate magnitude by averaging the central G_abs
                        mag = np.mean(y[inds])
                        # solve phase by least square fit
                        ui = (feedpos[di] - feedpos[0]) * (
                            1.0e6 * freq[fi]
                        ) / const.c  # position of this feed (relative to the first feed) in unit of wavelength
                        exp_factor = np.exp(2.0J * np.pi * np.dot(n0, ui))
                        ef = exp_factor
                        Gi = Gain.local_array[li:hi, ii]

                        #                        feedplotG += [Gi[10]]
                        #                        print(feedplotG)

                        e_phs = np.dot(ef[inds].conj(),
                                       Gi[inds] / y[inds]) / len(inds)
                        ea = np.abs(e_phs)

                        if np.abs(ea - 1.0) < 0.1:
                            # compute gain for this feed
                            lgain[ii] = mag * e_phs
                            if save_phs_change:
                                lphs[:, ii] = np.angle(
                                    np.exp(-2.0J * np.pi * np.dot(n0t, ui)) *
                                    Gain.local_array[:, ii])
                        else:
                            e_phs_conj = np.dot(ef[inds],
                                                Gi[inds] / y[inds]) / len(inds)
                            eac = np.abs(e_phs_conj)
                            if eac > ea:
                                if np.abs(eac - 1.0) < 0.01:
                                    print 'feedno = %d, fi = %d, pol = %s: may need to be conjugated' % (
                                        feedno[di], fi, gain_pd[pi])
                            else:
                                print 'feedno = %d, fi = %d, pol = %s: maybe wrong abs(e_phs): %s' % (
                                    feedno[di], fi, gain_pd[pi], ea)

#                import matplotlib.pyplot as plt
#                import os.path
#                feedplotG = np.array(feedplotG)
#                if mpiutil.rank0:
#                    if os.path.isfile('feedplotG%d.npy'%savenum):
#                        savenum += 1
#                    np.save('feedplotG%d'%savenum,feedplotG)
#                    plt.plot(feedplotG)
#                    plt.savefig('Gfig%d'%savenum)

# gather local gain
                gain = mpiutil.gather_array(lgain,
                                            axis=0,
                                            root=None,
                                            comm=ts.comm)
                del lgain
                gain = gain.reshape(nf, 2, nfeed)
                if save_phs_change:
                    phs = mpiutil.gather_array(lphs,
                                               axis=1,
                                               root=0,
                                               comm=ts.comm)
                    del lphs
                    if mpiutil.rank0:
                        phs = phs.reshape(nt, nf, 2, nfeed)

                # apply gain to vis
                if apply_gain:
                    for fi in range(nf):
                        for pi in [pol.index('xx'), pol.index('yy')]:
                            pi_ = gain_pd[pol[pi]]
                            for bi, (fd1, fd2) in enumerate(
                                    ts['blorder'].local_data):
                                g1 = gain[fi, pi_, feedno.index(fd1)]
                                g2 = gain[fi, pi_, feedno.index(fd2)]
                                if np.isfinite(g1) and np.isfinite(g2):
                                    ts.local_vis[:, fi, pi,
                                                 bi] /= (g1 * np.conj(g2))
                                else:
                                    # mask the un-calibrated vis
                                    ts.local_vis_mask[:, fi, pi, bi] = True

                # save gain to file
                if save_gain:
                    if tag_output_iter:
                        gain_file = output_path(gain_file,
                                                iteration=self.iteration)
                    else:
                        gain_file = output_path(gain_file)
                    if mpiutil.rank0:
                        with h5py.File(gain_file, 'w') as f:
                            # allocate space for Gain
                            dset = f.create_dataset('Gain', (nt, nf, 2, nfeed),
                                                    dtype=Gain.dtype)
                            dset.attrs['calibrator'] = calibrator
                            dset.attrs['dim'] = 'time, freq, pol, feed'
                            try:
                                dset.attrs['time'] = ts.time[start_ind:end_ind]
                            except RuntimeError:
                                f.create_dataset(
                                    'time', data=ts.time[start_ind:end_ind])
                                dset.attrs['time'] = '/time'
                            dset.attrs['freq'] = freq
                            dset.attrs['pol'] = np.array(['xx', 'yy'])
                            dset.attrs['feed'] = np.array(feedno)
                            # save gain
                            dset = f.create_dataset('gain', data=gain)
                            dset.attrs['calibrator'] = calibrator
                            dset.attrs['dim'] = 'freq, pol, feed'
                            dset.attrs['freq'] = freq
                            dset.attrs['pol'] = np.array(['xx', 'yy'])
                            dset.attrs['feed'] = np.array(feedno)
                            # save phs

                            if save_phs_change:
                                f.create_dataset('phs', data=phs)
                            # save transit index and transit time for the case do the ps cal first
                            f.attrs['transit_index'] = transit_inds[0]
                            #                            f.attrs['transit_time'] = ts['jul_date'][transit_inds[0]]
                            f.attrs['transit_jul'] = ts['jul_date'][
                                transit_inds[0]]
                            f.attrs['transit_time'] = ts['sec1970'][
                                transit_inds[0]]
                            #                            f.attrs['transit_time'] = ts.attrs['sec1970'] + ts.attrs['inttime']*transit_inds[0] # in sec1970 to improve numeric precision

                            if os.path.exists(output_path(
                                    ts.ns_gain_file)) and not ts.ps_first:
                                with h5py.File(output_path(ts.ns_gain_file),
                                               'r+') as ns_file:
                                    phs_only = not ('ns_cal_amp'
                                                    in ns_file.keys())
                                    #                                    exclude_bad = 'badchn' in ns_file['channo'].attrs.keys()
                                    new_gain = uni_gain(f,
                                                        ns_file,
                                                        phs_only=phs_only)
                                    ns_file.create_dataset('uni_gain',
                                                           data=new_gain)
                                    ns_file['uni_gain'].attrs[
                                        'dim'] = '(time, freq, bl)'

                    mpiutil.barrier()

                    # save Gain
                    for i in range(10):
                        try:
                            # NOTE: if write simultaneously, will loss data with processes distributed in several nodes
                            for ri in xrange(mpiutil.size):
                                if ri == mpiutil.rank:
                                    with h5py.File(gain_file, 'r+') as f:
                                        for ii, (ti, fi,
                                                 pi) in enumerate(tfp_linds):
                                            ti_ = ti - start_ind
                                            pi_ = gain_pd[pol[pi]]
                                            f['Gain'][ti_, fi, pi_] = lGain[ii]
                                mpiutil.barrier()
                            break
                        except IOError:
                            time.sleep(0.5)
                            continue
                    else:
                        raise RuntimeError('Could not open file: %s...' %
                                           gain_file)

                    mpiutil.barrier()

        # convert vis from intensity unit to temperature unit in K
        if temperature_convert:
            if 'unit' in ts.vis.attrs.keys() and ts.vis.attrs['unit'] == 'K':
                if mpiutil.rank0:
                    print 'vis is already in unit K, do nothing...'
            else:
                factor = 1.0e-26 * (const.c**2 / (2 * const.k_B *
                                                  (1.0e6 * freq)**2)
                                    )  # NOTE: 1Jy = 1.0e-26 W m^-2 Hz^-1
                ts.local_vis[:] *= factor[np.newaxis, :, np.newaxis,
                                          np.newaxis]
                ts.vis.attrs['unit'] = 'K'

        return super(PsCal, self).process(ts)
Esempio n. 10
0
    def load_tod_excl_main_data(self):
        """Load time ordered attributes and datasets (exclude the main data) from all files."""

        super(TimestreamCommon, self).load_tod_excl_main_data()

        if 'sec1970' not in self.iterkeys():
            # generate sec1970
            int_time = self.infiles[0].attrs['inttime']
            sec1970s = []
            nts = []
            for fh in mpiutil.mpilist(self.infiles, method='con', comm=self.comm):
                sec1970s.append(fh.attrs['sec1970'])
                nts.append(fh[self.main_data_name].shape[0])
            sec1970 = np.zeros(sum(nts), dtype=np.float64) # precision float32 is not enough
            cum_nts = np.cumsum([0] + nts)
            for idx, (nt, sec) in enumerate(zip(nts, sec1970s)):
                sec1970[cum_nts[idx]:cum_nts[idx+1]] = np.array([ sec + i*int_time for i in xrange(nt)], dtype=np.float64) # precision float32 is not enough
            # gather local sec1970
            sec1970 = mpiutil.gather_array(sec1970, root=None, comm=self.comm)
            # select the corresponding section
            sec1970 = sec1970[self.main_data_start:self.main_data_stop][self.main_data_select[0]]

            # if time is just the distributed axis, load sec1970 distributed
            if 'time' == self.main_data_axes[self.main_data_dist_axis]:
                sec1970 = mpiarray.MPIArray.from_numpy_array(sec1970)
            self.create_main_time_ordered_dataset('sec1970', data=sec1970)
            # create attrs of this dset
            self['sec1970'].attrs["unit"] = 'second'
            # determine if it is continuous in time
            sec_diff = np.diff(sec1970)
            break_inds = np.where(sec_diff>1.5*int_time)[0]
            if len(break_inds) > 0:
                self['sec1970'].attrs["continuous"] = False
                self['sec1970'].attrs["break_inds"] = break_inds + 1
            else:
                self['sec1970'].attrs["continuous"] = True

            # generate julian date
            jul_date = np.array([ date_util.get_juldate(datetime.fromtimestamp(s), tzone=self.infiles[0].attrs['timezone']) for s in sec1970 ], dtype=np.float64) # precision float32 is not enough
            if 'time' == self.main_data_axes[self.main_data_dist_axis]:
                jul_date = mpiarray.MPIArray.wrap(jul_date, axis=0)
            # if time is just the distributed axis, load jul_date distributed
            self.create_main_time_ordered_dataset('jul_date', data=jul_date)
            # create attrs of this dset
            self['jul_date'].attrs["unit"] = 'day'

            # generate local time in hour from 0 to 24.0
            def _hour(t):
                return t.hour + t.minute/60.0 + t.second/3600.0 + t.microsecond/3.6e8
            local_hour = np.array([ _hour(datetime.fromtimestamp(s).time()) for s in sec1970 ], dtype=np.float64)
            if 'time' == self.main_data_axes[self.main_data_dist_axis]:
                local_hour = mpiarray.MPIArray.wrap(local_hour, axis=0)
            # if time is just the distributed axis, load local_hour distributed
            self.create_main_time_ordered_dataset('local_hour', data=local_hour)
            # create attrs of this dset
            self['local_hour'].attrs["unit"] = 'hour'

            # generate az, alt
            az_alt = np.zeros((self['sec1970'].local_data.shape[0], 2), dtype=np.float32) # radians
            if self.is_dish:
                # antpointing = rt['antpointing'][-1, :, :] # degree
                # pointingtime = rt['pointingtime'][-1, :, :] # degree
                az_alt[:, 0] = 0.0 # az
                az_alt[:, 1] = np.pi/2 # alt
            elif self.is_cylinder:
                az_alt[:, 0] = np.pi/2 # az
                az_alt[:, 1] = np.pi/2 # alt
            else:
                raise RuntimeError('Unknown antenna type %s' % self.attrs['telescope'])

            # generate ra, dec of the antenna pointing
            aa = self.array
            ra_dec = np.zeros_like(az_alt) # radians
            for ti in xrange(az_alt.shape[0]):
                az, alt = az_alt[ti]
                az, alt = ephem.degrees(az), ephem.degrees(alt)
                aa.set_jultime(self['jul_date'].local_data[ti])
                ra_dec[ti] = aa.radec_of(az, alt) # in radians, a point in the sky above the observer

            if self.main_data_dist_axis == 0:
                az_alt = mpiarray.MPIArray.wrap(az_alt, axis=0)
                ra_dec = mpiarray.MPIArray.wrap(ra_dec, axis=0)
            # if time is just the distributed axis, create distributed datasets
            self.create_main_time_ordered_dataset('az_alt', data=az_alt)
            self['az_alt'].attrs['unit'] = 'radian'
            self.create_main_time_ordered_dataset('ra_dec', data=ra_dec)
            self['ra_dec'].attrs['unit'] = 'radian'

            # determin if it is the same pointing
            if self.main_data_dist_axis == 0:
                az_alt = az_alt.local_array
                ra_dec = ra_dec.local_array
            # gather local az_alt
            az_alt = mpiutil.gather_array(az_alt, root=None, comm=self.comm)
            if np.allclose(az_alt[:, 0], az_alt[0, 0]) and np.allclose(az_alt[:, 1], az_alt[0, 1]):
                self['az_alt'].attrs['same_pointing'] = True
            else:
                self['az_alt'].attrs['same_pointing'] = False
            # determin if it is the same dec
            # gather local ra_dec
            ra_dec = mpiutil.gather_array(ra_dec, root=None, comm=self.comm)
            if np.allclose(ra_dec[:, 1], ra_dec[0, 1]):
                self['ra_dec'].attrs['same_dec'] = True
            else:
                self['ra_dec'].attrs['same_dec'] = False
Esempio n. 11
0
    def load_tod_excl_main_data(self):
        """Load time ordered attributes and datasets (exclude the main data) from all files."""

        super(TimestreamCommon, self).load_tod_excl_main_data()

        if 'sec1970' not in self.iterkeys():
            # generate sec1970
            int_time = self.infiles[0].attrs['inttime']
            sec1970s = []
            nts = []
            for fh in mpiutil.mpilist(self.infiles, method='con', comm=self.comm):
                sec1970s.append(fh.attrs['sec1970'])
                nts.append(fh[self.main_data_name].shape[0])
            sec1970 = np.zeros(sum(nts), dtype=np.float64) # precision float32 is not enough
            cum_nts = np.cumsum([0] + nts)
            for idx, (nt, sec) in enumerate(zip(nts, sec1970s)):
                sec1970[cum_nts[idx]:cum_nts[idx+1]] = np.array([ sec + i*int_time for i in xrange(nt)], dtype=np.float64) # precision float32 is not enough
            # gather local sec1970
            sec1970 = mpiutil.gather_array(sec1970, root=None, comm=self.comm)
            # select the corresponding section
            sec1970 = sec1970[self.main_data_start:self.main_data_stop][self.main_data_select[0]]

            # if time is just the distributed axis, load sec1970 distributed
            if 'time' == self.main_data_axes[self.main_data_dist_axis]:
                sec1970 = mpiarray.MPIArray.from_numpy_array(sec1970)
            self.create_main_time_ordered_dataset('sec1970', data=sec1970)
            # create attrs of this dset
            self['sec1970'].attrs["unit"] = 'second'
            # determine if it is continuous in time
            sec_diff = np.diff(sec1970)
            break_inds = np.where(sec_diff>1.5*int_time)[0]
            if len(break_inds) > 0:
                self['sec1970'].attrs["continuous"] = False
                self['sec1970'].attrs["break_inds"] = break_inds + 1
            else:
                self['sec1970'].attrs["continuous"] = True

            # generate julian date
            jul_date = np.array([ date_util.get_juldate(datetime.fromtimestamp(s), tzone=self.infiles[0].attrs['timezone']) for s in sec1970 ], dtype=np.float64) # precision float32 is not enough
            if 'time' == self.main_data_axes[self.main_data_dist_axis]:
                jul_date = mpiarray.MPIArray.wrap(jul_date, axis=0)
            # if time is just the distributed axis, load jul_date distributed
            self.create_main_time_ordered_dataset('jul_date', data=jul_date)
            # create attrs of this dset
            self['jul_date'].attrs["unit"] = 'day'

            # generate local time in hour from 0 to 24.0
            def _hour(t):
                return t.hour + t.minute/60.0 + t.second/3600.0 + t.microsecond/3.6e8
            local_hour = np.array([ _hour(datetime.fromtimestamp(s).time()) for s in sec1970 ], dtype=np.float64)
            if 'time' == self.main_data_axes[self.main_data_dist_axis]:
                local_hour = mpiarray.MPIArray.wrap(local_hour, axis=0)
            # if time is just the distributed axis, load local_hour distributed
            self.create_main_time_ordered_dataset('local_hour', data=local_hour)
            # create attrs of this dset
            self['local_hour'].attrs["unit"] = 'hour'

            # generate az, alt
            az_alt = np.zeros((self['sec1970'].local_data.shape[0], 2), dtype=np.float32) # radians
            if self.is_dish:
                # antpointing = rt['antpointing'][-1, :, :] # degree
                # pointingtime = rt['pointingtime'][-1, :, :] # degree
                az_alt[:, 0] = 0.0 # az
                az_alt[:, 1] = np.pi/2 # alt
            elif self.is_cylinder:
                az_alt[:, 0] = np.pi/2 # az
                az_alt[:, 1] = np.pi/2 # alt
            else:
                raise RuntimeError('Unknown antenna type %s' % self.attrs['telescope'])

            # generate ra, dec of the antenna pointing
            aa = self.array
            ra_dec = np.zeros_like(az_alt) # radians
            for ti in xrange(az_alt.shape[0]):
                az, alt = az_alt[ti]
                az, alt = ephem.degrees(az), ephem.degrees(alt)
                aa.set_jultime(self['jul_date'].local_data[ti])
                ra_dec[ti] = aa.radec_of(az, alt) # in radians, a point in the sky above the observer

            if self.main_data_dist_axis == 0:
                az_alt = mpiarray.MPIArray.wrap(az_alt, axis=0)
                ra_dec = mpiarray.MPIArray.wrap(ra_dec, axis=0)
            # if time is just the distributed axis, create distributed datasets
            self.create_main_time_ordered_dataset('az_alt', data=az_alt)
            self['az_alt'].attrs['unit'] = 'radian'
            self.create_main_time_ordered_dataset('ra_dec', data=ra_dec)
            self['ra_dec'].attrs['unit'] = 'radian'

            # determin if it is the same pointing
            if self.main_data_dist_axis == 0:
                az_alt = az_alt.local_array
                ra_dec = ra_dec.local_array
            # gather local az_alt
            az_alt = mpiutil.gather_array(az_alt, root=None, comm=self.comm)
            if np.allclose(az_alt[:, 0], az_alt[0, 0]) and np.allclose(az_alt[:, 1], az_alt[0, 1]):
                self['az_alt'].attrs['same_pointing'] = True
            else:
                self['az_alt'].attrs['same_pointing'] = False
            # determin if it is the same dec
            # gather local ra_dec
            ra_dec = mpiutil.gather_array(ra_dec, root=None, comm=self.comm)
            if np.allclose(ra_dec[:, 1], ra_dec[0, 1]):
                self['ra_dec'].attrs['same_dec'] = True
            else:
                self['ra_dec'].attrs['same_dec'] = False