antenna_attributes.py

Go to the documentation of this file.

"""!
@file antenna_attributes.py
@brief `main()` shows how to use
[read_MI_antenna_geometry()](#antenna_attributes.read_MI_antenna_geometry) and
[read_MI_antenna_geometry()](#antenna_attributes.read_global_channel_mapping).
"""
 
from pathlib import Path
import polars as pl
import os
 
 
 
 
ALL_PUEO_ANTENNA_NAMES = (
    pl.Enum([
        "101", "201", "301", "401",
        "102", "202", "302", "402",
        "103", "203", "303", "403",
        "104", "204", "304", "404",
        "105", "205", "305", "405",
        "106", "206", "306", "406",
        "107", "207", "307", "407",
        "108", "208", "308", "408",
        "109", "209", "309", "409",
        "110", "210", "310", "410",
        "111", "211", "311", "411",
        "112", "212", "312", "412",
        "113", "213", "313", "413",
        "114", "214", "314", "414",
        "115", "215", "315", "415",
        "116", "216", "316", "416",
        "117", "217", "317", "417",
        "118", "218", "318", "418",
        "119", "219", "319", "419",
        "120", "220", "320", "420",
        "121", "221", "321", "421",
        "122", "222", "322", "422",
        "123", "223", "323", "423",
        "124", "224", "324", "424",
        "501", "505",
        "605", "609",
        "713",
        "813", "817", "821"
    ])
)
 
 
ANTENNA_POLARIZATION_TYPE = pl.Enum(["H", "V"])
 
 
PUEO_ANTENNA_TYPE = pl.Enum(["MI", "LF"])
 
 
NUM_PHI_SECTORS = 24
 
 
ASSUMED_PHI_SECTOR_APERTURE_WIDTH = 50
 
 
def read_global_channel_mapping() -> pl.DataFrame:
    """!
    @ingroup AAA
    @brief Generates a Polars DataFrame from
    [map.dat](https://github.com/PUEOCollaboration/pueo-data/blob/main/geometry/jun25/map.dat)
    in the pueo-data repository.
 
    *  `map.dat` contains the channel mapping between the indices `GlobalChan`, `AntNum` and `AntIdx`
        for all channels
    *   That is, the dataframe contains both polarizations (V and H) of both instruments (MI and LF).
    *  `GlobalChan` is the index used by the `pueoEvent` repository to access the waveforms,
        see its [example script](https://github.com/PUEOCollaboration/pueoEvent/tree/main/examples).
    *   The relevant line is
        ```
        chan = pueo.GeomTool.Instance().getChanIndexFromRingPhiPol(
            pueo.ring.ring_t.kTopRing, 10, pueo.pol.pol_t.kVertical
        )
        ```
    *   `chan` in the line above comes from the column `GlobalChan` in the DataFrame.
    * Sample output:
```
shape: (208, 10)
┌────────────┬────────┬────────┬──────┬─────────┬─────────┬─────────────┬───────────┬──────┬─────────┐
│ GlobalChan ┆ AntNum ┆ AntIdx ┆ Pol  ┆ AntType ┆ SurfNum ┆ SurfChanNum ┆ PhiSector ┆ Ring ┆ Comment │
│ ---        ┆ ---    ┆ ---    ┆ ---  ┆ ---     ┆ ---     ┆ ---         ┆ ---       ┆ ---  ┆ ---     │
│ u32        ┆ enum   ┆ u32    ┆ enum ┆ enum    ┆ u8      ┆ u8          ┆ u8        ┆ u8   ┆ str     │
╞════════════╪════════╪════════╪══════╪═════════╪═════════╪═════════════╪═══════════╪══════╪═════════╡
│ 0          ┆ 101    ┆ 0      ┆ H    ┆ MI      ┆ 0       ┆ 0           ┆ 1         ┆ 1    ┆ null    │
│ 1          ┆ 201    ┆ 1      ┆ H    ┆ MI      ┆ 0       ┆ 1           ┆ 1         ┆ 2    ┆ null    │
│ 2          ┆ 301    ┆ 2      ┆ H    ┆ MI      ┆ 0       ┆ 2           ┆ 1         ┆ 3    ┆ null    │
│ 3          ┆ 401    ┆ 3      ┆ H    ┆ MI      ┆ 0       ┆ 3           ┆ 1         ┆ 4    ┆ null    │
│ 4          ┆ 102    ┆ 4      ┆ H    ┆ MI      ┆ 0       ┆ 4           ┆ 2         ┆ 1    ┆ null    │
│ …          ┆ …      ┆ …      ┆ …    ┆ …       ┆ …       ┆ …           ┆ …         ┆ …    ┆ …       │
│ 203        ┆ 609    ┆ 99     ┆ V    ┆ LF      ┆ 25      ┆ 3           ┆ 9         ┆ 6    ┆ null    │
│ 204        ┆ 713    ┆ 100    ┆ V    ┆ LF      ┆ 25      ┆ 4           ┆ 13        ┆ 7    ┆ null    │
│ 205        ┆ 813    ┆ 101    ┆ V    ┆ LF      ┆ 25      ┆ 5           ┆ 13        ┆ 8    ┆ null    │
│ 206        ┆ 817    ┆ 102    ┆ V    ┆ LF      ┆ 25      ┆ 6           ┆ 17        ┆ 8    ┆ null    │
│ 207        ┆ 821    ┆ 103    ┆ V    ┆ LF      ┆ 25      ┆ 7           ┆ 21        ┆ 8    ┆ null    │
└────────────┴────────┴────────┴──────┴─────────┴─────────┴─────────────┴───────────┴──────┴─────────┘
```
    """
 
    map_dot_dat_file = os.environ.get("PUEO_UTIL_INSTALL_DIR") / Path("share/pueo/geometry/jun25/map.dat")
    all_channels: pl.DataFrame = (
        pl.read_csv(
            map_dot_dat_file,
            comment_prefix="#",
            schema_overrides={
                "GlobalChan": pl.UInt32,
                "AntIdx": pl.UInt32,
                "Pol": ANTENNA_POLARIZATION_TYPE,
                "AntType": PUEO_ANTENNA_TYPE,
                "SurfNum": pl.UInt8,
                "SurfChanNum": pl.UInt8,
                "PhiSector": pl.UInt8,
                "Ring": pl.UInt8
            }
        )
        .with_columns(
            pl.col("AntNum").cast(pl.String).cast(ALL_PUEO_ANTENNA_NAMES)
        )
    )
 
    return all_channels
 
 
