Example #1
0
class _ProjectFilterBase(task.SingleTask):
    """A base class for projecting data to/from a different basis.

    Attributes
    ----------
    mode : string
        Which projection to perform. Into the new basis (forward), out of the
        new basis (backward), and forward then backward in order to filter the
        data through the basis (filter).
    """

    mode = config.enum(["forward", "backward", "filter"], default="forward")

    def process(self, inp):
        """Project or filter the input data.

        Parameters
        ----------
        inp : memh5.BasicCont
            Data to process.

        Returns
        -------
        output : memh5.BasicCont
        """

        if self.mode == "forward":
            return self._forward(inp)

        if self.mode == "backward":
            return self._backward(inp)

        if self.mode == "filter":
            return self._backward(self._forward(inp))

    def _forward(self, inp):
        pass

    def _backward(self, inp):
        pass
Example #2
0
class DelayFilter(task.SingleTask):
    """Remove delays less than a given threshold.

    Attributes
    ----------
    delay_cut : float
        Delay value to filter at in seconds.
    za_cut : float
        Sine of the maximum zenith angle included in
        baseline-dependent delay filtering. Default is 1
        which corresponds to the horizon (ie: filters
        out all zenith angles). Setting to zero turns off
        baseline dependent cut.
    update_weight : bool
        Not implemented.
    weight_tol : float
        Maximum weight kept in the masked data, as a fraction of
        the largest weight in the original dataset.
    telescope_orientation : one of ('NS', 'EW', 'none')
        Determines if the baseline-dependent delay cut is based on
        the north-south component, the east-west component or the full
        baseline length. For cylindrical telescopes oriented in the
        NS direction (like CHIME) use 'NS'. The default is 'NS'.
    """

    delay_cut = config.Property(proptype=float, default=0.1)
    za_cut = config.Property(proptype=float, default=1.0)
    update_weight = config.Property(proptype=bool, default=False)
    weight_tol = config.Property(proptype=float, default=1e-4)
    telescope_orientation = config.enum(["NS", "EW", "none"], default="NS")

    def setup(self, telescope):
        """Set the telescope needed to obtain baselines.

        Parameters
        ----------
        telescope : TransitTelescope
        """
        self.telescope = io.get_telescope(telescope)

    def process(self, ss):
        """Filter out delays from a SiderealStream or TimeStream.

        Parameters
        ----------
        ss : containers.SiderealStream
            Data to filter.

        Returns
        -------
        ss_filt : containers.SiderealStream
            Filtered dataset.
        """
        tel = self.telescope

        if self.update_weight:
            raise NotImplemented("Weight updating is not implemented.")

        ss.redistribute(["input", "prod", "stack"])

        freq = ss.freq[:]

        ssv = ss.vis[:].view(np.ndarray)
        ssw = ss.weight[:].view(np.ndarray)

        ubase, uinv = np.unique(
            tel.baselines[:, 0] + 1.0j * tel.baselines[:, 1], return_inverse=True
        )
        ubase = ubase.view(np.float64).reshape(-1, 2)

        for lbi, bi in ss.vis[:].enumerate(axis=1):

            # Select the baseline length to use
            baseline = ubase[uinv[bi]]
            if self.telescope_orientation == "NS":
                baseline = abs(baseline[1])  # Y baseline
            elif self.telescope_orientation == "EW":
                baseline = abs(baseline[0])  # X baseline
            else:
                baseline = np.linalg.norm(baseline)  # Norm
            # In micro seconds
            baseline_delay_cut = self.za_cut * baseline / units.c * 1e6
            delay_cut = np.amax([baseline_delay_cut, self.delay_cut])

            weight_mask = np.median(ssw[:, lbi], axis=1)
            weight_mask = (weight_mask > (self.weight_tol * weight_mask.max())).astype(
                np.float64
            )
            NF = null_delay_filter(freq, delay_cut, weight_mask)

            ssv[:, lbi] = np.dot(NF, ssv[:, lbi])

        return ss
Example #3
0
class MaskBaselines(task.SingleTask):
    """Mask out baselines from a dataset.

    This task may produce output with shared datasets. Be warned that
    this can produce unexpected outputs if not properly taken into
    account.

    Attributes
    ----------
    mask_long_ns : float
        Mask out baselines longer than a given distance in the N/S direction.
    mask_short : float
        Mask out baselines shorter than a given distance.
    mask_short_ew : float
        Mask out baselines shorter then a given distance in the East-West
        direction. Useful for masking out intra-cylinder baselines for
        North-South oriented cylindrical telescopes.
    zero_data : bool, optional
        Zero the data in addition to modifying the noise weights
        (default is False).
    share : {"all", "none", "vis"}
        Which datasets should we share with the input. If "none" we create a
        full copy of the data, if "vis" we create a copy only of the modified
        weight dataset and the unmodified vis dataset is shared, if "all" we
        modify in place and return the input container.
    """

    mask_long_ns = config.Property(proptype=float, default=None)
    mask_short = config.Property(proptype=float, default=None)
    mask_short_ew = config.Property(proptype=float, default=None)

    zero_data = config.Property(proptype=bool, default=False)

    share = config.enum(["none", "vis", "all"], default="all")

    def setup(self, telescope):
        """Set the telescope model.

        Parameters
        ----------
        telescope : TransitTelescope
        """

        self.telescope = io.get_telescope(telescope)

        if self.zero_data and self.share == "vis":
            self.log.warn(
                "Setting `zero_data = True` and `share = vis` doesn't make much sense."
            )

    def process(self, ss):
        """Apply the mask to data.

        Parameters
        ----------
        ss : SiderealStream or TimeStream
            Data to mask. Applied in place.
        """

        ss.redistribute("freq")

        baselines = self.telescope.baselines
        mask = np.ones_like(ss.weight[:], dtype=bool)

        if self.mask_long_ns is not None:
            long_ns_mask = np.abs(baselines[:, 1]) < self.mask_long_ns
            mask *= long_ns_mask[np.newaxis, :, np.newaxis]

        if self.mask_short is not None:
            short_mask = np.sum(baselines**2, axis=1)**0.5 > self.mask_short
            mask *= short_mask[np.newaxis, :, np.newaxis]

        if self.mask_short_ew is not None:
            short_ew_mask = baselines[:, 0] > self.mask_short_ew
            mask *= short_ew_mask[np.newaxis, :, np.newaxis]

        if self.share == "all":
            ssc = ss
        elif self.share == "vis":
            ssc = ss.copy(shared=("vis", ))
        else:  # self.share == "all"
            ssc = ss.copy()

        # Apply the mask to the weight
        ssc.weight[:] *= mask
        # Apply the mask to the data
        if self.zero_data:
            ssc.vis[:] *= mask

        return ssc
Example #4
0
class RFISensitivityMask(task.SingleTask):
    """Slightly less crappy RFI masking.

    Attributes
    ----------
    mask_type : string, optional
        One of 'mad', 'sumthreshold' or 'combine'.
        Default is combine, which uses the sumthreshold everywhere
        except around the transits of the Sun, CasA and CygA where it
        applies the MAD mask to avoid masking out the transits.
    include_pol : list of strings, optional
        The list of polarisations to include. Default is to use all
        polarisations.
    remove_median : bool, optional
        Remove median accross times for each frequency?
        Recomended. Default: True.
    sir : bool, optional
        Apply scale invariant rank (SIR) operator on top of final mask?
        We find that this is advisable while we still haven't flagged
        out all the static bands properly. Default: True.
    sigma : float, optional
        The false positive rate of the flagger given as sigma value assuming
        the non-RFI samples are Gaussian.
        Used for the MAD and TV station flaggers.
    max_m : int, optional
        Maximum size of the SumThreshold window to use.
        The default (8) seems to work well with sensitivity data.
    start_threshold_sigma : float, optional
        The desired threshold for the SumThreshold algorythm at the
        final window size (determined by max m) given as a
        number of standard deviations (to be estimated from the
        sensitivity map excluding weight and static masks).
        The default (8) seems to work well with sensitivity data
        using the default max_m.
    tv_fraction : float, optional
        Number of bad samples in a digital TV channel that cause the whole
        channel to be flagged.
    tv_base_size : [int, int]
        The size of the region used to estimate the baseline for the TV channel
        detection.
    tv_mad_size : [int, int]
        The size of the region used to estimate the MAD for the TV channel detection.
    """

    mask_type = config.enum(["mad", "sumthreshold", "combine"],
                            default="combine")
    include_pol = config.list_type(str, default=None)
    remove_median = config.Property(proptype=bool, default=True)
    sir = config.Property(proptype=bool, default=True)

    sigma = config.Property(proptype=float, default=5.0)
    max_m = config.Property(proptype=int, default=8)
    start_threshold_sigma = config.Property(proptype=float, default=8)

    tv_fraction = config.Property(proptype=float, default=0.5)
    tv_base_size = config.list_type(int, length=2, default=(11, 3))
    tv_mad_size = config.list_type(int, length=2, default=(201, 51))

    def process(self, sensitivity):
        """Derive an RFI mask from sensitivity data.

        Parameters
        ----------
        sensitivity : containers.SystemSensitivity
            Sensitivity data to derive the RFI mask from.

        Returns
        -------
        rfimask : containers.RFIMask
            RFI mask derived from sensitivity.
        """
        ## Constants
        # Convert MAD to RMS
        MAD_TO_RMS = 1.4826

        # The difference between the exponents in the usual
        # scaling of the RMS (n**0.5) and the scaling used
        # in the sumthreshold algorithm (n**log2(1.5))
        RMS_SCALING_DIFF = np.log2(1.5) - 0.5

        # Distribute over polarisation as we need all times and frequencies
        # available simultaneously
        sensitivity.redistribute("pol")

        # Divide sensitivity to get a radiometer test
        radiometer = sensitivity.measured[:] * tools.invert_no_zero(
            sensitivity.radiometer[:])
        radiometer = mpiarray.MPIArray.wrap(radiometer, axis=1)

        freq = sensitivity.freq
        npol = len(sensitivity.pol)
        nfreq = len(freq)

        static_flag = ~self._static_rfi_mask_hook(freq)

        madmask = mpiarray.MPIArray((npol, nfreq, len(sensitivity.time)),
                                    axis=0,
                                    dtype=np.bool)
        madmask[:] = False
        stmask = mpiarray.MPIArray((npol, nfreq, len(sensitivity.time)),
                                   axis=0,
                                   dtype=np.bool)
        stmask[:] = False

        for li, ii in madmask.enumerate(axis=0):

            # Only process this polarisation if we should be including it,
            # otherwise skip and let it be implicitly set to False (i.e. not
            # masked)
            if self.include_pol and sensitivity.pol[ii] not in self.include_pol:
                continue

            # Initial flag on weights equal to zero.
            origflag = sensitivity.weight[:, ii] == 0.0

            # Remove median at each frequency, if asked.
            if self.remove_median:
                for ff in range(nfreq):
                    radiometer[ff, li] -= np.median(
                        radiometer[ff, li][~origflag[ff]].view(np.ndarray))

            # Combine weights with static flag
            start_flag = origflag | static_flag[:, None]

            # Obtain MAD and TV masks
            this_madmask, tvmask = self._mad_tv_mask(radiometer[:, li],
                                                     start_flag, freq)

            # combine MAD and TV masks
            madmask[li] = this_madmask | tvmask

            # Add TV channels to ST start flag.
            start_flag = start_flag | tvmask

            # Determine initial threshold
            med = np.median(radiometer[:, li][~start_flag].view(np.ndarray))
            mad = np.median(
                abs(radiometer[:, li][~start_flag].view(np.ndarray) - med))
            threshold1 = (mad * MAD_TO_RMS * self.start_threshold_sigma *
                          self.max_m**RMS_SCALING_DIFF)

            # SumThreshold mask
            stmask[li] = rfi.sumthreshold(
                radiometer[:, li],
                self.max_m,
                start_flag=start_flag,
                threshold1=threshold1,
                correct_for_missing=True,
            )

        # Perform an OR (.any) along the pol axis and reform into an MPIArray
        # along the freq axis
        madmask = mpiarray.MPIArray.wrap(madmask.redistribute(1).any(0), 0)
        stmask = mpiarray.MPIArray.wrap(stmask.redistribute(1).any(0), 0)

        # Pick which of the MAD or SumThreshold mask to use (or blend them)
        if self.mask_type == "mad":
            finalmask = madmask

        elif self.mask_type == "sumthreshold":
            finalmask = stmask

        else:
            # Combine ST and MAD masks
            madtimes = self._combine_st_mad_hook(sensitivity.time)
            finalmask = stmask
            finalmask[:, madtimes] = madmask[:, madtimes]

        # Collect all parts of the mask onto rank 1 and then broadcast to all ranks
        finalmask = mpiarray.MPIArray.wrap(finalmask, 0).allgather()

        # Apply scale invariant rank (SIR) operator, if asked for.
        if self.sir:
            finalmask = self._apply_sir(finalmask, static_flag)

        # Create container to hold mask
        rfimask = containers.RFIMask(axes_from=sensitivity)
        rfimask.mask[:] = finalmask

        return rfimask

    def _combine_st_mad_hook(self, times):
        """Override this function to add a custom blending mask between the
        SumThreshold and MAD flagged data.

        This is useful to use the MAD algorithm around bright source
        transits, where the SumThreshold begins to remove real signal.

        Parameters
        ----------
        times : np.ndarray[ntime]
            Times of the data at floating point UNIX time.

        Returns
        -------
        combine : np.ndarray[ntime]
            Mixing array as a function of time. If `True` that sample will be
            filled from the MAD, if `False` use the SumThreshold algorithm.
        """
        return np.ones_like(times, dtype=np.bool)

    def _static_rfi_mask_hook(self, freq):
        """Override this function to apply a static RFI mask to the data.

        Parameters
        ----------
        freq : np.ndarray[nfreq]
            1D array of frequencies in the data (in MHz).

        Returns
        -------
        mask : np.ndarray[nfreq]
            Mask array. True will include a frequency channel, False masks it out.
        """
        return np.ones_like(freq, dtype=np.bool)

    def _apply_sir(self, mask, baseflag, eta=0.2):
        """Expand the mask with SIR."""

        # Remove baseflag from mask and run SIR
        nobaseflag = np.copy(mask)
        nobaseflag[baseflag] = False
        nobaseflagsir = rfi.sir(nobaseflag[:, np.newaxis, :], eta=eta)[:, 0, :]

        # Make sure the original mask (including baseflag) is still masked
        flagsir = nobaseflagsir | mask

        return flagsir

    def _mad_tv_mask(self, data, start_flag, freq):
        """Use the specific scattered TV channel flagging."""
        # Make copy of data
        data = np.copy(data)

        # Calculate the scaled deviations
        data[start_flag] = 0.0
        maddev = mad(data,
                     start_flag,
                     base_size=self.tv_base_size,
                     mad_size=self.tv_mad_size)

        # Replace any NaNs (where too much data is missing) with a
        # large enough value to always be flagged
        maddev = np.where(np.isnan(maddev), 2 * self.sigma, maddev)

        # Reflag for scattered TV emission
        tvmask = tv_channels_flag(maddev,
                                  freq,
                                  sigma=self.sigma,
                                  f=self.tv_fraction)

        # Create MAD mask
        madmask = maddev > self.sigma

        # Ensure start flag is masked
        madmask = madmask | start_flag

        return madmask, tvmask
