コード例 #1
0
ファイル: sigproc.py プロジェクト: epic-astronomy/bifrost
    def on_sequence(self, iseq):
        ihdr = iseq.header
        itensor = ihdr['_tensor']

        axnames = list(itensor['labels'])
        shape = list(itensor['shape'])
        scales = list(itensor['scales'])
        units = list(itensor['units'])
        ndim = len(shape)
        dtype = DataType(itensor['dtype'])

        sigproc_hdr = {}
        _copy_item_if_exists(sigproc_hdr, ihdr, 'source_name')
        _copy_item_if_exists(sigproc_hdr, ihdr, 'rawdatafile')
        _copy_item_if_exists(sigproc_hdr, ihdr, 'az_start')
        _copy_item_if_exists(sigproc_hdr, ihdr, 'za_start')
        _copy_item_if_exists(sigproc_hdr, ihdr, 'raj', 'src_raj')
        _copy_item_if_exists(sigproc_hdr, ihdr, 'dej', 'src_dej')
        if 'telescope' in ihdr:
            sigproc_hdr['telescope_id'] = sigproc.telescope2id(
                ihdr['telescope'])
        if 'machine' in ihdr:
            sigproc_hdr['machine_id'] = sigproc.machine2id(ihdr['machine'])
        _copy_item_if_exists(sigproc_hdr, ihdr, 'telescope_id')
        _copy_item_if_exists(sigproc_hdr, ihdr, 'machine_id')
        _copy_item_if_exists(sigproc_hdr, ihdr, 'ibeam')
        _copy_item_if_exists(sigproc_hdr, ihdr, 'nbeams')
        sigproc_hdr['nbits'] = dtype.itemsize_bits
        _copy_item_if_exists(sigproc_hdr, ihdr, 'barycentric')
        _copy_item_if_exists(sigproc_hdr, ihdr, 'pulsarcentric')
        if dtype.is_integer and dtype.is_signed:
            sigproc_hdr['signed'] = True
        if 'coord_frame' in ihdr:
            coord_frame = ihdr['coord_frame']
        else:
            coord_frame = None
        sigproc_hdr['pulsarcentric'] = (coord_frame == 'pulsarcentric')
        sigproc_hdr['barycentric'] = (coord_frame == 'barycentric')

        filename = os.path.join(self.path, ihdr['name'])

        if ndim >= 3 and axnames[-3:] == ['time', 'pol', 'freq']:
            self.data_format = 'filterbank'
            assert (dtype.is_real)
            sigproc_hdr['data_type'] = 1
            sigproc_hdr['nifs'] = shape[-2]
            sigproc_hdr['nchans'] = shape[-1]
            sigproc_hdr['tstart'] = _unix2mjd(scales[-3][0])
            sigproc_hdr['tsamp'] = convert_units(scales[-3][1], units[-3], 's')
            sigproc_hdr['fch1'] = convert_units(scales[-1][0], units[-1],
                                                'MHz')
            sigproc_hdr['foff'] = convert_units(scales[-1][1], units[-1],
                                                'MHz')
            if 'refdm' in ihdr:
                sigproc_hdr['refdm'] = convert_units(ihdr['refdm'],
                                                     ihdr['refdm_units'],
                                                     'pc cm^-3')
            if ndim == 3:
                filename += '.fil'
                self.ofile = open(filename, 'wb')
                sigproc.write_header(sigproc_hdr, self.ofile)
            elif ndim == 4:
                if axnames[-4] != 'beam':
                    raise ValueError("Expected first axis to be 'beam'"
                                     " got '%s'" % axnames[-4])
                nbeam = shape[-4]
                sigproc_hdr['nbeams'] = nbeam
                filenames = [
                    filename + '.%06iof.%06i.fil' % (b + 1, nbeam)
                    for b in range(nbeam)
                ]
                self.ofiles = [open(fname, 'wb') for fname in filenames]
                for b in range(nbeam):
                    sigproc_hdr['ibeam'] = b
                    sigproc.write_header(sigproc_hdr, self.ofiles[b])
            else:
                raise ValueError("Too many dimensions")

        elif ndim >= 2 and 'time' in axnames and 'pol' in axnames:
            pol_axis = axnames.index('pol')
            if pol_axis != ndim - 1:
                # Need to move pol axis
                # Note: We support this because it tends to be convenient
                #         for rest of the pipeline to operate with pol being
                #         the first dim, and doing the transpose on the fly
                #         inside this block is unlikely to cost much relative
                #         to disk perf (and it's free if npol==1).
                axnames.append(axnames[pol_axis])
                del axnames[pol_axis]
                shape.append(shape[pol_axis])
                del shape[pol_axis]
                scales.append(scales[pol_axis])
                del scales[pol_axis]
                units.append(units[pol_axis])
                del units[pol_axis]
            self.pol_axis = pol_axis
            self.data_format = 'timeseries'
            assert (dtype.is_real)
            sigproc_hdr['data_type'] = 2
            sigproc_hdr['nchans'] = 1
            sigproc_hdr['nifs'] = shape[-2]
            sigproc_hdr['tstart'] = _unix2mjd(scales[-2][0])
            sigproc_hdr['tsamp'] = convert_units(scales[-2][1], units[-2], 's')
            if 'cfreq' in ihdr and 'bw' in ihdr:
                sigproc_hdr['fch1'] = convert_units(ihdr['cfreq'],
                                                    ihdr['cfreq_units'], 'MHz')
                sigproc_hdr['foff'] = convert_units(ihdr['bw'],
                                                    ihdr['bw_units'], 'MHz')
            # TODO: Write ndim separate output files, each with its own refdm
            if ndim == 2:
                if 'refdm' in ihdr:
                    sigproc_hdr['refdm'] = convert_units(
                        ihdr['refdm'], ihdr['refdm_units'], 'pc cm^-3')
                filename += '.tim'
                self.ofile = open(filename, 'wb')
                sigproc.write_header(sigproc_hdr, self.ofile)
            elif ndim == 3:
                if axnames[-3] != 'dispersion':
                    raise ValueError("Expected first axis to be 'dispersion'"
                                     " got '%s'" % axnames[-3])
                ndm = shape[-3]
                dm0 = scales[-3][0]
                ddm = scales[-3][1]
                dms = [dm0 + ddm * d for d in range(ndm)]
                dms = [convert_units(dm, units[-3], 'pc cm^-3') for dm in dms]
                filenames = [filename + '.%09.2f.tim' % dm for dm in dms]
                self.ofiles = [open(fname, 'wb') for fname in filenames]
                for d, dm in enumerate(dms):
                    sigproc_hdr['refdm'] = dm
                    sigproc.write_header(sigproc_hdr, self.ofiles[d])
            else:
                raise ValueError("Too many dimensions")

        elif ndim == 4 and axnames[-3:] == ['pol', 'freq', 'phase']:
            self.data_format = 'pulseprofile'
            assert (dtype.is_real)
            sigproc_hdr['data_type'] = 2
            sigproc_hdr['nifs'] = shape[-3]
            sigproc_hdr['nchans'] = shape[-2]
            sigproc_hdr['nbins'] = shape[-1]
            sigproc_hdr['tstart'] = _unix2mjd(scales[-4][0])
            sigproc_hdr['tsamp'] = convert_units(scales[-4][1], units[-4], 's')
            sigproc_hdr['fch1'] = convert_units(scales[-2][0], units[-2],
                                                'MHz')
            sigproc_hdr['foff'] = convert_units(scales[-2][1], units[-2],
                                                'MHz')
            if 'refdm' in ihdr:
                sigproc_hdr['refdm'] = convert_units(ihdr['refdm'],
                                                     ihdr['refdm_units'],
                                                     'pc cm^-3')
            _copy_item_if_exists(sigproc_hdr, ihdr, 'npuls')
            self.filename = filename
            self.sigproc_hdr = sigproc_hdr
            self.t0 = scales[-4][0]
            self.dt = scales[-4][1]

        else:
            raise ValueError(
                "Axis labels do not correspond to a known data format: " +
                str(axnames) + "\nKnown formats are:" +
                "\n  [time, pol, freq]\n  [beam, time, pol]\n" +
                "  [time, pol]\n  [dispersion, time, pol]\n" +
                "  [pol, freq, phase]")