def read_antenna_geometry(qrh_dot_dat: Path, offset_inSIM:Path | None = None, ant_type='MI') -> pl.DataFrame:
    """!
    @ingroup AAA
    @brief Reads in antenna positions (Cartesian coordinates, meters) and attitudes (degrees) from
    [qrh.dat](https://github.com/PUEOCollaboration/pueo-data/blob/main/geometry/jun25/qrh.dat)
    in the pueo-data repository.
    @param[in] qrh_dot_dat Path to `qrh.dat`
 
    *   Note that `qrh.dat` contains the **face centers** of all main-instrument (MI) antennas
        (quad-ridged horn antennas) and NOT the **phase center**.
 
    *   The function will error out if, although unlikely, `qrh.dat` contains an
        index that is not listed in #ALL_PUEO_ANTENNA_NAMES.
 
    * Sample output:
    ```
shape: (96, 10)
┌────────┬──────┬───────────┬────────┬───────┬────────┬───────┬────────────┬──────────┬───────────┐
│ AntNum ┆ Ring ┆ PhiSector ┆ AntIdx ┆ X[m]  ┆ Y[m]   ┆ Z[m]  ┆ Pitch[deg] ┆ Yaw[deg] ┆ Roll[deg] │
│ ---    ┆ ---  ┆ ---       ┆ ---    ┆ ---   ┆ ---    ┆ ---   ┆ ---        ┆ ---      ┆ ---       │
│ enum   ┆ u8   ┆ u8        ┆ u32    ┆ f64   ┆ f64    ┆ f64   ┆ f64        ┆ f64      ┆ f64       │
╞════════╪══════╪═══════════╪════════╪═══════╪════════╪═══════╪════════════╪══════════╪═══════════╡
│ 101    ┆ 1    ┆ 1         ┆ 0      ┆ 1.488 ┆ 0.0    ┆ 5.114 ┆ -10.0      ┆ 0.0      ┆ 0.0       │
│ 201    ┆ 2    ┆ 1         ┆ 1      ┆ 2.596 ┆ 0.0    ┆ 1.457 ┆ -10.0      ┆ 0.0      ┆ 0.0       │
│ 301    ┆ 3    ┆ 1         ┆ 2      ┆ 2.596 ┆ 0.0    ┆ 0.728 ┆ -10.0      ┆ 0.0      ┆ 0.0       │
│ 401    ┆ 4    ┆ 1         ┆ 3      ┆ 2.596 ┆ 0.0    ┆ 0.0   ┆ -10.0      ┆ 0.0      ┆ 0.0       │
│ 102    ┆ 1    ┆ 2         ┆ 4      ┆ 1.33  ┆ 0.356  ┆ 4.476 ┆ -10.0      ┆ 15.0     ┆ 0.0       │
│ …      ┆ …    ┆ …         ┆ …      ┆ …     ┆ …      ┆ …     ┆ …          ┆ …        ┆ …         │
│ 423    ┆ 4    ┆ 23        ┆ 91     ┆ 2.248 ┆ -1.298 ┆ 0.0   ┆ -10.0      ┆ -30.0    ┆ 0.0       │
│ 124    ┆ 1    ┆ 24        ┆ 92     ┆ 1.33  ┆ -0.357 ┆ 4.476 ┆ -10.0      ┆ -15.0    ┆ 0.0       │
│ 224    ┆ 2    ┆ 24        ┆ 93     ┆ 2.507 ┆ -0.672 ┆ 1.457 ┆ -10.0      ┆ -15.0    ┆ 0.0       │
│ 324    ┆ 3    ┆ 24        ┆ 94     ┆ 2.507 ┆ -0.672 ┆ 0.728 ┆ -10.0      ┆ -15.0    ┆ 0.0       │
│ 424    ┆ 4    ┆ 24        ┆ 95     ┆ 2.507 ┆ -0.672 ┆ 0.0   ┆ -10.0      ┆ -15.0    ┆ 0.0       │
└────────┴──────┴───────────┴────────┴───────┴────────┴───────┴────────────┴──────────┴───────────┘
    ```
    """
 
    # Checking Polars version.
    version = [int(i) for i in pl.__version__.split(".")]
    major = version[0]
    minor = version[1]
    if major == 1:
        assert minor > 26 and minor != 31, "Bad polars version"
    # Need 1.25.2+ for read_csv to enforce pl.Enum checks
    # see https://github.com/pola-rs/polars/issues/20251
    # Need NOT 1.31 due to a known bug related to `pl.concat_arr()`,
    # see https://stackoverflow.com/a/79682731/21955752
    
    if ant_type in ('MI', 'LF'):
        antenna_phi_sector_mapping = (
            read_global_channel_mapping()
            .filter(pl.col("AntType") == ant_type, pl.col("Pol") == "V")
            .select("AntNum", "AntIdx", "PhiSector", "Ring")
        )
    else: 
        print("Ant type is not supported!")
    
 
    ants: pl.DataFrame = (
        pl.read_csv(
            qrh_dot_dat,
            comment_prefix="#",
            schema_overrides={
                "AntNum": ALL_PUEO_ANTENNA_NAMES,
                "Pitch[deg]": pl.Float64,
                "Yaw[deg]": pl.Float64,
                "Roll[deg]": pl.Float64,
            }
        )
        .join(antenna_phi_sector_mapping, on="AntNum")
        .select(
            pl.col("AntNum", "Ring", "PhiSector", "AntIdx"),
            pl.col("X[m]", "Y[m]", "Z[m]", "Pitch[deg]", "Yaw[deg]", "Roll[deg]")
        )
 
    )
    if offset_inSIM:
        offset_df = pl.read_csv(
           offset_inSIM,
           separator=" ",
           has_header=False,
           skip_rows=1,
           new_columns=["Pol", "dx[m]", "dy[m]", "dz[m]"],
        )
 
        dx, dy, dz = offset_df.select(["dx[m]", "dy[m]", "dz[m]"]).row(0)
 
        # Apply offset to ants
        ants = ants.with_columns(
            (pl.col("X[m]") + dx).alias("X[m]"),
            (pl.col("Y[m]") + dy).alias("Y[m]"),
            (pl.col("Z[m]") + dz).alias("Z[m]"),
        )
 
    return ants.sort("AntNum")
 
 