Example #5
0
class TransitStacker(task.SingleTask):
    """Stack a number of transits together.

    The weights will be inverted and stacked as variances. The variance
    between transits is evaluated and recorded in the output container.

    All transits should be on a common grid in hour angle.

    Attributes
    ----------
    weight: str (default: uniform)
        The weighting to use in the stack. One of `uniform` or `inverse_variance`.
    """

    weight = config.enum(["uniform", "inverse_variance"], default="uniform")

    def setup(self):
        """Initialise internal variables."""

        self.stack = None
        self.variance = None
        self.pseudo_variance = None
        self.norm = None

    def process(self, transit):
        """Add a transit to the stack.

        Parameters
        ----------
        transit: draco.core.containers.TrackBeam
            A holography transit.
        """

        self.log.info("Weight is %s" % self.weight)

        if self.stack is None:
            self.log.info("Initializing transit stack.")
            self.stack = TrackBeam(
                axes_from=transit, distributed=transit.distributed, comm=transit.comm
            )

            self.stack.add_dataset("sample_variance")
            self.stack.add_dataset("nsample")
            self.stack.redistribute("freq")

            self.log.info("Adding %s to stack." % transit.attrs["tag"])

            # Copy over relevant attributes
            self.stack.attrs["filename"] = [transit.attrs["tag"]]
            self.stack.attrs["observation_id"] = [transit.attrs["observation_id"]]
            self.stack.attrs["transit_time"] = [transit.attrs["transit_time"]]
            self.stack.attrs["archivefiles"] = list(transit.attrs["archivefiles"])

            self.stack.attrs["dec"] = transit.attrs["dec"]
            self.stack.attrs["source_name"] = transit.attrs["source_name"]
            self.stack.attrs["icrs_ra"] = transit.attrs["icrs_ra"]
            self.stack.attrs["cirs_ra"] = transit.attrs["cirs_ra"]

            # Copy data for first transit
            flag = (transit.weight[:] > 0.0).astype(np.int)
            if self.weight == "inverse_variance":
                coeff = transit.weight[:]
            else:
                coeff = flag.astype(np.float32)

            self.stack.beam[:] = coeff * transit.beam[:]
            self.stack.weight[:] = (coeff**2) * invert_no_zero(transit.weight[:])
            self.stack.nsample[:] = flag.astype(np.int)

            self.variance = coeff * np.abs(transit.beam[:]) ** 2
            self.pseudo_variance = coeff * transit.beam[:] ** 2
            self.norm = coeff

        else:
            if list(transit.beam.shape) != list(self.stack.beam.shape):
                self.log.error(
                    "Transit has different shape than stack: {}, {}".format(
                        transit.beam.shape, self.stack.beam.shape
                    )
                    + " Skipping."
                )
                return None

            self.log.info("Adding %s to stack." % transit.attrs["tag"])

            self.stack.attrs["filename"].append(transit.attrs["tag"])
            self.stack.attrs["observation_id"].append(transit.attrs["observation_id"])
            self.stack.attrs["transit_time"].append(transit.attrs["transit_time"])
            self.stack.attrs["archivefiles"] += list(transit.attrs["archivefiles"])

            # Accumulate transit data
            flag = (transit.weight[:] > 0.0).astype(np.int)
            if self.weight == "inverse_variance":
                coeff = transit.weight[:]
            else:
                coeff = flag.astype(np.float32)

            self.stack.beam[:] += coeff * transit.beam[:]
            self.stack.weight[:] += (coeff**2) * invert_no_zero(transit.weight[:])
            self.stack.nsample[:] += flag

            self.variance += coeff * np.abs(transit.beam[:]) ** 2
            self.pseudo_variance += coeff * transit.beam[:] ** 2
            self.norm += coeff

        return None

    def process_finish(self):
        """Normalise the stack and return the result.

        Includes the sample variance over transits within the stack.

        Returns
        -------
        stack: draco.core.containers.TrackBeam
            Stacked transits.
        """
        # Divide by norm to get average transit
        inv_norm = invert_no_zero(self.norm)
        self.stack.beam[:] *= inv_norm
        self.stack.weight[:] = invert_no_zero(self.stack.weight[:]) * self.norm**2

        self.variance = self.variance * inv_norm - np.abs(self.stack.beam[:]) ** 2
        self.pseudo_variance = self.pseudo_variance * inv_norm - self.stack.beam[:] ** 2

        # Calculate the covariance between the real and imaginary component
        # from the accumulated variance and psuedo-variance
        self.stack.sample_variance[0] = 0.5 * (
            self.variance + self.pseudo_variance.real
        )
        self.stack.sample_variance[1] = 0.5 * self.pseudo_variance.imag
        self.stack.sample_variance[2] = 0.5 * (
            self.variance - self.pseudo_variance.real
        )

        # Create tag
        time_range = np.percentile(self.stack.attrs["transit_time"], [0, 100])
        self.stack.attrs["tag"] = "{}_{}_to_{}".format(
            self.stack.attrs["source_name"],
            ephem.unix_to_datetime(time_range[0]).strftime("%Y%m%dT%H%M%S"),
            ephem.unix_to_datetime(time_range[1]).strftime("%Y%m%dT%H%M%S"),
        )

        return self.stack
Example #6
0
class ConstructStackedBeam(task.SingleTask):
    """Construct the effective beam for stacked baselines.

    Parameters
    ----------
    weight : string ('uniform' or 'inverse_variance')
        How to weight the baselines when stacking:
            'uniform' - each baseline given equal weight
            'inverse_variance' - each baseline weighted by the weight attribute
    """

    weight = config.enum(["uniform", "inverse_variance"], default="uniform")

    def setup(self, tel):
        """Set the Telescope instance to use.

        Parameters
        ----------
        tel : TransitTelescope
        """
        self.telescope = io.get_telescope(tel)

    def process(self, beam, data):
        """Stack

        Parameters
        ----------
        beam : TrackBeam
            The beam that will be stacked.
        data : VisContainer
            Must contain `prod` index map and `stack` reverse map
            that will be used to stack the beam.

        Returns
        -------
        stacked_beam: VisContainer
            The input `beam` stacked in the same manner as

        """
        # Distribute over frequencies
        data.redistribute("freq")
        beam.redistribute("freq")

        # Grab the stack specifications from the input sidereal stream
        prod = data.index_map["prod"]
        reverse_stack = data.reverse_map["stack"][:]

        input_flags = data.input_flags[:]
        if not np.any(input_flags):
            input_flags = np.ones_like(input_flags)

        # Create output container
        if isinstance(data, SiderealStream):
            OutputContainer = SiderealStream
            output_kwargs = {"ra": data.ra[:]}
        else:
            OutputContainer = TimeStream
            output_kwargs = {"time": data.time[:]}

        stacked_beam = OutputContainer(
            axes_from=data,
            attrs_from=beam,
            distributed=True,
            comm=data.comm,
            **output_kwargs
        )

        stacked_beam.vis[:] = 0.0
        stacked_beam.weight[:] = 0.0

        stacked_beam.attrs["tag"] = "_".join([beam.attrs["tag"], data.attrs["tag"]])

        # Dereference datasets
        bv = beam.beam[:].view(np.ndarray)
        bw = beam.weight[:].view(np.ndarray)

        ov = stacked_beam.vis[:]
        ow = stacked_beam.weight[:]

        pol_filter = {
            "X": "X",
            "Y": "Y",
            "E": "X",
            "S": "Y",
            "co": "co",
            "cross": "cross",
        }
        pol = [pol_filter.get(pp, None) for pp in self.telescope.polarisation]
        beam_pol = [pol_filter.get(pp, None) for pp in beam.index_map["pol"][:]]

        # Compute the fractional variance of the beam measurement
        frac_var = invert_no_zero(bw * np.abs(bv) ** 2)

        # Create counter to increment during the stacking.
        # This will be used to normalize at the end.
        counter = np.zeros_like(ow)

        # Construct stack
        for pp, (ss, conj) in enumerate(reverse_stack):

            aa, bb = prod[pp]
            if conj:
                aa, bb = bb, aa

            try:
                aa_pol, bb_pol = self._resolve_pol(pol[aa], pol[bb], beam_pol)
            except ValueError:
                continue

            cross = bv[:, aa_pol, aa, :] * bv[:, bb_pol, bb, :].conj()

            weight = (
                input_flags[np.newaxis, aa, :]
                * input_flags[np.newaxis, bb, :]
                * invert_no_zero(
                    np.abs(cross) ** 2
                    * (frac_var[:, aa_pol, aa, :] + frac_var[:, bb_pol, bb, :])
                )
            )

            if self.weight == "inverse_variance":
                wss = weight
            else:
                wss = (weight > 0.0).astype(np.float32)

            # Accumulate variances in quadrature.  Save in the weight dataset.
            ov[:, ss, :] += wss * cross
            ow[:, ss, :] += wss**2 * invert_no_zero(weight)

            # Increment counter
            counter[:, ss, :] += wss

        # Divide through by counter to get properly weighted visibility average
        ov[:] *= invert_no_zero(counter)
        ow[:] = counter**2 * invert_no_zero(ow[:])

        return stacked_beam

    @staticmethod
    def _resolve_pol(pol1, pol2, pol_axis):

        if "co" in pol_axis:

            if pol1 == pol2:
                ipol = pol_axis.index("co")
            else:
                ipol = pol_axis.index("cross")

            return ipol, ipol

        else:

            if pol1 == pol2:
                ipol1 = pol_axis.index(pol1)
                ipol2 = pol_axis.index(pol2)
            else:
                ipol1 = pol_axis.index(pol2)
                ipol2 = pol_axis.index(pol1)

            return ipol1, ipol2
