Source code for realtime.receive.core.baseline_utils

# -*- coding: utf-8 -*-
"""Baseline-related functions """

import itertools
import logging
from collections.abc import Reversible
from typing import Sequence, TypeVar, Union

logger = logging.getLogger(__name__)




def baselines(stations: Union[int, Sequence], autocorr: bool):
    """
    Return the number of baselines for the given number of stations, depending
    on whether stations are autocorrelated or not.

    :param stations: The number of stations, or a sequence of stations, in which
      case its length is used.
    :param autocorr: Whether autocorrelated baselines should be accounted for.
    """
    if not isinstance(stations, int):
        stations = len(stations)
    if autocorr:
        return (stations * (stations + 1)) // 2
    else:
        return (stations * (stations - 1)) // 2





def baseline_indices(
    num_stations: int, lower_triangular: bool, autocorr: bool = True
) -> tuple[tuple[int, int]]:
    """
    Generates baseline indices for the given number of stations.

    :param num_stations: Number of stations to generate baselines for.
    :param lower_triangular: Whether baseline indices should be generated as if
     they were laid in a lower or upper triangular matrix.
    :param autocorr: Whether to generate indices for autocorrelation.
    :return: A tuple of 2-tuples containing the station1/station2 indices
     making up each baseline.
    """
    return generate_baselines(
        range(num_stations), lower_triangular=lower_triangular, autocorr=autocorr
    )



AntennaT = TypeVar("AntennaT")




def generate_baselines(
    antennas: Sequence[AntennaT] | Reversible[AntennaT],
    lower_triangular: bool,
    autocorr: bool = True,
) -> tuple[tuple[AntennaT, AntennaT]]:
    """
    Generates baselines for the given elements.

    Baselines are generated differently depending on how antennas are
    traversed. The ``lower_triangular`` argument controls this, with a ``True``
    value indicating a sequencing that goes in row-major order through the
    following matrix::

     A1|
     0:| 0
     1:| 1 2
     2:| 3 4 5
     3:| 6 7 8 9
     --+--------
     A2| 0 1 2 3

    while a ``False`` value indicates the following sequencing::

     A1|
     0:| 0 1 2 3
     1:|   4 5 6
     2:|     7 8
     3:|       9
     --+--------
     A2| 0 1 2 3

    ``autocorr`` indicates whether baseline indices for autocorrelation (i.e.,
    the diagonal in the matrices above) should be generated or not.

    :param antennas: Antennas to generate baselines for.
    :param lower_triangular: Whether baseline indices should be generated as if
     they were laid in a lower or upper triangular matrix.
    :param autocorr: Whether to generate indices for autocorrelation.
    :return: A tuple of 2-tuples containing the antennas making up each
     baseline.
    """
    if autocorr:
        combinations = itertools.combinations_with_replacement
    else:
        combinations = itertools.combinations
    if lower_triangular:
        antenna_indices = reversed(antennas)
        indices = tuple(combinations(antenna_indices, 2))
        indices = tuple(reversed(indices))
        # indices == (0,0), (1,0), (1,1), (2,0), (2,1), ...
    else:
        antenna_indices = antennas
        indices = tuple(combinations(antenna_indices, 2))
        # indices == (0,0), (0,1), (0,2), (0,3), (1,1), ...
    assert len(indices) == baselines(len(antennas), autocorr)
    return indices





def are_autocorrelated_baselines(num_stations: int, num_baselines: int):
    """
    Returns `True` if the number of baselines corresponds to an autocorrelated
    configuration for the given number of stations.

    :param num_stations: The number of stations
    :param num_baselines: The number of baselines
    """
    if num_baselines == baselines(num_stations, True):
        return True
    if num_baselines == baselines(num_stations, False):
        return False
    else:
        raise ValueError(
            f"unexpected number of baselines ({num_baselines}) for number of stations ({num_stations})"
        )





def validate_antenna_and_baseline_counts(num_stations, num_baselines):
    """
    Check if the given number of baselines corresponds to a valid number of
    baselines for the given number of stations, and raise an exception if it
    doesn't.

    :param num_stations: The number of stations
    :param num_baselines: The number of baselines
    """
    are_autocorrelated_baselines(num_stations, num_baselines)