spectra.py

import copy

import numpy as np
import scipy.signal
import psr_utils

class Spectra(object):
    """A class to store spectra. This is mainly to provide
        reusable functionality.
    """
    def __init__(self, freqs, dt, data, starttime=0, dm=0):
        """Spectra constructor.
            
            Inputs:
                freqs: Observing frequencies for each channel.
                dt: Sample time (in seconds).
                data: A 2D numpy array containing pulsar data.
                        Axis 0 should contain channels. (e.g. data[0,:])
                        Axis 1 should contain spectra. (e.g. data[:,0])
                starttime: Start time (in seconds) of the spectra
                        with respect to the start of the observation.
                        (Default: 0).
                dm: Dispersion measure (in pc/cm^3). (Default: 0)

            Output:
                spectra_obj: Spectrum object.
        """
        self.numchans, self.numspectra = data.shape
        assert len(freqs)==self.numchans

        self.freqs = freqs
        self.data = data.astype('float')
        self.dt = dt
        self.starttime = starttime
        self.dm = 0

    def __str__(self):
        return str(self.data)

    def __getitem__(self, key):
        return self.data[key]
    
    def __setitem__(self, key, value):
        self.data[key] = value
    
    def get_chan(self, channum):
        return self.data[channum,:]

    def get_spectrum(self, specnum):
        return self.data[:,specnum]

    def shift_channels(self, bins, padval=0):
        """Shift each channel to the left by the corresponding 
            value in bins, an array.

            Inputs:
                bins: An array containing the number of bins
                    to shift each channel by.
                padval: Value to use when shifting near the edge
                    of a channel. This can be a numeric value,
                    'median', 'mean', or 'rotate'. 
                    
                    The values 'median' and 'mean' refer to the 
                    median and mean of the channel. The value 
                    'rotate' takes values from one end of the 
                    channel and shifts them to the other.

            Outputs:
                None

            *** Shifting happens in-place ***
        """
        assert self.numchans == len(bins)
        for ii in range(self.numchans):
            chan = self.get_chan(ii)
            # Use 'chan[:]' so update happens in-place
            # this way the change effects self.data
            chan[:] = psr_utils.rotate(chan, bins[ii])
            if padval!='rotate':
                # Get padding value
                if padval=='mean':
                    pad = np.mean(chan)
                elif padval=='median':
                    pad = np.median(chan)
                else:
                    pad = padval
                
                # Replace rotated values with padval
                if bins[ii]>0:
                    chan[-bins[ii]:] = pad
                elif bins[ii]<0:
                    chan[:-bins[ii]] = pad

    def subband(self, nsub, subdm=None, padval=0):
        """Reduce the number of channels to 'nsub' by subbanding.
            The channels within a subband are combined using the
            DM 'subdm'. 'padval' is passed to the call to
            'Spectra.shift_channels'.

            Inputs:
                nsub: Number of subbands. Must be a factor of 
                    the number of channels.
                subdm: The DM with which to combine channels within
                    each subband (Default: don't shift channels 
                    within each subband)
                padval: The padding value to use when shifting
                    channels during dedispersion. See documentation
                    of Spectra.shift_channels. (Default: 0)

            Outputs:
                None

            *** Subbanding happens in-place ***
        """
        assert (self.numchans % nsub) == 0
        assert (subdm is None) or (subdm >= 0)
        nchan_per_sub = self.numchans/nsub
        sub_hifreqs = self.freqs[np.arange(nsub)*nchan_per_sub]
        sub_lofreqs = self.freqs[(1+np.arange(nsub))*nchan_per_sub-1]
        sub_ctrfreqs = 0.5*(sub_hifreqs+sub_lofreqs)
        
        if subdm is not None:
            # Compute delays
            ref_delays = psr_utils.delay_from_DM(subdm-self.dm, sub_ctrfreqs)
            delays = psr_utils.delay_from_DM(subdm-self.dm, self.freqs)
            rel_delays = delays-ref_delays.repeat(nchan_per_sub) # Relative delay
            rel_bindelays = np.round(rel_delays/self.dt).astype('int')
            # Shift channels
            self.shift_channels(rel_bindelays, padval)

        # Subband
        self.data = np.array([np.sum(sub, axis=0) for sub in \
                                np.vsplit(self.data, nsub)])
        self.freqs = sub_ctrfreqs
        self.numchans = nsub

    def scaled(self, indep=False):
        """Return a scaled version of the Spectra object.
            When scaling subtract the median from each channel,
            and divide by global std deviation (if indep==False), or
            divide by std deviation of each row (if indep==True).

            Input:
                indep: Boolean. If True, scale each row
                    independantly (Default: False).

            Output:
                scaled_spectra: A scaled version of the
                    Spectra object.
        """
        other = copy.deepcopy(self)
        if not indep:
            std = other.data.std()
        for ii in range(other.numchans):
            chan = other.get_chan(ii)
            median = np.median(chan)
            if indep:
                std = chan.std()
            chan[:] = (chan-median)/std
        return other
    
    def scaled2(self, indep=False):
        """Return a scaled version of the Spectra object.
            When scaling subtract the min from each channel,
            and divide by global max (if indep==False), or
            divide by max of each row (if indep==True).

            Input:
                indep: Boolean. If True, scale each row
                    independantly (Default: False).

            Output:
                scaled_spectra: A scaled version of the
                    Spectra object.
        """
        other = copy.deepcopy(self)
        if not indep:
            max = other.data.max()
        for ii in range(other.numchans):
            chan = other.get_chan(ii)
            min = chan.min()
            if indep:
                max = chan.max()
            chan[:] = (chan-min)/max
        return other

    def masked(self, mask, maskval='median-mid80'):
        """Replace masked data with 'maskval'. Returns
            a masked copy of the Spectra object.
            
            Inputs:
                mask: An array of boolean values of the same size and shape
                    as self.data. True represents an entry to be masked.
                maskval: Value to use when masking. This can be a numeric
                    value, 'median', 'mean', or 'median-mid80'.

                    The values 'median' and 'mean' refer to the median and
                    mean of the channel, respectively. The value 'median-mid80'
                    refers to the median of the channel after the top and bottom
                    10% of the sorted channel is removed.
                    (Default: 'median-mid80')

            Output:
                maskedspec: A masked version of the Spectra object.
        """
        assert self.data.shape == mask.shape
        maskvals = np.ones(self.numchans)
        for ii in range(self.numchans):
            chan = self.get_chan(ii)
            # Use 'chan[:]' so update happens in-place
            if maskval=='mean':
                maskvals[ii]=np.mean(chan)
            elif maskval=='median':
                maskvals[ii]=np.median(chan)
            elif maskval=='median-mid80':
                n = int(np.round(0.1*self.numspectra))
                maskvals[ii]=np.median(sorted(chan)[n:-n])
            else:
                maskvals[ii]=maskval
            if np.all(mask[ii]):
                self.data[ii] = np.ones_like(self.data[ii])*(maskvals[:,np.newaxis][ii])
        return self

    def dedisperse(self, dm=0, padval=0):
        """Shift channels according to the delays predicted by
            the given DM.
            
            Inputs:
                dm: The DM (in pc/cm^3) to use.
                padval: The padding value to use when shifting
                    channels during dedispersion. See documentation
                    of Spectra.shift_channels. (Default: 0)

            Outputs:
                None

            *** Dedispersion happens in place ***
        """
        assert dm >= 0
        ref_delay = psr_utils.delay_from_DM(dm-self.dm, np.max(self.freqs))
        delays = psr_utils.delay_from_DM(dm-self.dm, self.freqs)
        rel_delays = delays-ref_delay # Relative delay
        rel_bindelays = np.round(rel_delays/self.dt).astype('int')
        # Shift channels
        self.shift_channels(rel_bindelays, padval)

        self.dm=dm

    def smooth(self, width=1, padval=0):
        """Smooth each channel by convolving with a top hat
            of given width. The height of the top had is
            chosen shuch that RMS=1 after smoothing. 
            Overlap values are determined by 'padval'.

            Inputs:
                width: Number of bins to smooth by (Default: no smoothing)
                padval: Padding value to use. Possible values are
                    float-value, 'mean', 'median', 'wrap'.
                    (Default: 0).

            Ouputs:
                None

            This bit of code is taken from Scott Ransom's
            PRESTO's single_pulse_search.py (line ~ 423).
            
            *** Smoothing is done in place. ***
        """
        if width > 1:
            kernel = np.ones(width, dtype='float32')/np.sqrt(width)
            for ii in range(self.numchans):
                chan = self.get_chan(ii)
                if padval=='wrap':
                    tosmooth = np.concatenate([chan[-width:], \
                                chan, chan[:width]])
                elif padval=='mean':
                    tosmooth = np.ones(self.numspectra+width*2) * \
                                np.mean(chan)
                    tosmooth[width:-width] = chan
                elif padval=='median':
                    tosmooth = np.ones(self.numspectra+width*2) * \
                                np.median(chan)
                    tosmooth[width:-width] = chan
                else: # padval is a float
                    tosmooth = np.ones(self.numspectra+width*2) * \
                                padval
                    tosmooth[width:-width] = chan
                    
                smoothed = scipy.signal.convolve(tosmooth, kernel, 'same')
                chan[:] = smoothed[width:-width]
                    
    def trim(self, bins=0):
        """Trim the end of the data by 'bins' spectra.
            
            Input:
                bins: Number of spectra to trim off the end of the observation.
                    If bins is negative trim spectra off the beginning of the
                    observation.

            Outputs:
                None

            *** Trimming is irreversible ***
        """
        assert bins < self.numspectra
        if bins == 0:
            return
        elif bins > 0:
            self.data = self.data[:,:-bins]
            self.numspectra = self.numspectra-bins
        elif bins < 0:
            self.data = self.data[:,bins:]
            self.numspectra = self.numspectra-bins
            self.starttime = self.starttime+bins*self.dt

    def downsample(self, factor=1, trim=True):
        """Downsample (in-place) the spectra by co-adding
            'factor' adjacent bins.

            Inputs:
                factor: Reduce the number of spectra by this
                    factor. Must be a factor of the number of
                    spectra if 'trim' is False.
                trim: Trim off excess bins.

            Ouputs:
                None

            *** Downsampling is done in place ***
        """
        assert trim or not (self.numspectra % factor)
        new_num_spectra = self.numspectra/factor
        num_to_trim = self.numspectra%factor
        self.trim(num_to_trim)
        self.data = np.array(np.column_stack([np.sum(subint, axis=1) for \
                        subint in np.hsplit(self.data,new_num_spectra)]))
        self.numspectra = new_num_spectra
        self.dt = self.dt*factor