Example #7
0
class QuadraticPSEstimation(task.SingleTask):
    """Estimate a power spectrum from a set of KLModes.

    Attributes
    ----------
    psname : str
        Name of power spectrum to use. Must be precalculated in the driftscan
        products.
    pstype : str
        Type of power spectrum estimate to calculate. One of 'unwindowed',
        'minimum_variance' or 'uncorrelated'.
    """

    psname = config.Property(proptype=str)

    pstype = config.enum(["unwindowed", "minimum_variance", "uncorrelated"],
                         default="unwindowed")

    def setup(self, manager):
        """Set the ProductManager instance to use.

        Parameters
        ----------
        manager : ProductManager
        """
        self.manager = manager

    def process(self, klmodes):
        """Estimate the power spectrum from the given data.

        Parameters
        ----------
        klmodes : containers.KLModes

        Returns
        -------
        ps : containers.PowerSpectrum
        """

        import scipy.linalg as la

        if not isinstance(klmodes, containers.KLModes):
            raise ValueError("Input container must be instance of "
                             "KLModes (received %s)" % klmodes.__class__)

        klmodes.redistribute("m")

        pse = self.manager.psestimators[self.psname]
        pse.genbands()

        q_list = []

        for mi, m in klmodes.vis[:].enumerate(axis=0):
            ps_single = pse.q_estimator(m, klmodes.vis[m, :klmodes.nmode[m]])
            q_list.append(ps_single)

        q = klmodes.comm.allgather(np.array(q_list).sum(axis=0))
        q = np.array(q).sum(axis=0)

        # reading from directory
        fisher, bias = pse.fisher_bias()

        ps = containers.Powerspectrum2D(kperp_edges=pse.kperp_bands,
                                        kpar_edges=pse.kpar_bands)

        npar = len(ps.index_map["kpar"])
        nperp = len(ps.index_map["kperp"])

        # Calculate the right unmixing matrix for each ps type
        if self.pstype == "unwindowed":
            M = la.pinv(fisher, rcond=1e-8)
        elif self.pstype == "uncorrelated":
            Fh = la.cholesky(fisher)
            M = la.inv(Fh) / Fh.sum(axis=1)[:, np.newaxis]
        elif self.pstype == "minimum_variance":
            M = np.diag(fisher.sum(axis=1)**-1)

        ps.powerspectrum[:] = np.dot(M, q - bias).reshape(nperp, npar)
        ps.C_inv[:] = fisher.reshape(nperp, npar, nperp, npar)

        return ps