def get_nominal_phase_center(face_centers: pl.DataFrame, coordinates: str = None, ant_type='MI') -> pl.DataFrame:
    """!Takes the antenna **face centers** and returns the (nominal) antenna **phase centers**
    @ingroup AAA
    @param[in] face_centers  The output of #read_MI_antenna_geometry
    @param[in] coordinates   can be "cylindrical" or default (Cartesian)
    @retval    phase_centers Same as `face_centers` except `X[m]`, `Y[m]`, and `Z[m]` now represent
                             phase centers.
 
* Essential columns: `face_centers` **must** contain `X[m]` and `Y[m]`
* The nominal phase center is obtained by taking the face center and moving radially inward by 30 cm
  (cylindrical coordinates).
* If `coordinates` is not specified then Cartesian coordinates is used, and the output schema is
```
$ AntNum     <enum>
$ Ring         <u8>
$ PhiSector    <u8>
$ AntIdx      <u32>
$ X[m]        <f64>
$ Y[m]        <f64>
$ Z[m]        <f64>
$ Pitch[deg]  <f64>
$ Yaw[deg]    <f64>
$ Roll[deg]   <f64>
```
* If `coordinates` is `cylindrical`, then the ouput schema is
```
$ AntNum     <enum>
$ Ring         <u8>
$ PhiSector    <u8>
$ AntIdx      <u32>
$ Z[m]        <f64>
$ Pitch[deg]  <f64>
$ Yaw[deg]    <f64>
$ Roll[deg]   <f64>
$ Rho[m]      <f64>
$ Phi[rad]    <f64>
```
    """
 
    x = pl.col("X[m]")
    y = pl.col("Y[m]")
    
    if ant_type == 'MI':
        rho = (x ** 2 + y ** 2).sqrt() 
    else:
        rho = (x ** 2 + y ** 2).sqrt()
 
    phi = pl.arctan2(y, x)
 
    new_x = rho * phi.cos()
    new_y = rho * phi.sin()
 
    if coordinates is None:  # Cartesian
        phase_centers = face_centers.with_columns(new_x.alias("X[m]"), new_y.alias("Y[m]"))
 
    elif coordinates == "cylindrical":
        phase_centers = face_centers.select(
            pl.all().exclude("X[m]", "Y[m]"),
            rho.alias("Rho[m]"),
            phi.alias("Phi[rad]")
        )
    else:
        print(f"I can't handle this coordinate system: {coordinates}")
        exit(1)
 
    return phase_centers
 
 
if __name__ == "__main__":
    pl.Config().set_tbl_cols(10)
 
    print(read_global_channel_mapping())
 
    #ants = read_antenna_geometry(
    #    os.environ.get("PUEO_UTIL_INSTALL_DIR") / Path("share/pueo/geometry/jun25/qrh.dat")
    #)
 
    ants = read_antenna_geometry(
        os.environ.get("PUEO_UTIL_INSTALL_DIR") / Path("share/pueo/geometry/jun25/qrh.dat"),
        'MI'
    )
    print(ants)
 
    print(get_nominal_phase_center(ants, coordinates='cylindrical', ant_type='MI'))