コード例 #2
0
    def on_sequence(self, iseq):
        ihdr = iseq.header
        itensor = ihdr['_tensor']

        axnames = tuple(itensor['labels'])
        shape = itensor['shape']
        scales = itensor['scales']
        units = itensor['units']
        ndim = len(shape)
        dtype = DataType(itensor['dtype'])

        sigproc_hdr = {}
        copy_item_if_exists(sigproc_hdr, ihdr, 'source_name')
        copy_item_if_exists(sigproc_hdr, ihdr, 'rawdatafile')
        copy_item_if_exists(sigproc_hdr, ihdr, 'az_start')
        copy_item_if_exists(sigproc_hdr, ihdr, 'za_start')
        copy_item_if_exists(sigproc_hdr, ihdr, 'raj', 'src_raj')
        copy_item_if_exists(sigproc_hdr, ihdr, 'dej', 'src_dej')
        if 'telescope' in ihdr:
            sigproc_hdr['telescope_id'] = sigproc.telescope2id(
                ihdr['telescope'])
        if 'machine' in ihdr:
            sigproc_hdr['machine_id'] = sigproc.machine2id(ihdr['machine'])
        copy_item_if_exists(sigproc_hdr, ihdr, 'telescope_id')
        copy_item_if_exists(sigproc_hdr, ihdr, 'machine_id')
        copy_item_if_exists(sigproc_hdr, ihdr, 'ibeam')
        copy_item_if_exists(sigproc_hdr, ihdr, 'nbeams')
        sigproc_hdr['nbits'] = dtype.itemsize_bits
        copy_item_if_exists(sigproc_hdr, ihdr, 'barycentric')
        copy_item_if_exists(sigproc_hdr, ihdr, 'pulsarcentric')
        if dtype.is_integer and dtype.is_signed:
            sigproc_hdr['signed'] = True
        if 'coord_frame' in ihdr:
            coord_frame = ihdr['coord_frame']
        else:
            coord_frame = None
        sigproc_hdr['pulsarcentric'] = (coord_frame == 'pulsarcentric')
        sigproc_hdr['barycentric'] = (coord_frame == 'barycentric')

        filename = os.path.join(self.path, ihdr['name'])

        if ndim >= 3 and axnames[-3:] == ('time', 'pol', 'freq'):
            self.data_format = 'filterbank'
            assert (dtype.is_real)
            sigproc_hdr['data_type'] = 1
            sigproc_hdr['nifs'] = shape[-2]
            sigproc_hdr['nchans'] = shape[-1]
            sigproc_hdr['tstart'] = unix2mjd(scales[-3][0])
            sigproc_hdr['tsamp'] = convert_units(scales[-3][1], units[-3], 's')
            sigproc_hdr['fch1'] = convert_units(scales[-1][0], units[-1],
                                                'MHz')
            sigproc_hdr['foff'] = convert_units(scales[-1][1], units[-1],
                                                'MHz')
            if 'refdm' in ihdr:
                sigproc_hdr['refdm'] = convert_units(ihdr['refdm'],
                                                     ihdr['refdm_units'],
                                                     'pc cm^-3')
            if ndim == 3:
                filename += '.fil'
                self.ofile = open(filename, 'wb')
                sigproc.write_header(sigproc_hdr, self.ofile)
            elif ndim == 4:
                if axnames[-4] != 'beam':
                    raise ValueError("Expected first axis to be 'beam'")
                nbeam = shape[-4]
                sigproc_hdr['nbeams'] = nbeam
                filenames = [
                    filename + '.%06iof.%06i.fil' % (b + 1, nbeam)
                    for b in xrange(nbeam)
                ]
                self.ofiles = [open(fname, 'wb') for fname in filenames]
                for b in xrange(nbeam):
                    sigproc_hdr['ibeam'] = b
                    sigproc.write_header(sigproc_hdr, self.ofiles[b])
            else:
                raise ValueError("Too many dimensions")

        elif ndim >= 2 and axnames[-2:] == ('time', 'pol'):
            self.data_format = 'timeseries'
            assert (dtype.is_real)
            sigproc_hdr['data_type'] = 2
            sigproc_hdr['nchans'] = 1
            sigproc_hdr['nifs'] = shape[-2]
            sigproc_hdr['tstart'] = unix2mjd(scales[-2][0])
            sigproc_hdr['tsamp'] = convert_units(scales[-2][1], units[-2], 's')
            if 'cfreq' in ihdr and 'bw' in ihdr:
                sigproc_hdr['fch1'] = convert_units(ihdr['cfreq'],
                                                    ihdr['cfreq_units'], 'MHz')
                sigproc_hdr['foff'] = convert_units(ihdr['bw'],
                                                    ihdr['bw_units'], 'MHz')
            # TODO: Write ndim separate output files, each with its own refdm
            if ndim == 2:
                if 'refdm' in ihdr:
                    sigproc_hdr['refdm'] = convert_units(
                        ihdr['refdm'], ihdr['refdm_units'], 'pc cm^-3')
                filename += '.tim'
                self.ofile = open(filename, 'wb')
                sigproc.write_header(sigproc_hdr, self.ofile)
            elif ndim == 3:
                if axnames[-3] != 'dispersion measure':
                    raise ValueError(
                        "Expected first axis to be 'dispersion measure'")
                ndm = shape[-3]
                dm0 = scales[-3][0]
                ddm = scales[-3][1]
                dms = [dm0 + ddm * d for d in xrange(ndm)]
                dms = [convert_units(dm, units[-3], 'pc cm^-3') for dm in dms]
                filenames = [filename + '.%09.2f.tim' % dm for dm in dms]
                self.ofiles = [open(fname, 'wb') for fname in filenames]
                for d, dm in enumerate(dms):
                    sigproc_hdr['refdm'] = dm
                    sigproc.write_header(sigproc_hdr, self.ofiles[d])
            else:
                raise ValueError("Too many dimensions")

        elif ndim == 4 and axnames[-3:] == ('pol', 'freq', 'phase'):
            self.data_format = 'pulseprofile'
            assert (dtype.is_real)
            sigproc_hdr['data_type'] = 2
            sigproc_hdr['nifs'] = shape[-3]
            sigproc_hdr['nchans'] = shape[-2]
            sigproc_hdr['nbins'] = shape[-1]
            sigproc_hdr['tstart'] = unix2mjd(scales[-4][0])
            sigproc_hdr['tsamp'] = convert_units(scales[-4][1], units[-4], 's')
            sigproc_hdr['fch1'] = convert_units(scales[-2][0], units[-2],
                                                'MHz')
            sigproc_hdr['foff'] = convert_units(scales[-2][1], units[-2],
                                                'MHz')
            if 'refdm' in ihdr:
                sigproc_hdr['refdm'] = convert_units(ihdr['refdm'],
                                                     ihdr['refdm_units'],
                                                     'pc cm^-3')
            copy_item_if_exists(sigproc_hdr, ihdr, 'npuls')
            self.filename = filename
            self.sigproc_hdr = sigproc_hdr
            self.t0 = scales[-4][0]
            self.dt = scales[-4][1]

        else:
            raise ValueError(
                "Axis labels do not correspond to a known data format: " +
                str(axnames))