Example #8
0
class TransitTelescope(with_metaclass(abc.ABCMeta, config.Reader, ctime.Observer)):
    """Base class for simulating any transit interferometer.

    This is an abstract class, and several methods must be implemented before it
    is usable. These are:

    * `feedpositions` - a property which contains the positions of all the feeds
    * `_get_unique` -  calculates which baselines are identical
    * `_transfer_single` - calculate the beam transfer for a single baseline+freq
    * `_make_matrix_array` - makes an array of the right size to hold the
      transfer functions
    * `_copy_transfer_into_single` - copy a single transfer matrix into a
      collection.

    The last two are required for supporting polarised beam functions.

    Properties
    ----------
    zenith : [theta, phi]
        The position of the zenith spherical polars (in radians). Read only.
    freq_lower, freq_higher : scalar
        The center of the lowest and highest frequency bands. Deprecated, use
        `freq_start`, `freq_end` instead.
    freq_start, freq_end : scalar
        The start and end frequencies in MHz.
    num_freq : scalar
        The number of frequency bands (only use for setting up the frequency
        binning). Generally using `nfreq` is preferred.
    freq_mode : {"centre", "edge"}
        Choose if `freq_start` and `freq_end` are the edges of the band
        ("edge"), or whether they are the central frequencies of the first
        and last channel, in this case the last (Nyquist) frequency can
        either be skipped ("centre", default) or included ("centre_nyquist").
        The behaviour of the "centre" mode matches the output of the CASPER
        PFB-FIR block.
    channel_bin : int, optional
        Number of channels to bin together. This must exactly devide the total number.
        Binning is performed prior to selection of any subset. Default: 1.
    channel_list : list, optional
        List of channel indices to select. If set, this takes priority over
        `channel_range`. Currently this is not implemented.
    channel_range : list, optional
        Select subset of frequencies using a range of frequency channel indices,
        either [start, stop, step], [start, stop], or [stop] is acceptable.
        Default selects all channels.
    tsys_flat : scalar
        The system temperature (in K). Override `tsys` for anything more
        sophisticated.
    positive_m_only: boolean
        Whether to only deal with half the `m` range. In many cases we are
        much less sensitive to negative-m (depending on the hemisphere, and
        baseline alignment). This does not affect the beams calculated, only
        how they're used in further calculation. Default: False
    minlength, maxlength : scalar
        Minimum and maximum baseline lengths to include (in metres).
    local_origin : bool
        If set the observers location is the terrestrial origin, and so the
        rotation angle corresponds to the right ascension that is overhead
        (Local Stellar Angle in `caput.time`). If not the origin is Greenwich,
        so the rotation angle is what is overhead at Greenwich (Earth Rotation
        Angle).
    """

    freq_lower = config.Property(proptype=float, default=None)
    freq_upper = config.Property(proptype=float, default=None)

    freq_start = config.Property(proptype=float, default=800.0)
    freq_end = config.Property(proptype=float, default=400.0)
    num_freq = config.Property(proptype=int, default=1024)

    freq_mode = config.enum(["centre", "centre_nyquist", "edge"], default="centre")

    channel_bin = config.Property(proptype=int, default=1)
    channel_range = config.Property(proptype=list)
    channel_list = config.Property(proptype=list)

    tsys_flat = config.Property(proptype=float, default=50.0, key="tsys")
    ndays = config.Property(proptype=int, default=733)

    accuracy_boost = config.Property(proptype=float, default=1.0)
    l_boost = config.Property(proptype=float, default=1.0)

    minlength = config.Property(proptype=float, default=0.0)
    maxlength = config.Property(proptype=float, default=1.0e7)

    auto_correlations = config.Property(proptype=bool, default=False)

    local_origin = config.Property(proptype=bool, default=True)

    def __init__(self, latitude=45, longitude=0, **kwargs):
        """Initialise a telescope object.

        Parameters
        ----------
        latitude, longitude : scalar
            Position on the Earths surface of the telescope (in degrees).
        """

        # Set the observers position on the Earth
        ctime.Observer.__init__(self, longitude, latitude, **kwargs)

    _pickle_keys = []

    def __getstate__(self):

        state = self.__dict__.copy()

        for key in self.__dict__:
            if (key not in self._pickle_keys) and (key[0] == "_"):
                del state[key]

        return state

    @property
    def zenith(self):
        """The zenith vector in spherical polars."""

        # Set polar angle
        theta = np.pi / 2.0 - np.radians(self.latitude)

        # Set azimuthal angle
        phi = np.remainder(np.radians(self.longitude), 2 * np.pi)

        # If we want a local origin, the observers location is the terrestrial
        # origin, so the zenith should be at phi=0. Otherwise the origin is
        # Greenwich, so we need the longitude.
        phi = 0.0 if self.local_origin else phi

        return np.array([theta, phi])

    # ========= Properties related to baselines =========

    _baselines = None

    @property
    def baselines(self):
        """The unique baselines in the telescope."""
        if self._baselines is None:
            self.calculate_feedpairs()

        return self._baselines

    _redundancy = None

    @property
    def redundancy(self):
        """The redundancy of each baseline (corresponds to entries in
        cyl.baselines)."""
        if self._redundancy is None:
            self.calculate_feedpairs()

        return self._redundancy

    @property
    def nbase(self):
        """The number of unique baselines."""
        return self.npairs

    @property
    def npairs(self):
        """The number of unique feed pairs."""
        return self.uniquepairs.shape[0]

    _uniquepairs = None

    @property
    def uniquepairs(self):
        """An (npairs, 2) array of the feed pairs corresponding to each baseline."""
        if self._uniquepairs is None:
            self.calculate_feedpairs()
        return self._uniquepairs

    _feedmap = None

    @property
    def feedmap(self):
        """An (nfeed, nfeed) array giving the mapping between feedpairs and
        the calculated baselines. Each entry is an index into the arrays of unique pairs."""

        if self._feedmap is None:
            self.calculate_feedpairs()

        return self._feedmap

    _feedmask = None

    @property
    def feedmask(self):
        """An (nfeed, nfeed) array giving the entries that have been
        calculated. This allows to mask out pairs we want to ignore."""

        if self._feedmask is None:
            self.calculate_feedpairs()

        return self._feedmask

    _feedconj = None

    @property
    def feedconj(self):
        """An (nfeed, nfeed) array giving the feed pairs which must be complex
        conjugated."""

        if self._feedconj is None:
            self.calculate_feedpairs()

        return self._feedconj

    # ===================================================

    # ======== Properties related to frequencies ========

    _frequencies = None

    @property
    def frequencies(self):
        """The centre of each frequency band (in MHz)."""
        if self._frequencies is None:
            self.calculate_frequencies()

        return self._frequencies

    def calculate_frequencies(self):

        if self.freq_lower or self.freq_upper:
            import warnings

            warnings.warn(
                "`freq_lower` and `freq_upper` parameters are deprecated",
                DeprecationWarning,
            )
            self.freq_start = self.freq_lower
            self.freq_end = self.freq_upper

        if self.freq_mode == "centre":
            df = abs(self.freq_end - self.freq_start) / self.num_freq
            frequencies = np.linspace(
                self.freq_start, self.freq_end, self.num_freq, endpoint=False
            )
        elif self.freq_mode == "centre_nyquist":
            df = abs(self.freq_end - self.freq_start) / (self.num_freq - 1)
            frequencies = np.linspace(
                self.freq_start, self.freq_end, self.num_freq, endpoint=True
            )
        else:
            df = abs(self.freq_end - self.freq_start) / self.num_freq
            frequencies = self.freq_start + df * (np.arange(self.num_freq) + 0.5)

        # Rebin frequencies if needed
        if self.channel_bin > 1:

            if self.num_freq % self.channel_bin != 0:
                raise ValueError(
                    "Channel binning must exactly divide the total number of channels"
                )

            frequencies = frequencies.reshape(-1, self.channel_bin).mean(axis=1)
            df = df * self.channel_bin

        # Select a subset of channels if required
        if self.channel_list is not None:
            raise NotImplementedError(
                "`channel_list` is not yet supported, as sparse channel selections "
                "may break things downstream."
            )
        if self.channel_range is not None:
            frequencies = frequencies[self.channel_range[0] : self.channel_range[1]]

        # TODO: do something with the channel width `df` as well
        self._frequencies = frequencies

    @property
    def wavelengths(self):
        """The central wavelength of each frequency band (in metres)."""
        return units.c / (1e6 * self.frequencies)

    @property
    def nfreq(self):
        """The number of frequency bins."""
        return self.frequencies.shape[0]

    # ===================================================

    # ======== Properties related to the feeds ==========

    @property
    def input_index(self):
        """Override to add custom labelling of the inputs, e.g. serial numbers.

        This should give an identifier that uniquely labels a correlator input and so
        can be used to match inputs through subsetting and reordering.

        There are two conventional fields used in the output, either a `chan_id`
        field for an integer label, or a `correlator_input` for a string labelling
        (useful for serial number strings). If both are present, `correlator_input`
        is used.
        """
        return np.array(np.arange(self.nfeed), dtype=[("chan_id", "u2")])

    @property
    def nfeed(self):
        """The number of feeds."""
        return self.feedpositions.shape[0]

    # ===================================================

    # ======= Properties related to polarisation ========

    @property
    def num_pol_sky(self):
        """The number of polarisation combinations on the sky that we are
        considering. Should be either 1 (T=I only), 3 (T, Q, U) or 4 (T, Q, U and V).
        """
        return self._npol_sky_

    # ===================================================

    # ===== Properties related to harmonic spread =======

    @property
    def lmax(self):
        """The maximum l the telescope is sensitive to."""
        lmax, mmax = max_lm(
            self.baselines, self.wavelengths.min(), self.u_width, self.v_width
        )
        return int(np.ceil(lmax.max() * self.l_boost))

    @property
    def mmax(self):
        """The maximum m the telescope is sensitive to."""
        lmax, mmax = max_lm(
            self.baselines, self.wavelengths.min(), self.u_width, self.v_width
        )
        return int(np.ceil(mmax.max() * self.l_boost))

    # ===================================================

    # == Methods for calculating the unique baselines ===

    def calculate_feedpairs(self):
        """Calculate all the unique feedpairs and their redundancies, and set
        the internal state of the object.
        """

        # Get unique pairs, and create mapping arrays
        self._feedmap, self._feedmask, self._feedconj = self._get_unique()

        # Reorder and conjugate baselines such that the default feedpair
        # points W->E (to ensure we use positive-m)
        self._make_ew()

        # Sort baselines into order
        self._sort_pairs()

        # Create mask of included pairs, that are not conjugated
        tmask = np.logical_and(self._feedmask, np.logical_not(self._feedconj))

        self._uniquepairs = _get_indices(self._feedmap, mask=tmask)
        self._redundancy = np.bincount(
            self._feedmap[np.where(tmask)]
        )  # Triangle mask to avoid double counting
        self._baselines = (
            self.feedpositions[self._uniquepairs[:, 0]]
            - self.feedpositions[self._uniquepairs[:, 1]]
        )

    def _make_ew(self):
        # Reorder baselines pairs, such that the baseline vector always points E (or pure N)

        tmask = np.logical_and(self._feedmask, np.logical_not(self._feedconj))
        uniq = _get_indices(self._feedmap, mask=tmask)

        conj_map = np.zeros(uniq.shape[0] + 1, dtype=np.bool)

        for i in range(uniq.shape[0]):
            sep = self.feedpositions[uniq[i, 0]] - self.feedpositions[uniq[i, 1]]

            if sep[0] < 0.0 or (sep[0] == 0.0 and sep[1] < 0.0):
                # Note down that we need to flip feedconj
                conj_map[i] = True

        # Flip the feedpairs
        self._feedconj = np.logical_xor(self._feedconj, conj_map[self._feedmap])

    # Tolerance used when comparing baselines. See np.around documentation for details.
    _bl_tol = 6

    def _unique_baselines(self):
        """Map of equivalent baseline lengths, and mask of ones to exclude."""
        # Construct array of indices
        fshape = [self.nfeed, self.nfeed]
        f_ind = np.indices(fshape)

        # Construct array of baseline separations in complex representation
        bl1 = self.feedpositions[f_ind[0]] - self.feedpositions[f_ind[1]]
        bl2 = np.around(bl1[..., 0] + 1.0j * bl1[..., 1], self._bl_tol)

        # Flip sign if required to get common direction to correctly find redundant baselines
        # flip_sign = np.logical_or(bl2.real < 0.0, np.logical_and(bl2.real == 0, bl2.imag < 0))
        # bl2 = np.where(flip_sign, -bl2, bl2)

        # Construct array of baseline lengths
        blen = np.sum(bl1 ** 2, axis=-1) ** 0.5

        # Create mask of included baselines
        mask = np.logical_and(blen >= self.minlength, blen <= self.maxlength)

        # Remove the auto correlated baselines between all polarisations
        if not self.auto_correlations:
            mask = np.logical_and(blen > 0.0, mask)

        return _remap_keyarray(bl2, mask), mask

    def _unique_beams(self):
        """Map of unique beam pairs, and mask of ones to exclude."""
        # Construct array of indices
        fshape = [self.nfeed, self.nfeed]

        bci, bcj = np.broadcast_arrays(
            self.beamclass[:, np.newaxis], self.beamclass[np.newaxis, :]
        )

        beam_map = _merge_keyarray(bci, bcj)

        if self.auto_correlations:
            beam_mask = np.ones(fshape, dtype=np.bool)
        else:
            beam_mask = np.logical_not(np.identity(self.nfeed, dtype=np.bool))

        return beam_map, beam_mask

    def _get_unique(self):
        """Calculate the unique baseline pairs.

        All feeds are assumed to be identical. Baselines are identified if
        they have the same length, and are selected such that they point East
        (to ensure that the sensitivity ends up in positive-m modes).

        It is also possible to select only baselines within a particular
        length range by setting the `minlength` and `maxlength` properties.

        Parameters
        ----------
        fpairs : np.ndarray
            An array of all the feed pairs, packed as [[i1, i2, ...], [j1, j2, ...] ].

        Returns
        -------
        baselines : np.ndarray
            An array of all the unique pairs. Packed as [ [i1, i2, ...], [j1, j2, ...]].
        redundancy : np.ndarray
            For each unique pair, give the number of equivalent pairs.
        """

        # Fetch and merge map of unique feed pairs
        base_map, base_mask = self._unique_baselines()
        beam_map, beam_mask = self._unique_beams()
        comb_map, comb_mask = _merge_keyarray(
            base_map, beam_map, mask1=base_mask, mask2=beam_mask
        )

        # Take into account conjugation by identifying the indices of conjugate pairs
        conj_map = comb_map > comb_map.T
        comb_map = np.dstack((comb_map, comb_map.T)).min(axis=-1)
        comb_map = _remap_keyarray(comb_map, comb_mask)

        return comb_map, comb_mask, conj_map

    def _sort_pairs(self):
        """Re-order keys into a desired sort order.

        By default the order is lexicographic in (baseline u, baselines v,
        beamclass i, beamclass j).
        """

        # Create mask of included pairs, that are not conjugated
        tmask = np.logical_and(self._feedmask, np.logical_not(self._feedconj))
        uniq = _get_indices(self._feedmap, mask=tmask)

        fi, fj = uniq[:, 0], uniq[:, 1]

        # Fetch keys by which to sort (lexicographically)
        bx = self.feedpositions[fi, 0] - self.feedpositions[fj, 0]
        by = self.feedpositions[fi, 1] - self.feedpositions[fj, 1]
        ci = self.beamclass[fi]
        cj = self.beamclass[fj]

        ## Sort by constructing a numpy array with the keys as fields, and use
        ## np.argsort to get the indices

        # Create array of keys to sort
        dt = np.dtype("f8,f8,i4,i4")
        sort_arr = np.zeros(fi.size, dtype=dt)
        sort_arr["f0"] = bx
        sort_arr["f1"] = by
        sort_arr["f2"] = cj
        sort_arr["f3"] = ci

        # Get map which sorts
        sort_ind = np.argsort(sort_arr)

        # Invert mapping
        tmp_sort_ind = sort_ind.copy()
        sort_ind[tmp_sort_ind] = np.arange(sort_ind.size)

        # Remap feedmap entries
        fm_copy = self._feedmap.copy()
        wmask = np.where(self._feedmask)
        fm_copy[wmask] = sort_ind[self._feedmap[wmask]]

        self._feedmap = fm_copy

    # ===================================================

    # ==== Methods for calculating Transfer matrices ====

    def transfer_matrices(self, bl_indices, f_indices, global_lmax=True):
        """Calculate the spherical harmonic transfer matrices for baseline and
        frequency combinations.

        Parameters
        ----------
        bl_indices : array_like
            Indices of baselines to calculate.
        f_indices : array_like
            Indices of frequencies to calculate. Must be broadcastable against
            `bl_indices`.
        global_lmax : boolean, optional
            If set (default), the output size `lside` in (l,m) is big enough to
            hold the maximum for the entire telescope. If not set it is only big
            enough for the requested set.

        Returns
        -------
        transfer : np.ndarray, dtype=np.complex128
            An array containing the transfer functions. The shape is somewhat
            complicated, the first indices correspond to the broadcast size of
            `bl_indices` and `f_indices`, then there may be some polarisation
            indices, then finally the (l,m) indices, range (lside, 2*lside-1).
        """

        # Broadcast arrays against each other
        bl_indices, f_indices = np.broadcast_arrays(bl_indices, f_indices)

        ## Check indices are all in range
        if out_of_range(bl_indices, 0, self.npairs):
            raise Exception("Baseline indices aren't valid")

        if out_of_range(f_indices, 0, self.nfreq):
            raise Exception("Frequency indices aren't valid")

        # Fetch the set of lmax's for the baselines (in order to reduce time
        # regenerating Healpix maps)
        lmax, mmax = np.ceil(
            self.l_boost
            * np.array(
                max_lm(
                    self.baselines[bl_indices],
                    self.wavelengths[f_indices],
                    self.u_width,
                    self.v_width,
                )
            )
        ).astype(np.int64)
        # lmax, mmax = lmax * self.l_boost, mmax * self.l_boost
        # Set the size of the (l,m) array to write into
        lside = self.lmax if global_lmax else lmax.max()

        # Generate the array for the Transfer functions

        tshape = bl_indices.shape + (self.num_pol_sky, lside + 1, 2 * lside + 1)
        print(
            "Size: %i elements. Memory %f GB."
            % (np.prod(tshape), 2 * np.prod(tshape) * 8.0 / 2 ** 30)
        )
        tarray = np.zeros(tshape, dtype=np.complex128)

        # Sort the baselines by ascending lmax and iterate through in that
        # order, calculating the transfer matrices
        i_arr = np.argsort(lmax.flat)

        for iflat in np.argsort(lmax.flat):
            ind = np.unravel_index(iflat, lmax.shape)
            trans = self._transfer_single(
                bl_indices[ind], f_indices[ind], lmax[ind], lside
            )

            ## Iterate over pol combinations and copy into transfer array
            for pi in range(self.num_pol_sky):
                islice = ind + (pi,) + (slice(None), slice(None))
                tarray[islice] = trans[pi]

        return tarray

    def transfer_for_frequency(self, freq):
        """Fetch all transfer matrices for a given frequency.

        Parameters
        ----------
        freq : integer
            The frequency index.

        Returns
        -------
        transfer : np.ndarray
            The transfer matrices. Packed as in `TransitTelescope.transfer_matrices`.
        """
        bi = np.arange(self.npairs)
        fi = freq * np.ones_like(bi)

        return self.transfer_matrices(bi, fi)

    def transfer_for_baseline(self, baseline):
        """Fetch all transfer matrices for a given baseline.

        Parameters
        ----------
        baseline : integer
            The baseline index.

        Returns
        -------
        transfer : np.ndarray
            The transfer matrices. Packed as in `TransitTelescope.transfer_matrices`.
        """
        fi = np.arange(self.nfreq)
        bi = baseline * np.ones_like(fi)

        return self.transfer_matrices(bi, fi)

    # ===================================================

    # ======== Noise properties of the telescope ========

    def tsys(self, f_indices=None):
        """The system temperature.

        Currenty has a flat T_sys across the whole bandwidth. Override for
        anything more complicated.

        Parameters
        ----------
        f_indices : array_like
            Indices of frequencies to get T_sys at.

        Returns
        -------
        tsys : array_like
            System temperature at requested frequencies.
        """
        if f_indices is None:
            freq = self.frequencies
        else:
            freq = self.frequencies[f_indices]
        return np.ones_like(freq) * self.tsys_flat

    def noisepower(self, bl_indices, f_indices, ndays=None):
        """Calculate the instrumental noise power spectrum.

        Assume we are still within the regime where the power spectrum is white
        in `m` modes.

        Parameters
        ----------
        bl_indices : array_like
            Indices of baselines to calculate.
        f_indices : array_like
            Indices of frequencies to calculate. Must be broadcastable against
            `bl_indices`.
        ndays : integer
            The number of sidereal days observed.

        Returns
        -------
        noise_ps : np.ndarray
            The noise power spectrum.
        """

        ndays = self.ndays if not ndays else ndays  # Set to value if not set.

        # Broadcast arrays against each other
        bl_indices, f_indices = np.broadcast_arrays(bl_indices, f_indices)

        bw = np.abs(self.frequencies[1] - self.frequencies[0]) * 1e6
        delnu = units.t_sidereal * bw / (2 * np.pi)
        noisepower = self.tsys(f_indices) ** 2 / (2 * np.pi * delnu * ndays)
        noisebase = noisepower / self.redundancy[bl_indices]

        return noisebase

    def noisepower_feedpairs(self, fi, fj, f_indices, m, ndays=None):
        ndays = self.ndays if not ndays else ndays

        bw = np.abs(self.frequencies[1] - self.frequencies[0]) * 1e6
        delnu = units.t_sidereal * bw / (2 * np.pi)
        noisepower = self.tsys(f_indices) ** 2 / (2 * np.pi * delnu * ndays)

        return (
            np.ones_like(fi) * np.ones_like(fj) * np.ones_like(m) * noisepower / 2.0
        )  # For unpolarised only at the moment.

    # ===================================================

    _nside = None

    def _init_trans(self, nside):
        ## Internal function for generating some common Healpix maps (position,
        ## horizon). These should need to be generated only when nside changes.

        # Angular positions in healpix map of nside
        self._nside = nside
        self._angpos = hputil.ang_positions(nside)

        # The horizon function
        self._horizon = visibility.horizon(self._angpos, self.zenith)

    # ===================================================
    # ================ ABSTRACT METHODS =================
    # ===================================================

    # Implement to specify feed positions in the telescope.
    @abc.abstractproperty
    def feedpositions(self):
        """An (nfeed,2) array of the feed positions relative to an arbitary point (in m)"""
        return

    # Implement to specify the beams of the telescope
    @abc.abstractproperty
    def beamclass(self):
        """An nfeed array of the class of each beam (identical labels are
        considered to have identical beams)."""
        return

    # Implement to specify feed positions in the telescope.
    @abc.abstractproperty
    def u_width(self):
        """The approximate physical width (in the u-direction) of the dish/telescope etc, for
        calculating the maximum (l,m)."""
        return

    # Implement to specify feed positions in the telescope.
    @abc.abstractproperty
    def v_width(self):
        """The approximate physical length (in the v-direction) of the dish/telescope etc, for
        calculating the maximum (l,m)."""
        return

    # The work method which does the bulk of calculating all the transfer matrices.
    @abc.abstractmethod
    def _transfer_single(self, bl_index, f_index, lmax, lside):
        """Calculate transfer matrix for a single baseline+frequency.

        **Abstract method** must be implemented.

        Parameters
        ----------
        bl_index : integer
            The index of the baseline to calculate.
        f_index : integer
            The index of the frequency to calculate.
        lmax : integer
            The maximum *l* we are interested in. Determines accuracy of
            spherical harmonic transforms.
        lside : integer
            The size of array to embed the transfer matrix within.

        Returns
        -------
        transfer : np.ndarray
            The transfer matrix, an array of shape (pol_indices, lside,
            2*lside-1). Where the `pol_indices` are usually only present if
            considering the polarised case.
        """
        return
