Source code for low_comm_tools.sdp.utils

from __future__ import annotations

from pathlib import Path
from typing import Any, NamedTuple

import everybeam as eb
import numpy as np
import numpy.typing as npt
from astropy.coordinates import ITRS, AltAz, Angle, EarthLocation, SkyCoord
from astropy.time import Time

try:
    from ska_sdp_datamodels.visibility import Visibility
except ImportError as e:

[docs]
    msg = "SKA-SDP tools are required. Install with [ska-sdp] pip extra"

    raise ImportError(msg) from e



[docs]
def phase(z: np.complexfloating[Any, Any]) -> np.floating[Any]:
    # return np.unwrap(np.angle(z)) * 180 / np.pi
    return np.angle(z) * 180 / np.pi  # type: ignore[no-any-return]

    # phi += 360 * (phi < -90)



[docs]
def radec_to_xyz(
    ra: Angle, dec: Angle, mjds: npt.NDArray[np.floating[Any]]
) -> npt.NDArray[np.floating[Any]]:
    """
    Convert RA and Dec ICRS coordinates to ITRS cartesian coordinates.
    See the Everybeam docs.

    Args:
        ra (astropy.coordinates.Angle): Right ascension
        dec (astropy.coordinates.Angle): Declination
        mjds (float): MJD time in seconds

    Returns:
        pointing_xyz (ndarray): NumPy array containing the ITRS X, Y and Z
        coordinates
    """
    obstime = Time(mjds / 86400.0, scale="utc", format="mjd")
    dir_pointing = SkyCoord(ra, dec)
    dir_pointing_itrs = dir_pointing.transform_to(ITRS(obstime=obstime))
    return np.asarray(dir_pointing_itrs.cartesian.xyz.transpose())




[docs]
class MetaData(NamedTuple):

[docs]
    time: Time


[docs]
    mjds: npt.NDArray[np.floating[Any]]


[docs]
    location: EarthLocation


[docs]
    nstation: int


[docs]
    stations: list[str]


[docs]
    zen_itrf: npt.NDArray[np.floating[Any]]


[docs]
    telescope: eb.Telescope


[docs]
    cos_term: npt.NDArray[np.floating[Any]]


[docs]
    beam_itrf: npt.NDArray[np.floating[Any]]


[docs]
    ant1: npt.NDArray[np.integer[Any]]


[docs]
    ant2: npt.NDArray[np.integer[Any]]





[docs]
def pre_calculate_metadata(
    dataset: Path,
    vis: Visibility,
) -> MetaData:
    # ============================================================================ #
    # pre-calculate some metadata for later

    location: EarthLocation = vis.configuration.location
    stations: list[str] = vis.configuration.stations.data
    nstation = len(stations)
    ant1 = vis.antenna1.data[vis.antenna1.data != vis.antenna2.data]
    ant2 = vis.antenna2.data[vis.antenna1.data != vis.antenna2.data]

    telescope = eb.load_telescope(dataset.as_posix())

    # metadata for beam at the central time step
    time = np.mean(Time(vis.datetime.data))
    mjds = time.mjd * 86400

    altaz = vis.phasecentre.transform_to(AltAz(obstime=time, location=location))
    # these are used in beam models and should be done separately for each station location
    # for our purposes though, just use a common location
    theta = np.pi / 2 - altaz.alt.radian
    cos_term = np.cos(theta)

    beam_itrf = radec_to_xyz(vis.phasecentre.ra, vis.phasecentre.dec, mjds)

    # ITRS coordinates for zenith at the central time step
    pointing = SkyCoord(
        alt=90,
        az=0,
        unit="deg",
        frame="altaz",
        obstime=time,
        location=location,
    ).transform_to(ITRS(obstime=time))
    zen_itrf = np.asarray(pointing.cartesian.xyz.transpose())

    return MetaData(
        time=time,
        mjds=mjds,
        location=location,
        stations=stations,
        nstation=nstation,
        zen_itrf=zen_itrf,
        telescope=telescope,
        cos_term=cos_term,
        beam_itrf=beam_itrf,
        ant1=ant1,
        ant2=ant2,
    )