Exemplo n.º 1
0
def setup_uvdata(
    array_layout=None,
    telescope_location=None,
    telescope_name=None,
    Nfreqs=None,
    start_freq=None,
    bandwidth=None,
    freq_array=None,
    Ntimes=None,
    time_cadence=None,
    start_time=None,
    time_array=None,
    bls=None,
    anchor_ant=None,
    antenna_nums=None,
    no_autos=True,
    pols=["xx"],
    make_full=False,
    redundancy=None,
    run_check=True,
):
    """
    Setup a UVData object for simulating.

    Args:
        array_layout : str
            Filepath to array layout in ENU coordinates [meters]
        telescope_location : len-3 tuple
            Telescope location on Earth in LatLonAlt coordinates [deg, deg, meters]
        telescope_name : str
            Name of telescope
        Nfreqs : int
            Number of frequency channels
        start_freq : float
            Starting frequency [Hz]
        bandwidth : float
            Total frequency bandwidth of spectral window [Hz]
        freq_array : ndarray
            frequency array [Hz], cannot be specified if start_freq, Nfreqs or bandwidth is specified
        Ntimes : int
            Number of integration bins
        time_cadence : float
            Cadence of time bins [seconds]
        start_time : float
            Time of the first integration bin [Julian Date]
        time_array : ndarray
            time array [Julian Date], cannot be specified if start_time, Ntimes and time_cadence is specified
        bls : list
            List of antenna-pair tuples for baseline selection
        anchor_ant: int
            Selects baselines such that one of the pair is a specified antenna number
        redundancy: float
            Redundant baseline selection tolerance for selection
        antenna_nums : list
            List of antenna numbers to keep in array
        make_full : Generate the full UVData object, includes arrays of length Nblts.
                    Default behavior creates an invalid UVData object, where baseline_array has length Nbls, etc.
                    This is to avoid excessive memory usage up front when it's not necessary.

    Returns:
        UVData object with zeroed data_array
    """

    # get antenna information
    tele_dict = parse_telescope_params(
        dict(
            array_layout=array_layout,
            telescope_location=telescope_location,
            telescope_name=telescope_name,
        ))
    lat, lon, alt = uvutils.LatLonAlt_from_XYZ(tele_dict["telescope_location"])
    antpos = tele_dict["antenna_positions"]
    enu = uvutils.ENU_from_ECEF(
        tele_dict["antenna_positions"] + tele_dict["telescope_location"], lat,
        lon, alt)
    anums = tele_dict["antenna_numbers"]
    antnames = tele_dict["antenna_names"]
    Nants = tele_dict["Nants_data"]

    # setup object and fill in basic info
    uv_obj = UVData()
    uv_obj.telescope_location = tele_dict["telescope_location"]
    uv_obj.telescope_location_lat_lon_alt = (lat, lon, alt)
    uv_obj.telescope_location_lat_lon_alt_degrees = (
        np.degrees(lat),
        np.degrees(lon),
        alt,
    )
    uv_obj.antenna_numbers = anums
    uv_obj.antenna_names = antnames
    uv_obj.antenna_positions = antpos
    uv_obj.Nants_telescope = Nants

    # fill in frequency and time info: wait to fill in len-Nblts arrays until later
    if freq_array is not None:
        if Nfreqs is not None or start_freq is not None or bandwidth is not None:
            raise ValueError(
                "Cannot specify freq_array as well as Nfreqs, start_freq or bandwidth"
            )
        if freq_array.ndim == 1:
            freq_array = freq_array.reshape(1, -1)
            Nfreqs = freq_array.size
    else:
        freq_dict = parse_frequency_params(
            dict(Nfreqs=Nfreqs, start_freq=start_freq, bandwidth=bandwidth))
        freq_array = freq_dict["freq_array"]

    if time_array is not None:
        if Ntimes is not None or start_time is not None or time_cadence is not None:
            raise ValueError(
                "Cannot specify time_array as well as Ntimes, start_time or time_cadence"
            )
        Ntimes = time_array.size
    else:
        time_dict = parse_time_params(
            dict(Ntimes=Ntimes,
                 start_time=start_time,
                 time_cadence=time_cadence))
        time_array = time_dict["time_array"]

    uv_obj.freq_array = freq_array
    uv_obj.Nfreqs = Nfreqs
    uv_obj.Ntimes = Ntimes

    # fill in other attributes
    uv_obj.spw_array = np.array([0], dtype=int)
    uv_obj.Nspws = 1
    uv_obj.polarization_array = np.array(
        [uvutils.polstr2num(pol) for pol in pols], dtype=int)
    uv_obj.Npols = uv_obj.polarization_array.size
    if uv_obj.Nfreqs > 1:
        uv_obj.channel_width = np.diff(uv_obj.freq_array[0])[0]
    else:
        uv_obj.channel_width = 1.0
    uv_obj._set_drift()
    uv_obj.telescope_name = tele_dict["telescope_name"]
    uv_obj.instrument = "simulator"
    uv_obj.object_name = "zenith"
    uv_obj.vis_units = "Jy"
    uv_obj.history = ""

    if redundancy is not None:
        red_tol = redundancy
        reds, vec_bin_centers, lengths = uvutils.get_antenna_redundancies(
            anums, enu, tol=red_tol, include_autos=not bool(no_autos))
        bls = [r[0] for r in reds]
        bls = [uvutils.baseline_to_antnums(bl_ind, Nants) for bl_ind in bls]

    # Setup array and implement antenna/baseline selections.
    bl_array = []
    _bls = [(a1, a2) for a1 in anums for a2 in anums if a1 <= a2]
    if bls is not None:
        if isinstance(bls, str):
            bls = ast.literal_eval(bls)
        bls = [bl for bl in _bls if bl in bls]
    else:
        bls = _bls
    if anchor_ant is not None:
        bls = [bl for bl in bls if anchor_ant in bl]

    if bool(no_autos):
        bls = [bl for bl in bls if bl[0] != bl[1]]
    if antenna_nums is not None:
        if isinstance(antenna_nums, str):
            antenna_nums = ast.literal_eval(antenna_nums)
        if isinstance(antenna_nums, int):
            antenna_nums = [antenna_nums]
        bls = [(a1, a2) for (a1, a2) in bls
               if a1 in antenna_nums or a2 in antenna_nums]
    bls = sorted(bls)
    for (a1, a2) in bls:
        bl_array.append(uvutils.antnums_to_baseline(a1, a2, 1))
    bl_array = np.asarray(bl_array)
    print("Nbls: {}".format(bl_array.size))
    if bl_array.size == 0:
        raise ValueError("No baselines selected.")
    uv_obj.time_array = time_array  # Keep length Ntimes
    uv_obj.baseline_array = bl_array  # Length Nbls

    if make_full:
        uv_obj = complete_uvdata(uv_obj, run_check=run_check)

    return uv_obj