class TimingErrors(gain.BaseGains):
    r"""Simulate timing distribution errors and propagate them to gain errors.

    This task simulates errors in the delay and calculates the errors in the
    phase of the gain according to

    .. math::
        \delta \phi_i = 2 \pi \nu \delta \tau_i

    The phase errors are generated for times matching the input timestream file.

    There are severeal options to run the timing distribution simulations.
    Choosing a `sim_type` defines if the delay errors are non common-mode or
    common-mode within a cylinder or iceboard.

    Furthermore there are two different options to simulate delay errors that
    are common-mode within a cylinder (`common_mode_cyl`). `Random` generates
    random fluctuations in phase whereas 'sinusoidal' as the name suggests
    simulates sinusoids for each cylinder with frequencies specified in
    the attribute `sinusoidal_period`.

    Attributes
    ----------
    corr_length_delay : float
        Correlation length for delay fluctuations in seconds.
    sigma_delay : float
        Size of fluctuations for delay fluctuations (s).
    sim_type: string, optional
        Timing error simulation type. List of allowed options are `relative`
        (non common-mode delay errors), `common_mode_cyl` (common-mode within
        a cylinder) and `common_mode_iceboard` (common-mode within an iceboard)
    common_mode_type : string, optional
        Options are 'random' and 'sinusoidal' if sim_type `common_mode_cyl` is chosen.
    sinusoidal_period : list, optional
        Specify the periods of the sinusoids for each cylinder. Needs to be specified
        when simulating `sinusoidal` `common_mode_cyl` timing errors.
    """
    ndays = config.Property(proptype=float, default=733)
    corr_length_delay = config.Property(proptype=float, default=3600)
    sigma_delay = config.Property(proptype=float, default=1e-12)

    sim_type = config.enum(
        ["relative", "common_mode_cyl", "common_mode_iceboard"],
        default="relative")

    common_mode_type = config.enum(["random", "sinusoidal"], default="random")

    # Default periods of chime specific timing jitter with the clock
    # distribution system informed by data see doclib 704
    sinusoidal_period = config.Property(proptype=list, default=[333, 500])

    _prev_delay = None
    _prev_time = None

    amp = False

    nchannel = 16
    ncyl = 2

    def _generate_phase(self, time):
        ntime = len(time)
        freq = self.freq
        nfreq = len(freq)

        # Generate the correlation function
        cf_delay = self._corr_func(self.corr_length_delay, self.sigma_delay)

        # Check if we are simulating relative delays or common mode delays
        if self.sim_type == "relative":
            n_realisations = self.ninput_local

            # Generate delay fluctuations
            self.delay_error = gain.generate_fluctuations(
                time, cf_delay, n_realisations, self._prev_time,
                self._prev_delay)

            gain_phase = (2.0 * np.pi * freq[:, np.newaxis, np.newaxis] * 1e6 *
                          self.delay_error[np.newaxis, :, :] /
                          np.sqrt(self.ndays))

        if self.sim_type == "common_mode_cyl":
            n_realisations = 1
            ninput = self.ninput_global

            # Generates as many random delay errors as there are cylinders
            if self.comm.rank == 0:
                if self.common_mode_type == "sinusoidal":
                    P1 = self.sinusoidal_period[0]
                    P2 = self.sinusoidal_period[1]
                    omega1 = 2 * np.pi / P1
                    omega2 = 2 * np.pi / P2

                    delay_error = (self.sigma_delay *
                                   (np.sin(omega1 * time) -
                                    np.sin(omega2 * time))[np.newaxis, :])

                if self.common_mode_type == "random":
                    delay_error = gain.generate_fluctuations(
                        time,
                        cf_delay,
                        n_realisations,
                        self._prev_time,
                        self._prev_delay,
                    )
            else:
                delay_error = None

            # Broadcast to other ranks
            self.delay_error = self.comm.bcast(delay_error, root=0)

            # Split frequencies to processes.
            lfreq, sfreq, efreq = mpiutil.split_local(nfreq)

            # Create an array to hold all inputs, which are common-mode within
            # a cylinder
            gain_phase = np.zeros((lfreq, ninput, ntime), dtype=complex)
            # Since we have 2 cylinders populate half of them with a delay)
            # TODO: generalize this for 3 or even 4 cylinders in the future.
            gain_phase[:, ninput // self.ncyl:, :] = (
                2.0 * np.pi * freq[sfreq:efreq, np.newaxis, np.newaxis] * 1e6 *
                self.delay_error[np.newaxis, :, :] / np.sqrt(self.ndays))

            gain_phase = mpiarray.MPIArray.wrap(gain_phase,
                                                axis=0,
                                                comm=self.comm)
            # Redistribute over input to match rest of the code
            gain_phase = gain_phase.redistribute(axis=1)
            gain_phase = gain_phase.view(np.ndarray)

        if self.sim_type == "common_mode_iceboard":
            nchannel = self.nchannel
            ninput = self.ninput_global
            # Number of channels on a board
            nboards = ninput // nchannel

            # Generates as many random delay errors as there are iceboards
            if self.comm.rank == 0:
                delay_error = gain.generate_fluctuations(
                    time, cf_delay, nboards, self._prev_time, self._prev_delay)
            else:
                delay_error = None

            # Broadcast to other ranks
            self.delay_error = self.comm.bcast(delay_error, root=0)

            # Calculate the corresponding phase by multiplying with frequencies
            phase = (2.0 * np.pi * freq[:, np.newaxis, np.newaxis] * 1e6 *
                     self.delay_error[np.newaxis, :] / np.sqrt(self.ndays))

            # Create an array to hold all inputs, which are common-mode within
            # one iceboard
            gain_phase = mpiarray.MPIArray((nfreq, ninput, ntime),
                                           axis=1,
                                           dtype=np.complex128,
                                           comm=self.comm)
            gain_phase[:] = 0.0

            # Loop over inputs and and group common-mode phases on every board
            for il, ig in gain_phase.enumerate(axis=1):
                # Get the board number bi
                bi = int(ig / nchannel)
                gain_phase[:, il] = phase[:, bi]

            gain_phase = gain_phase.view(np.ndarray)

        self._prev_delay = self.delay_error
        self._prev_time = time

        return gain_phase
Example #10
0
class BeamFormBase(task.SingleTask):
    """Base class for beam forming tasks.

    Defines a few useful methods. Not to be used directly
    but as parent class for BeamForm and BeamFormCat.

    Attributes
    ----------
    collapse_ha : bool
        Wether or not to sum over hour-angle/time to complete
        the beamforming. Default is True, which sums over.
    polarization : string
        One of:
        'I' : Stokes I only.
        'full' : 'XX', 'XY', 'YX' and 'YY' in this order.
        'copol' : 'XX' and 'YY' only.
        'stokes' : 'I', 'Q', 'U' and 'V' in this order. Not implemented.
    weight : string ('natural', 'uniform', or 'inverse_variance')
        How to weight the redundant baselines when adding:
            'natural' - each baseline weighted by its redundancy (default)
            'uniform' - each baseline given equal weight
            'inverse_variance' - each baseline weighted by the weight attribute
    timetrack : float
        How long (in seconds) to track sources at each side of transit.
        Total transit time will be ~ 2 * timetrack.
    freqside : int
        Number of frequencies to process at each side of the source.
        Default (None) processes all frequencies.
    """

    collapse_ha = config.Property(proptype=bool, default=True)
    polarization = config.enum(["I", "full", "copol", "stokes"],
                               default="full")
    weight = config.enum(["natural", "uniform", "inverse_variance"],
                         default="natural")
    timetrack = config.Property(proptype=float, default=900.0)
    freqside = config.Property(proptype=int, default=None)

    def setup(self, manager):
        """Generic setup method.

        To be complemented by specific
        setup methods in daughter tasks.

        Parameters
        ----------
        manager : either `ProductManager`, `BeamTransfer` or `TransitTelescope`
            Contains a TransitTelescope object describing the telescope.
        """
        # Get the TransitTelescope object
        self.telescope = io.get_telescope(manager)
        # Polarizations.
        if self.polarization == "I":
            self.process_pol = ["XX", "YY"]
            self.return_pol = ["I"]
        elif self.polarization == "full":
            self.process_pol = ["XX", "XY", "YX", "YY"]
            self.return_pol = self.process_pol
        elif self.polarization == "copol":
            self.process_pol = ["XX", "YY"]
            self.return_pol = self.process_pol
        elif self.polarization == "stokes":
            self.process_pol = ["XX", "XY", "YX", "YY"]
            self.return_pol = ["I", "Q", "U", "V"]
            msg = "Stokes parameters are not implemented"
            raise RuntimeError(msg)
        else:
            # This should never happen. config.enum should bark first.
            msg = "Invalid polarization parameter: {0}"
            msg = msg.format(self.polarization)
            raise ValueError(msg)
        # Number of polarizations to process
        self.npol = len(self.process_pol)
        self.latitude = np.deg2rad(self.telescope.latitude)

    def process(self):
        """Generic process method.

        Performs all the beamforming,
        but not the data parsing. To be complemented by specific
        process methods in daughter tasks.

        Returns
        -------
        formed_beam : `containers.FormedBeam` or `containers.FormedBeamHA`
            Formed beams at each source. Shape depends on parameter
            `collapse_ha`.
        """
        # Contruct containers for formed beams
        if self.collapse_ha:
            # Container to hold the formed beams
            formed_beam = containers.FormedBeam(
                freq=self.freq,
                object_id=self.source_cat.index_map["object_id"],
                pol=np.array(self.return_pol),
                distributed=True,
            )
        else:
            # Container to hold the formed beams
            formed_beam = containers.FormedBeamHA(
                freq=self.freq,
                ha=np.arange(self.nha, dtype=np.int),
                object_id=self.source_cat.index_map["object_id"],
                pol=np.array(self.return_pol),
                distributed=True,
            )
            # Initialize container to zeros.
            formed_beam.ha[:] = 0.0

        # Initialize container to zeros.
        formed_beam.beam[:] = 0.0
        formed_beam.weight[:] = 0.0
        # Copy catalog information
        formed_beam["position"][:] = self.source_cat["position"][:]
        if "redshift" in self.source_cat:
            formed_beam["redshift"][:] = self.source_cat["redshift"][:]
        else:
            # TODO: If there is not redshift information,
            # should I have a different formed_beam container?
            formed_beam["redshift"]["z"][:] = 0.0
            formed_beam["redshift"]["z_error"][:] = 0.0
        # Ensure container is distributed in frequency
        formed_beam.redistribute("freq")

        if self.freqside is None:
            # Indices of local frequency axis. Full axis if freqside is None.
            f_local_indices = np.arange(self.ls, dtype=np.int32)
            f_mask = np.zeros(self.ls, dtype=bool)

        fbb = formed_beam.beam[:]
        fbw = formed_beam.weight[:]

        # For each source, beamform and populate container.
        for src in range(self.nsource):

            if src % 1000 == 0:
                self.log.debug(f"Source {src}/{self.nsource}")

            # Declination of this source
            dec = np.radians(self.sdec[src])

            if self.freqside is not None:
                # Get the frequency bin this source is closest to.
                freq_diff = abs(self.freq["centre"] - self.sfreq[src])
                sfreq_index = np.argmin(freq_diff)
                # Start and stop indices to process in global frequency axis
                freq_idx0 = np.amax([0, sfreq_index - self.freqside])
                freq_idx1 = np.amin(
                    [self.nfreq, sfreq_index + self.freqside + 1])
                # Mask in full frequency axis
                f_mask = np.ones(self.nfreq, dtype=bool)
                f_mask[freq_idx0:freq_idx1] = False
                # Restrict frequency mask to local range
                f_mask = f_mask[self.lo:(self.lo + self.ls)]

                # TODO: In principle I should be able to skip
                # sources that have no indices to be processed
                # in this rank. I am getting a NaN error, however.
                # I may need an mpiutil.barrier() call before the
                # return statement.
                if f_mask.all():
                    # If there are no indices to be processed in
                    # the local frequency range, skip source.
                    continue

                # Frequency indices to process in local range
                f_local_indices = np.arange(self.ls,
                                            dtype=np.int32)[np.invert(f_mask)]

            if self.is_sstream:
                # Get RA bin this source is closest to.
                # Phasing will actually be done at src position.
                sra_index = np.searchsorted(self.ra, self.sra[src])
            else:
                # Cannot use searchsorted, because RA might not be
                # monotonically increasing. Slower.
                # Notice: in case there is more than one transit,
                # this will pick a single transit quasi-randomly!
                transit_diff = abs(self.ra - self.sra[src])
                sra_index = np.argmin(transit_diff)
                # For now, skip sources that do not transit in the data
                ra_cadence = self.ra[1] - self.ra[0]
                if transit_diff[sra_index] > 1.5 * ra_cadence:
                    continue

            # Compute hour angle array
            ha_array, ra_index_range, ha_mask = self._ha_array(
                self.ra, sra_index, self.sra[src], self.ha_side,
                self.is_sstream)

            # Arrays to store beams and weights for this source
            # for all polarizations prior to combining polarizations
            if self.collapse_ha:
                formed_beam_full = np.zeros((self.npol, self.ls),
                                            dtype=np.float64)
                weight_full = np.zeros((self.npol, self.ls), dtype=np.float64)
            else:
                formed_beam_full = np.zeros((self.npol, self.ls, self.nha),
                                            dtype=np.float64)
                weight_full = np.zeros((self.npol, self.ls, self.nha),
                                       dtype=np.float64)
            # For each polarization
            for pol in range(self.npol):

                # Compute primary beams to be used in the weighting
                primary_beam = self._beamfunc(
                    ha_array[np.newaxis, :],
                    self.process_pol[pol],
                    self.freq_local[:, np.newaxis],
                    dec,
                )

                # Fringestop and sum over products
                # 'beamform' does not normalize sum.
                this_formed_beam = beamform(
                    self.vis[pol],
                    self.sumweight[pol],
                    dec,
                    self.latitude,
                    np.cos(ha_array),
                    np.sin(ha_array),
                    self.bvec[pol][0],
                    self.bvec[pol][1],
                    f_local_indices,
                    ra_index_range,
                )

                sumweight_inrange = self.sumweight[pol][:, ra_index_range, :]
                visweight_inrange = self.visweight[pol][:, ra_index_range, :]

                if self.collapse_ha:
                    # Sum over RA. Does not multiply by weights because
                    # this_formed_beam was never normalized (this avoids
                    # re-work and makes code more efficient).
                    this_sumweight = np.sum(
                        np.sum(sumweight_inrange, axis=-1) * primary_beam**2,
                        axis=1)

                    formed_beam_full[pol] = np.sum(
                        this_formed_beam * primary_beam,
                        axis=1) * invert_no_zero(this_sumweight)

                    if self.weight != "inverse_variance":
                        this_weight2 = np.sum(
                            np.sum(
                                sumweight_inrange**2 *
                                invert_no_zero(visweight_inrange),
                                axis=-1,
                            ) * primary_beam**2,
                            axis=1,
                        )

                        weight_full[pol] = this_sumweight**2 * invert_no_zero(
                            this_weight2)

                    else:
                        weight_full[pol] = this_sumweight

                else:
                    # Need to divide by weight here for proper
                    # normalization because it is not done in
                    # beamform()
                    this_sumweight = np.sum(sumweight_inrange, axis=-1)
                    # Populate only where ha_mask is true. Zero otherwise.
                    formed_beam_full[
                        pol][:, ha_mask] = this_formed_beam * invert_no_zero(
                            this_sumweight)
                    if self.weight != "inverse_variance":
                        this_weight2 = np.sum(
                            sumweight_inrange**2 *
                            invert_no_zero(visweight_inrange),
                            axis=-1,
                        )
                        # Populate only where ha_mask is true. Zero otherwise.
                        weight_full[
                            pol][:,
                                 ha_mask] = this_sumweight**2 * invert_no_zero(
                                     this_weight2)
                    else:
                        weight_full[pol][:, ha_mask] = this_sumweight

                # Ensure weights are zero for non-processed frequencies
                weight_full[pol][f_mask] = 0.0

            # Combine polarizations if needed.
            # TODO: For now I am ignoring differences in the X and
            # Y beams and just adding them as is.
            if self.polarization == "I":
                formed_beam_full = np.sum(formed_beam_full * weight_full,
                                          axis=0) * invert_no_zero(
                                              np.sum(weight_full, axis=0))
                weight_full = np.sum(weight_full, axis=0)
                # Add an axis for the polarization
                if self.collapse_ha:
                    formed_beam_full = np.reshape(formed_beam_full,
                                                  (1, self.ls))
                    weight_full = np.reshape(weight_full, (1, self.ls))
                else:
                    formed_beam_full = np.reshape(formed_beam_full,
                                                  (1, self.ls, self.nha))
                    weight_full = np.reshape(weight_full,
                                             (1, self.ls, self.nha))
            elif self.polarization == "stokes":
                # TODO: Not implemented
                pass

            # Populate container.
            fbb[src] = formed_beam_full
            fbw[src] = weight_full
            if not self.collapse_ha:
                if self.is_sstream:
                    formed_beam.ha[src, :] = ha_array
                else:
                    # Populate only where ha_mask is true.
                    formed_beam.ha[src, ha_mask] = ha_array

        return formed_beam

    def _ha_side(self, data, timetrack=900.0):
        """Number of RA/time bins to track the source at each side of transit.

        Parameters
        ----------
        data : `containers.SiderealStream` or `containers.TimeStream`
            Data to read time from.
        timetrack : float
            Time in seconds to track at each side of transit.
            Default is 15 minutes.

        Returns
        -------
        ha_side : int
            Number of RA bins to track the source at each side of transit.
        """
        # TODO: Instead of a fixed time for transit, I could have a minimum
        # drop in the beam at a conventional distance from the NCP.
        if "ra" in data.index_map:
            # In seconds
            approx_time_perbin = 24.0 * 3600.0 / float(
                len(data.index_map["ra"]))
        else:
            approx_time_perbin = data.time[1] - data.time[0]

        # Track for `timetrack` seconds at each side of transit
        return int(timetrack / approx_time_perbin)

    def _ha_array(self,
                  ra,
                  source_ra_index,
                  source_ra,
                  ha_side,
                  is_sstream=True):
        """Hour angle for each RA/time bin to be processed.

        Also return the indices of these bins in the full RA/time axis.

        Parameters
        ----------
        ra : array
            RA axis in the data
        source_ra_index : int
            Index in data.index_map['ra'] closest to source_ra
        source_ra : float
            RA of the quasar
        ha_side : int
            Number of RA/HA bins on each side of transit.
        is_sstream : bool
            True if data is sidereal stream. Flase if time stream

        Returns
        -------
        ha_array : np.ndarray
            Hour angle array in the range -180. to 180
        ra_index_range : np.ndarray of int
            Indices (in data.index_map['ra']) corresponding
            to ha_array.
        """
        # RA range to track this quasar through the beam.
        ra_index_range = np.arange(source_ra_index - ha_side,
                                   source_ra_index + ha_side + 1,
                                   dtype=np.int32)
        # Number of RA bins in data.
        nra = len(ra)

        if is_sstream:
            # Wrap RA indices around edges.
            ra_index_range[ra_index_range < 0] += nra
            ra_index_range[ra_index_range >= nra] -= nra
            # Hour angle array (convert to radians)
            ha_array = np.deg2rad(ra[ra_index_range] - source_ra)
            # For later convenience it is better if `ha_array` is
            # in the range -pi to pi instead of 0 to 2pi.
            ha_array = (ha_array + np.pi) % (2.0 * np.pi) - np.pi
            # In this case the ha_mask is trivial
            ha_mask = np.ones(len(ra_index_range), dtype=bool)
        else:
            # Mask-out indices out of range
            ha_mask = (ra_index_range >= 0) & (ra_index_range < nra)
            # Return smaller HA range, and mask.
            ra_index_range = ra_index_range[ha_mask]
            # Hour angle array (convert to radians)
            ha_array = np.deg2rad(ra[ra_index_range] - source_ra)
            # For later convenience it is better if `ha_array` is
            # in the range -pi to pi instead of 0 to 2pi.
            ha_array = (ha_array + np.pi) % (2.0 * np.pi) - np.pi

        return ha_array, ra_index_range, ha_mask

    # TODO: This is very CHIME specific. Should probably be moved somewhere else.
    def _beamfunc(self, ha, pol, freq, dec, zenith=0.70999994):
        """Simple and fast beam model to be used as beamforming weights.

        Parameters
        ----------
        ha : array or float
            Hour angle (in radians) to compute beam at.
        freq : array or float
            Frequency in MHz
        dec : array or float
            Declination in radians
        pol : int or string
            Polarization index. 0: X, 1: Y, >=2: XY
            or one of 'XX', 'XY', 'YX', 'YY'
        zenith : float
            Polar angle of the telescope zenith in radians.
            Equal to pi/2 - latitude

        Returns
        -------
        beam : array or float
            The beam at the designated hhour angles, frequencies
            and declinations. This is the beam 'power', that is,
            voltage squared. To get the beam voltage, take the
            square root.
        """

        pollist = ["XX", "YY", "XY", "YX"]
        if pol in pollist:
            pol = pollist.index(pol)

        def _sig(pp, freq, dec):
            sig_amps = [14.87857614, 9.95746878]
            return sig_amps[pp] / freq / np.cos(dec)

        def _amp(pp, dec, zenith):
            def _flat_top_gauss6(x, A, sig, x0):
                """Flat-top gaussian. Power of 6."""
                return A * np.exp(-abs((x - x0) / sig)**6)

            def _flat_top_gauss3(x, A, sig, x0):
                """Flat-top gaussian. Power of 3."""
                return A * np.exp(-abs((x - x0) / sig)**3)

            prm_ns_x = np.array([9.97981768e-01, 1.29544939e00, 0.0])
            prm_ns_y = np.array([9.86421047e-01, 8.10213326e-01, 0.0])

            if pp == 0:
                return _flat_top_gauss6(dec - (0.5 * np.pi - zenith),
                                        *prm_ns_x)
            else:
                return _flat_top_gauss3(dec - (0.5 * np.pi - zenith),
                                        *prm_ns_y)

        ha0 = 0.0
        if pol < 2:
            # XX or YY
            return _amp(pol, dec, zenith) * np.exp(-((
                (ha - ha0) / _sig(pol, freq, dec))**2))
        else:
            # XY or YX
            return (_amp(0, dec, zenith) *
                    np.exp(-(((ha - ha0) / _sig(0, freq, dec))**2)) *
                    _amp(1, dec, zenith) *
                    np.exp(-(((ha - ha0) / _sig(1, freq, dec))**2)))**0.5

    def _process_data(self, data):
        """Store code for parsing and formating data prior to beamforming."""
        # Easy access to communicator
        self.comm_ = data.comm

        # Extract data info
        if "ra" in data.index_map:
            self.is_sstream = True
            self.ra = data.index_map["ra"]

            # Calculate the epoch for the data so we can calculate the correct
            # CIRS coordinates
            if "lsd" not in data.attrs:
                raise ValueError(
                    "SiderealStream must have an LSD attribute to calculate the epoch."
                )

            # This will be a float for a single sidereal day, or a list of
            # floats for a stack
            lsd = (data.attrs["lsd"][0] if isinstance(
                data.attrs["lsd"], np.ndarray) else data.attrs["lsd"])
            self.epoch = self.telescope.lsd_to_unix(lsd)

        else:
            self.is_sstream = False
            # Convert data timestamps into LSAs (degrees)
            self.ra = self.telescope.unix_to_lsa(data.time)
            self.epoch = data.time.mean()

        self.freq = data.index_map["freq"]
        self.nfreq = len(self.freq)
        # Ensure data is distributed in freq axis
        data.redistribute(0)

        # Number of RA bins to track each source at each side of transit
        self.ha_side = self._ha_side(data, self.timetrack)
        self.nha = 2 * self.ha_side + 1

        # polmap: indices of each vis product in
        # polarization list: ['XX', 'XY', 'YX', 'YY']
        polmap = polarization_map(data.index_map, self.telescope)
        # Baseline vectors in meters
        bvec_m = baseline_vector(data.index_map, self.telescope)

        # MPI distribution values
        self.lo = data.vis.local_offset[0]
        self.ls = data.vis.local_shape[0]
        self.freq_local = self.freq["centre"][self.lo:self.lo + self.ls]
        # These are to be used when gathering results in the end.
        # Tuple (not list!) of number of frequencies in each rank
        self.fsize = tuple(data.comm.allgather(self.ls))
        # Tuple (not list!) of displacements of each rank array in full array
        self.foffset = tuple(data.comm.allgather(self.lo))

        fullpol = ["XX", "XY", "YX", "YY"]
        # Save subsets of the data for each polarization, changing
        # the ordering to 'C' (needed for the cython part).
        # This doubles the memory usage.
        self.vis, self.visweight, self.bvec, self.sumweight = [], [], [], []
        for pol in self.process_pol:
            pol = fullpol.index(pol)
            polmask = polmap == pol
            # Swap order of product(1) and RA(2) axes, to reduce striding
            # through memory later on.
            self.vis.append(
                np.copy(np.moveaxis(data.vis[:, polmask, :], 1, 2), order="C"))
            # Restrict visweight to the local frequencies
            self.visweight.append(
                np.copy(
                    np.moveaxis(
                        data.weight[self.lo:self.lo + self.ls][:, polmask, :],
                        1, 2).astype(np.float64),
                    order="C",
                ))
            # Multiply bvec_m by frequencies to get vector in wavelengths.
            # Shape: (2, nfreq_local, nvis), for each pol.
            self.bvec.append(
                np.copy(
                    bvec_m[:, np.newaxis, polmask] *
                    self.freq_local[np.newaxis, :, np.newaxis] * 1e6 / C,
                    order="C",
                ))
            if self.weight == "inverse_variance":
                # Weights for sum are just the visibility weights
                self.sumweight.append(self.visweight[-1])
            else:
                # Ensure zero visweights result in zero sumweights
                this_sumweight = (self.visweight[-1] > 0.0).astype(np.float64)
                ssi = data.input_flags[:]
                ssp = data.index_map["prod"][:]
                sss = data.reverse_map["stack"]["stack"][:]
                nstack = data.vis.shape[1]
                # this redundancy takes into account input flags.
                # It has shape (nstack, ntime)
                redundancy = np.moveaxis(
                    calculate_redundancy(ssi, ssp, sss,
                                         nstack)[polmask].astype(np.float64),
                    0,
                    1,
                )[np.newaxis, :, :]
                # redundancy = (self.telescope.redundancy[polmask].
                #        astype(np.float64)[np.newaxis, np.newaxis, :])
                this_sumweight *= redundancy
                if self.weight == "uniform":
                    this_sumweight = (this_sumweight > 0.0).astype(np.float64)
                self.sumweight.append(np.copy(this_sumweight, order="C"))

    def _process_catalog(self, catalog):
        """Process the catalog to get CIRS coordinates at the correct epoch.

        Note that `self._process_data` must have been called before this.
        """

        if "position" not in catalog:
            raise ValueError("Input is missing a position table.")

        self.sra, self.sdec = icrs_to_cirs(catalog["position"]["ra"],
                                           catalog["position"]["dec"],
                                           self.epoch)

        if self.freqside is not None:
            if "redshift" not in catalog:
                raise ValueError("Input is missing a required redshift table.")
            self.sfreq = NU21 / (catalog["redshift"]["z"][:] + 1.0)  # MHz

        self.source_cat = catalog
        self.nsource = len(self.sra)
Example #11
0
class CHIME(telescope.PolarisedTelescope):
    """Model telescope for the CHIME/Pathfinder.

    This class currently uses a simple Gaussian model for the primary beams.

    Attributes
    ----------
    layout : datetime or int
        Specify which layout to use.
    correlator : string
        Restrict to a specific correlator.
    skip_non_chime : boolean
        Ignore non CHIME feeds in the BeamTransfers.
    stack_type : string, optional
        Stacking type.
        `redundant`: feeds of same polarization have same beam class (default).
        `redundant_cyl`: feeds of same polarization and cylinder have same beam
        class.
        `unique`: Each feed has a unique beam class.
    use_pathfinder_freq: boolean
        Use the pathfinder channelization of 1024 frequencies between 400 and
        800 MHz.  Setting this to True also enables the specification of a
        subset of these frequencies through the four attributes below.  Default
        is True.
    channel_bin : int, optional
        Number of channels to bin together. Must exactly divide the total
        number. Binning is performed prior to selection of any subset. Default
        is 1.
    freq_physical : list, optional
        Select subset of frequencies using a list of physical frequencies in
        MHz. Finds the closests pathfinder channel.
    channel_range : list, optional
        Select subset of frequencies using a range of frequency channel indices,
        either [start, stop, step], [start, stop], or [stop] is acceptable.
    channel_index : list, optional
        Select subset of frequencies using a list of frequency channel indices.
    input_sel : list, optional
        Select a reduced set of feeds to use. Useful for generating small
        subsets of the data.
    baseline_masking_type : string, optional
        Select a subset of baselines.  `total_length` selects baselines according to
        their total length. Need to specify `minlength` and `maxlength` properties
        (defined in baseclass).  `individual_length` selects baselines according to
        their seperation in the North-South (specify `minlength_ns` and `maxlength_ns`)
        or the East-West (specify `minlength_ew` and `maxlength_ew`).
    minlength_ns, maxlength_ns : float
        Minimum and maximum North-South baseline lengths to include (in metres)
    minlength_ew, maxlength_ew: float
        Minimum and maximum East-West baseline lengths to include (in metres)
    dec_normalized: float, optional
        Normalize the beam by its magnitude at transit at this declination
        in degrees.
    skip_pol_pair : list
        List of antenna polarisation pairs to skip. Valid entries are "XX", "XY", "YX"
        or "YY". Like the skipped frequencies these pol pairs will have entries
        generated but their beam transfer matrices are implicitly zero and thus not
        calculated.
    """

    # Configure which feeds and layout to use
    layout = config.Property(default=None)
    correlator = config.Property(proptype=str, default=None)
    skip_non_chime = config.Property(proptype=bool, default=False)

    # Redundancy settings
    stack_type = config.enum(["redundant", "redundant_cyl", "unique"],
                             default="redundant")

    # Configure frequency properties
    use_pathfinder_freq = config.Property(proptype=bool, default=True)
    channel_bin = config.Property(proptype=int, default=1)
    freq_physical = config.Property(proptype=list, default=[])
    channel_range = config.Property(proptype=list, default=[])
    channel_index = config.Property(proptype=list, default=[])

    # Input selection
    input_sel = config.Property(proptype=list, default=None)

    # Baseline masking options
    baseline_masking_type = config.enum(["total_length", "individual_length"],
                                        default="individual_length")
    minlength_ew = config.Property(proptype=float, default=0.0)
    maxlength_ew = config.Property(proptype=float, default=1.0e7)
    minlength_ns = config.Property(proptype=float, default=0.0)
    maxlength_ns = config.Property(proptype=float, default=1.0e7)

    # Auto-correlations setting (overriding default in baseclass)
    auto_correlations = config.Property(proptype=bool, default=True)

    # Beam normalization
    dec_normalized = config.Property(proptype=float, default=None)
    # Skipping frequency/baseline parameters
    skip_pol_pair = config.list_type(type_=str, maxlength=4, default=[])

    # Fix base properties
    cylinder_width = 20.0
    cylinder_spacing = tools._PF_SPACE

    _exwidth = [0.7]
    _eywidth = _exwidth

    _hxwidth = [1.2]
    _hywidth = _hxwidth

    _pickle_keys = ["_feeds"]

    #
    # === Initialisation routines ===
    #

    def __init__(self, feeds=None):
        import datetime

        self._feeds = feeds

        # Set location properties
        self.latitude = ephemeris.CHIMELATITUDE
        self.longitude = ephemeris.CHIMELONGITUDE
        self.altitude = ephemeris.CHIMEALTITUDE

        # Set the LSD start epoch (i.e. CHIME/Pathfinder first light)
        self.lsd_start_day = datetime.datetime(2013, 11, 15)

        # Set the overall normalization of the beam
        self._set_beam_normalization()

    @classmethod
    def from_layout(cls, layout, correlator=None, skip=False):
        """Create a CHIME/Pathfinder telescope description for the specified layout.

        Parameters
        ----------
        layout : integer or datetime
            Layout id number (corresponding to one in database), or datetime
        correlator : string, optional
            Name of the specific correlator. Needed to return a unique config
            in some cases.
        skip : boolean, optional
            Whether to skip non-CHIME antennas. If False, leave them in but
            set them to infinite noise (unsupported at the moment).

        Returns
        -------
        tel : CHIME
        """

        tel = cls()

        tel.layout = layout
        tel.correlator = correlator
        tel.skip_non_chime = skip
        tel._load_layout()

        return tel

    def _load_layout(self):
        """Load the CHIME/Pathfinder layout from the database.

        Generally this routine shouldn't be called directly. Use
        :method:`CHIME.from_layout` or configure from a YAML file.
        """
        if self.layout is None:
            raise Exception("Layout attributes not set.")

        # Fetch feed layout from database
        feeds = tools.get_correlator_inputs(self.layout, self.correlator)

        if mpiutil.size > 1:
            feeds = mpiutil.world.bcast(feeds, root=0)

        if self.skip_non_chime:
            raise Exception("Not supported.")

        self._feeds = feeds

    def _finalise_config(self):
        # Override base method to implement automatic loading of layout when
        # configuring from YAML.

        if self.layout is not None:
            logger.debug("Loading layout: %s", str(self.layout))
            self._load_layout()

        # Set the overall normalization of the beam
        self._set_beam_normalization()

    #
    # === Redefine properties of the base class ===
    #

    # Tweak the following two properties to change the beam width
    @cached_property
    def fwhm_ex(self):
        """Full width half max of the E-plane antenna beam for X polarization."""

        return np.polyval(
            np.array(self._exwidth) * 2.0 * np.pi / 3.0, self.frequencies)

    @cached_property
    def fwhm_hx(self):
        """Full width half max of the H-plane antenna beam for X polarization."""

        return np.polyval(
            np.array(self._hxwidth) * 2.0 * np.pi / 3.0, self.frequencies)

    @cached_property
    def fwhm_ey(self):
        """Full width half max of the E-plane antenna beam for Y polarization."""

        return np.polyval(
            np.array(self._eywidth) * 2.0 * np.pi / 3.0, self.frequencies)

    @cached_property
    def fwhm_hy(self):
        """Full width half max of the H-plane antenna beam for Y polarization."""

        return np.polyval(
            np.array(self._hywidth) * 2.0 * np.pi / 3.0, self.frequencies)

    # Set the approximate uv feed sizes
    @property
    def u_width(self):
        return self.cylinder_width

    # v-width property override
    @property
    def v_width(self):
        return 1.0

    # Set non-zero rotation angle for pathfinder and chime
    @property
    def rotation_angle(self):
        if self.correlator == "pathfinder":
            return tools._PF_ROT
        elif self.correlator == "chime":
            return tools._CHIME_ROT
        else:
            return 0.0

    def calculate_frequencies(self):
        """Override default version to give support for specifying by frequency
        channel number.
        """
        if self.use_pathfinder_freq:
            # Use pathfinder channelization of 1024 bins between 400 and 800 MHz.
            basefreq = np.linspace(800.0, 400.0, 1024, endpoint=False)

            # Bin the channels together
            if len(basefreq) % self.channel_bin != 0:
                raise Exception(
                    "Channel binning must exactly divide the total number of channels"
                )

            basefreq = basefreq.reshape(-1, self.channel_bin).mean(axis=-1)

            # If requested, select subset of frequencies.
            if self.freq_physical:
                basefreq = basefreq[[
                    np.argmin(np.abs(basefreq - freq))
                    for freq in self.freq_physical
                ]]

            elif self.channel_range and (len(self.channel_range) <= 3):
                basefreq = basefreq[slice(*self.channel_range)]

            elif self.channel_index:
                basefreq = basefreq[self.channel_index]

            # Save to object
            self._frequencies = np.unique(basefreq)[::-1]

        else:
            # Otherwise use the standard method
            telescope.TransitTelescope.calculate_frequencies(self)

    @property
    def feeds(self):
        """Return a description of the feeds as a list of :class:`tools.CorrInput` instances."""

        if self.input_sel is None:
            feeds = self._feeds
        else:
            feeds = [self._feeds[fi] for fi in self.input_sel]

        return feeds

    @property
    def input_index(self):
        """An index_map describing the inputs known to the telescope. Useful
        for generating synthetic datasets.
        """
        # Extract lists of channel ID and serial numbers
        channels, feed_sn = list(
            zip(*[(feed.id, feed.input_sn) for feed in self.feeds]))

        # Create an input index map and return it.
        from ch_util import andata

        return andata._generate_input_map(feed_sn, channels)

    _pos = None

    @property
    def feedpositions(self):
        """The set of feed positions on *all* cylinders.

        This is constructed for the given layout and includes all rotations of
        the cylinder axis.

        Returns
        -------
        feedpositions : np.ndarray
            The positions in the telescope plane of the receivers. Packed as
            [[u1, v1], [u2, v2], ...].
        """

        if self._pos is None:
            # Fetch cylinder relative positions
            pos = tools.get_feed_positions(self.feeds)

            # The above routine returns NaNs for non CHIME feeds. This is a bit
            # messy, so turn them into zeros.
            self._pos = np.nan_to_num(pos)

        return self._pos

    @property
    def beamclass(self):
        """Beam class definition for the CHIME/Pathfinder.

        When `self.stack_type` is `redundant`, the X-polarisation feeds get
        `beamclass = 0`, and the Y-polarisation gets `beamclass = 1`.
        When `self.stack_type` is `redundant_cyl`, feeds of same polarisation
        and cylinder have same beam class. The beam class is given by
        `beamclass = 2*cyl + pol` where `cyl` is the cylinder number according to
        `ch_util.tools` convention and `pol` is the polarisation (0 for X and 1
        for Y polarisation)
        When `self.stack_type` is `unique`, then the feeds are just given an
        increasing unique class.
        In all cases, any other type of feed gets set to `-1` and should be
        ignored.
        """

        # Make beam class just channel number.

        def _feedclass(f, redundant_cyl=False):
            if tools.is_array(f):
                if tools.is_array_x(f):  # feed is X polarisation
                    pol = 0
                else:  # feed is Y polarisation
                    pol = 1

                if redundant_cyl:
                    return 2 * f.cyl + pol
                else:
                    return pol
            return -1

        if self.stack_type == "redundant":
            return np.array([_feedclass(f) for f in self.feeds])
        elif self.stack_type == "redundant_cyl":
            return np.array(
                [_feedclass(f, redundant_cyl=True) for f in self.feeds])
        else:
            beamclass = [
                fi if tools.is_array(feed) else -1
                for fi, feed in enumerate(self.feeds)
            ]
            return np.array(beamclass)

    @property
    def polarisation(self):
        """
        Polarisation map.

        Returns
        -------
        pol : np.ndarray
            One-dimensional array with the polarization for each feed ('X' or 'Y').
        """
        def _pol(f):
            if tools.is_array(f):
                if tools.is_array_x(f):  # feed is X polarisation
                    return "X"
                else:  # feed is Y polarisation
                    return "Y"
            return "N"

        return np.asarray([_pol(f) for f in self.feeds], dtype=np.str)

    #
    # === Setup the primary beams ===
    #

    def beam(self, feed, freq, angpos=None):
        """Primary beam implementation for the CHIME/Pathfinder.

        This only supports normal CHIME cylinder antennas. Asking for the beams
        for other types of inputs will cause an exception to be thrown. The
        beams from this routine are rotated by `self.rotation_angle` to account
        for the CHIME/Pathfinder rotation.

        Parameters
        ----------
        feed : int
            Index for the feed.
        freq : int
            Index for the frequency.
        angpos : np.ndarray[nposition, 2], optional
            Angular position on the sky (in radians).
            If not provided, default to the _angpos
            class attribute.

        Returns
        -------
        beam : np.ndarray[nposition, 2]
            Amplitude vector of beam at each position on the sky.
        """
        # # Fetch beam parameters out of config database.

        feed_obj = self.feeds[feed]

        # Check that feed exists and is a CHIME cylinder antenna
        if feed_obj is None:
            raise ValueError(
                "Craziness. The requested feed doesn't seem to exist.")

        if not tools.is_array(feed_obj):
            raise ValueError("Requested feed is not a CHIME antenna.")

        # If the angular position was not provided, then use the values in the
        # class attribute.
        if angpos is None:
            angpos = self._angpos

        # Get the beam rotation parameters.
        yaw = -self.rotation_angle
        pitch = 0.0
        roll = 0.0

        rot = np.radians([yaw, pitch, roll])

        # We can only support feeds angled parallel or perp to the cylinder
        # axis. Check for these and throw exception for anything else.
        if tools.is_array_y(feed_obj):
            beam = cylbeam.beam_y(
                angpos,
                self.zenith,
                self.cylinder_width / self.wavelengths[freq],
                self.fwhm_ey[freq],
                self.fwhm_hy[freq],
                rot=rot,
            )
        elif tools.is_array_x(feed_obj):
            beam = cylbeam.beam_x(
                angpos,
                self.zenith,
                self.cylinder_width / self.wavelengths[freq],
                self.fwhm_ex[freq],
                self.fwhm_hx[freq],
                rot=rot,
            )
        else:
            raise RuntimeError(
                "Given polarisation (feed.pol=%s) not supported." %
                feed_obj.pol)

        # Normalize the beam
        if self._beam_normalization is not None:
            beam *= self._beam_normalization[freq, feed, np.newaxis, :]

        return beam

    #
    # === Override methods determining the feed pairs we should calculate ===
    #
    # These should probably get ported back into `driftscan` as options.

    def _sort_pairs(self):
        # Reimplemented sort pairs to ensure that returned array is in
        # channel order.

        # Create mask of included pairs, that are not conjugated
        tmask = np.logical_and(self._feedmask, np.logical_not(self._feedconj))
        uniq = telescope._get_indices(self._feedmap, mask=tmask)

        # Get channel id for each feed in the pair, this will be used for the sort
        ci, cj = np.array([(self.feeds[fi].id, self.feeds[fj].id)
                           for fi, fj in uniq]).T

        # # Sort by constructing a numpy array with the keys as fields, and use
        # # np.argsort to get the indices

        # Create array of keys to sort
        dt = np.dtype("i4,i4")
        sort_arr = np.zeros(ci.size, dtype=dt)
        sort_arr["f0"] = ci
        sort_arr["f1"] = cj

        # Get map which sorts
        sort_ind = np.argsort(sort_arr)

        # Invert mapping
        tmp_sort_ind = sort_ind.copy()
        sort_ind[tmp_sort_ind] = np.arange(sort_ind.size)

        # Remap feedmap entries
        fm_copy = self._feedmap.copy()
        wmask = np.where(self._feedmask)
        fm_copy[wmask] = sort_ind[self._feedmap[wmask]]

        self._feedmap = fm_copy

    def _make_ew(self):
        # # Reimplemented to make sure entries we always pick the upper
        # # triangle (and do not reorder to make EW baselines)
        if self.stack_type != "unique":
            super(CHIME, self)._make_ew()

    def _unique_baselines(self):
        # Reimplement unique baselines in order to mask out either according to total
        # baseline length or maximum North-South and East-West baseline seperation.

        from drift.core import telescope

        # Construct array of indices
        fshape = [self.nfeed, self.nfeed]
        f_ind = np.indices(fshape)

        # Construct array of baseline separations
        bl1 = self.feedpositions[f_ind[0]] - self.feedpositions[f_ind[1]]
        bl2 = np.around(bl1[..., 0] + 1.0j * bl1[..., 1], self._bl_tol)

        # Construct array of baseline lengths
        blen = np.sum(bl1**2, axis=-1)**0.5

        if self.baseline_masking_type == "total_length":
            # Create mask of included baselines
            mask = np.logical_and(blen >= self.minlength,
                                  blen <= self.maxlength)
        else:
            mask_ew = np.logical_and(
                abs(bl1[..., 0]) >= self.minlength_ew,
                abs(bl1[..., 0]) <= self.maxlength_ew,
            )
            mask_ns = np.logical_and(
                abs(bl1[..., 1]) >= self.minlength_ns,
                abs(bl1[..., 1]) <= self.maxlength_ns,
            )
            mask = np.logical_and(mask_ew, mask_ns)

        # Remove the auto correlated baselines between all polarisations
        if not self.auto_correlations:
            mask = np.logical_and(blen > 0.0, mask)

        return telescope._remap_keyarray(bl2, mask), mask

    def _unique_beams(self):
        # Override to mask out any feed where the beamclass is less than zero.
        # This is used to get exclude feeds which are not normal CHIME cylinder
        # feeds

        beam_map, beam_mask = telescope.TransitTelescope._unique_beams(self)

        # Construct a mask including only the feeds where the beam class is
        # greater than zero
        bc_mask = self.beamclass >= 0
        bc_mask = np.logical_and(bc_mask[:, np.newaxis],
                                 bc_mask[np.newaxis, :])

        beam_mask = np.logical_and(beam_mask, bc_mask)

        return beam_map, beam_mask

    def _set_beam_normalization(self):
        """Determine the beam normalization for each feed and frequency.

        The beam will be normalized by its value at transit at the declination
        provided in the dec_normalized config parameter.  If this config parameter
        is set to None, then there is no additional normalization applied.
        """

        self._beam_normalization = None

        if self.dec_normalized is not None:

            angpos = np.array([(0.5 * np.pi - np.radians(self.dec_normalized)),
                               0.0]).reshape(1, -1)

            beam = np.ones((self.nfreq, self.nfeed, 2), dtype=np.float64)

            beam_lookup = {}

            for fe, feed in enumerate(self.feeds):

                if not tools.is_array(feed):
                    continue

                beamclass = self.beamclass[fe]

                if beamclass not in beam_lookup:

                    beam_lookup[beamclass] = np.ones((self.nfreq, 2),
                                                     dtype=np.float64)
                    for fr in range(self.nfreq):
                        beam_lookup[beamclass][fr] = self.beam(fe, fr,
                                                               angpos)[0]

                beam[:, fe, :] = beam_lookup[beamclass]

            self._beam_normalization = tools.invert_no_zero(
                np.sqrt(np.sum(beam**2, axis=-1)))

    def _skip_baseline(self, bl_ind):
        """Override to skip baselines based on which polarisation pair they are."""

        # Pull in baseline skip choice from parent class
        skip_bl = super()._skip_baseline(bl_ind)

        pol_i, pol_j = self.polarisation[self.uniquepairs[bl_ind]]
        pol_pair = pol_i + pol_j

        skip_pol = pol_pair in self.skip_pol_pair

        return skip_bl or skip_pol