Parse_HIP_Catalog.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Process SIMBAD HIP results\n",
    "\n",
    "Many of these HIP stars are in a binary system, unresolved by Hipparcos but resolved by Gaia which will cause cross-matching issue. Usually by the fact that they are missing `FLUX_G` but not `FLUX_V`, `FLUX_J` because Gaia did resolve them. The underlying issue is these binary stars are closed enough to be observed by 2MASS/Hipparcos as a single source hence having V, JHK band magnitude but resolved by Gaia hence not having a G-band magnitude because such single object does not exist in Gaia.\n",
    "\n",
    "## Data Loading and cross-match HIP and Gaia\n",
    "\n",
    "Every HIP stars should have a Gaia counterpart unless they are very bright or not a stars.\n",
    "\n",
    "*Note*: These HIP ids are not strictly stars so we will delete them\n",
    "- HIP 1902 is a globular cluster NGC 104\n",
    "- HIP 59018 is an artifact\n",
    "- HIP 60450 is an artifact\n",
    "- HIP 95723 is an artifact\n",
    "- HIP 103992 is a planetary nebular NGC 7009\n",
    "- HIP 108802 is a composite object / blend\n",
    "- HIP 54948 is a cluster of star (cant be sure which one they are referring, doubt Hipparcos has resolved them)\n",
    "- HIP 24647 is HIP 24648\n",
    "- HIP 29116 is HIP 29119\n",
    "- HIP 35194 is HIP 35195\n",
    "- HIP 91906 is HIP 91924\n",
    "- HIP 98623 is HIP 98625\n",
    "- HIP 88759 is HIP 88762\n",
    "- And a few more...\n",
    "\n",
    "And this is a list of NASA Scientific Visualization Studio missing naked eyes visible stars (HR 2366 is not actually missing from HIP)\n",
    "\n",
    "HR 2950 is a binary stars resolved in Gaia, (v) means manually checked Gaia has the stars.\n",
    "\n",
    "missing_stars = HR + [4210 (v), 4375 (v), 4374 (v), 5978 (v), 5977 (v), 4729 (v), 2322 (v), 5343 (v), 2950 (v+v), 1982 (v), 5034 (v), 4619 (v), 1704 (v), 6660 (v), 6263 (v), 2341 (v), 6848 (v)]\n",
    "\n",
    "To get Gaia data, the ADQL query used is as follow:\n",
    "\n",
    "```sql\n",
    "with x as\n",
    "(\n",
    "\tSELECT G.source_id, H2.original_ext_source_id as hip, G.ra, G.dec, G.parallax, G.parallax_over_error, G.pmra, G.pmdec, G.phot_g_mean_mag, G.bp_rp, G.radial_velocity, G.radial_velocity_error, G.astrometric_params_solved, G.ruwe, G.rv_expected_sig_to_noise, G.phot_g_mean_mag - (0.01426 * POWER(G.bp_rp, 3) - 0.2156 * POWER(G.bp_rp, 2) + 0.01424 * POWER(G.bp_rp, 1) - 0.02704) as v_mag \n",
    "    FROM gaiadr3.gaia_source AS G\n",
    "\tLEFT JOIN gaiadr3.hipparcos2_best_neighbour AS H2 ON H2.source_id = G.source_id\n",
    "\tWHERE G.phot_g_mean_mag IS NOT NULL OR H2.source_id IS NOT NULL\n",
    ")\n",
    "SELECT * \n",
    "FROM x\n",
    "WHERE (bp_rp IS NOT NULL AND v_mag <= 15.5) OR (bp_rp IS NULL AND phot_g_mean_mag <= 15.5)\n",
    "```\n",
    "\n",
    "correspond to file named `1733546104637O-result.fits` used in here (You will get a different filename if you do the same query yourself)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Resolving potential binary: 100%|██████████| 4049/4049 [01:31<00:00, 44.11it/s]\n",
      "Resolving Gaia DR3 ID xmatch: 100%|██████████| 566/566 [16:29<00:00,  1.75s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Gaia DR2 239863001382455424 not found in DR3\n"
     ]
    }
   ],
   "source": [
    "import hashlib\n",
    "import pathlib\n",
    "import re\n",
    "import struct\n",
    "import time\n",
    "import warnings\n",
    "\n",
    "import astropy.units as u\n",
    "import pandas as pd\n",
    "import tqdm\n",
    "from astropy.coordinates.name_resolve import NameResolveError\n",
    "from astropy.table import Table\n",
    "from astroquery.gaia import Gaia\n",
    "\n",
    "from py.utils import custom_simbad\n",
    "\n",
    "base_path = pathlib.Path(\"simbad_query_results\")\n",
    "hip_combined_path = base_path / \"hip_combined.dat\"\n",
    "sao_combined_path = base_path / \"sao_combined.dat\"\n",
    "hd_combined_path = base_path / \"hd_combined.dat\"\n",
    "hr_combined_path = base_path / \"hr_combined.dat\"\n",
    "# check if the combined file already exists, if not then raise\n",
    "for combined_path in [\n",
    "    hip_combined_path,\n",
    "    sao_combined_path,\n",
    "    hd_combined_path,\n",
    "    hr_combined_path,\n",
    "]:\n",
    "    if not combined_path.exists():\n",
    "        raise FileNotFoundError(\n",
    "            f\"{combined_path} does not exist. Please run simbad_query_hipsaohdhr.py first.\"\n",
    "        )\n",
    "hip_combined_table = Table.read(hip_combined_path, format=\"ascii\")\n",
    "sao_combined_table = Table.read(sao_combined_path, format=\"ascii\")\n",
    "hd_combined_table = Table.read(hd_combined_path, format=\"ascii\")\n",
    "hr_combined_table = Table.read(hr_combined_path, format=\"ascii\")\n",
    "\n",
    "gaia_id_path = base_path / \"gaia_source_id_lookup.csv\"\n",
    "try:\n",
    "    gaia_id_df = pd.read_csv(gaia_id_path, dtype=\"Int64\")\n",
    "except FileNotFoundError:\n",
    "    gaia_id_df = pd.DataFrame(columns=[\"dr2_sourceid\", \"dr3_sourceid\"], dtype=\"Int64\")\n",
    "    gaia_id_df.to_csv(gaia_id_path, index=False)\n",
    "simbad_heriarchy_cache = base_path / \"heriarchy_cache\"\n",
    "simbad_heriarchy_cache.mkdir(parents=True, exist_ok=True)\n",
    "# remove the known bad HIP stars\n",
    "hip_to_remove = [\n",
    "    1902,\n",
    "    24647,\n",
    "    29116,\n",
    "    35194,\n",
    "    54948,\n",
    "    88759,\n",
    "    91906,\n",
    "    98623,\n",
    "]  # indices of rows\n",
    "hip_combined_table.remove_rows([i - 1 for i in hip_to_remove])\n",
    "\n",
    "# regex\n",
    "re_gaiadr2 = re.compile(r\"Gaia DR2\\s+(\\d+)\")\n",
    "re_gaiadr3 = re.compile(r\"Gaia DR3\\s+(\\d+)\")\n",
    "re_hip = re.compile(r\"HIP\\s+(\\d+)\")\n",
    "re_sao = re.compile(r\"SAO\\s+(\\d+)\")\n",
    "re_hd = re.compile(r\"HD\\s+(\\d+)\")\n",
    "re_hr = re.compile(r\"HR\\s+(\\d+)\")\n",
    "\n",
    "# create an empty dataframe for cross ID of integers\n",
    "cross_id_df = pd.DataFrame(\n",
    "    columns=[\"hip\", \"gaia_dr3\", \"component\", \"sao\", \"hd\", \"hr\"], dtype=\"Int64\"\n",
    ")\n",
    "cross_id_df[\"component\"] = cross_id_df[\"component\"].astype(pd.StringDtype())\n",
    "\n",
    "\n",
    "def extract_gaia_number(text, strict=False, cone_search=False):\n",
    "    \"\"\"\n",
    "    This function will extract the Gaia DR number from the text. Try Gaia DR3 first, if not found try Gaia DR2.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    text : str\n",
    "        The text to extract the Gaia DR number from\n",
    "    strict : bool\n",
    "        If True, only extract Gaia DR3 number. Otherwise, try to extract other Gaia DR number if Gaia DR3 is not found\n",
    "    cone_search : bool\n",
    "        If True, try to cone search the source_id if not found in the lookup table\n",
    "    \"\"\"\n",
    "    global gaia_id_df\n",
    "\n",
    "    if pd.isna(text) or text == \"\":  # return NaN if the text is NaN or empty\n",
    "        return pd.NA\n",
    "\n",
    "    match = re_gaiadr3.search(text)\n",
    "    if match:\n",
    "        return int(match.group(1))\n",
    "    else:\n",
    "        if strict and not cone_search:  # if strict and not cone_search, then return NaN\n",
    "            return pd.NA\n",
    "\n",
    "    # otherwise try to find Gaia DR2\n",
    "    if not cone_search and (\n",
    "        match := re_gaiadr2.search(text)\n",
    "    ):  # still need to make sure the source_id is valid\n",
    "        # attempt to lookup the source_id in the lookup table but if not then try to cone search the source_id\n",
    "        in_lookup_row = gaia_id_df[gaia_id_df[\"dr2_sourceid\"] == int(match.group(1))]\n",
    "        if len(in_lookup_row) > 0:\n",
    "            # already in the lookup table, return the corresponding dr3_sourceid\n",
    "            return int(in_lookup_row[\"dr3_sourceid\"].values[0])\n",
    "        # else:\n",
    "        #     print(f\"Gaia DR2 {match.group(1)} not found in DR3 in lookup table\")\n",
    "        if (\n",
    "            len(\n",
    "                Gaia.launch_job(\n",
    "                    f\"SELECT source_id FROM gaiadr3.gaia_source WHERE source_id={match.group(1)}\"\n",
    "                ).results\n",
    "            )\n",
    "            > 0\n",
    "        ):\n",
    "            # the case where the dr2==dr3, so no need to clone search\n",
    "            # add the source_id to the lookup table\n",
    "            gaia_id_df = pd.concat(\n",
    "                [\n",
    "                    gaia_id_df,\n",
    "                    pd.DataFrame(\n",
    "                        {\n",
    "                            \"dr2_sourceid\": [int(match.group(1))],\n",
    "                            \"dr3_sourceid\": [int(match.group(1))],\n",
    "                        }\n",
    "                    ),\n",
    "                ]\n",
    "            )\n",
    "            return int(match.group(1))\n",
    "        else:\n",
    "            # in case we are upgrading to DR4, this need to be upgraded to DR3\n",
    "            gaia_dr2 = Gaia.launch_job(\n",
    "                f\"SELECT designation, source_id, ra, dec, parallax, pmra, pmdec, phot_g_mean_mag, bp_rp, radial_velocity FROM gaiadr2.gaia_source WHERE source_id={match.group(1)}\"\n",
    "            ).results.filled(0)[0]\n",
    "\n",
    "            # remember to change the epoch when updating to DR4\n",
    "            query = rf\"\"\"SELECT TOP 10 designation, source_id, ra, dec, parallax, pmra, pmdec, ruwe, phot_g_mean_mag, bp_rp, radial_velocity,\n",
    "DISTANCE(POINT('ICRS' , gaiadr3.gaia_source.ra, gaiadr3.gaia_source.dec), \n",
    "POINT('ICRS', \n",
    "COORD1(EPOCH_PROP_POS({gaia_dr2[\"ra\"]}, {gaia_dr2[\"dec\"]}, {gaia_dr2[\"parallax\"]}, {gaia_dr2[\"pmra\"]}, {gaia_dr2[\"pmdec\"]}, {gaia_dr2[\"radial_velocity\"]}, 2015.5, 2016.0)), \n",
    "COORD2(EPOCH_PROP_POS({gaia_dr2[\"ra\"]}, {gaia_dr2[\"dec\"]}, {gaia_dr2[\"parallax\"]}, {gaia_dr2[\"pmra\"]}, {gaia_dr2[\"pmdec\"]}, {gaia_dr2[\"radial_velocity\"]}, 2015.5, 2016.0)))) AS \"target_separation (deg)\"\n",
    "FROM gaiadr3.gaia_source \n",
    "WHERE \n",
    "CONTAINS(\n",
    "\tPOINT('ICRS',gaiadr3.gaia_source.ra,gaiadr3.gaia_source.dec),\n",
    "\tCIRCLE(\n",
    "\t\t'ICRS',\n",
    "\t\tCOORD1(EPOCH_PROP_POS({gaia_dr2[\"ra\"]}, {gaia_dr2[\"dec\"]}, {gaia_dr2[\"parallax\"]}, {gaia_dr2[\"pmra\"]}, {gaia_dr2[\"pmdec\"]}, {gaia_dr2[\"radial_velocity\"]}, 2015.5, 2016.0)),\n",
    "\t\tCOORD2(EPOCH_PROP_POS({gaia_dr2[\"ra\"]}, {gaia_dr2[\"dec\"]}, {gaia_dr2[\"parallax\"]}, {gaia_dr2[\"pmra\"]}, {gaia_dr2[\"pmdec\"]}, {gaia_dr2[\"radial_velocity\"]}, 2015.5, 2016.0)),\n",
    "\t\t{(15 * u.arcsecond).to(u.deg).value})\n",
    ")=1 ORDER BY \"target_separation (deg)\"\n",
    "\"\"\"\n",
    "            tmp = Gaia.launch_job(query).results\n",
    "            # tmp = Gaia.cone_search(f\"Gaia DR2 {match.group(1)}\", radius=15 * u.arcsecond).results\n",
    "            if len(tmp) == 0:\n",
    "                print(f\"Gaia DR2 {match.group(1)} not found in DR3\")\n",
    "                return pd.NA\n",
    "            else:\n",
    "                # add the source_id to the lookup table\n",
    "                gaia_id_df = pd.concat(\n",
    "                    [\n",
    "                        gaia_id_df,\n",
    "                        pd.DataFrame(\n",
    "                            {\n",
    "                                \"dr2_sourceid\": [int(match.group(1))],\n",
    "                                \"dr3_sourceid\": [int(tmp[0][\"source_id\"])],\n",
    "                            }\n",
    "                        ),\n",
    "                    ]\n",
    "                )\n",
    "                return int(tmp[0][\"source_id\"])\n",
    "    else:\n",
    "        if cone_search and text:\n",
    "            # try to cone search in case SIMBAD missed the Gaia DR number, stricter in radius to prevent false positive\n",
    "            gaia_source = Gaia.cone_search(\n",
    "                text.split(\"|\")[1], radius=7.5 * u.arcsecond\n",
    "            ).results\n",
    "            if len(gaia_source) > 0:\n",
    "                return int(gaia_source[0][\"source_id\"])\n",
    "        return pd.NA\n",
    "\n",
    "\n",
    "def parse_result(simbad_table):\n",
    "    \"\"\"\n",
    "    This function will parse the result from the query:\n",
    "    - Remove rows that are not stars\n",
    "    - Add additional rows to the table for stars that are potentially in a binary system\n",
    "\n",
    "    To attempt to resolve potential binary system, the logic is as follows:\n",
    "    1. If the stars have V and J magnitude dimmer than 3.0, then it should have Gaia source_id\n",
    "    2. When no Gaia source_id is found, query the children of the star because it might be a binary system resolved by Gaia hence no Gaia source_id\n",
    "    3. If childen have no Gaia source_id, then assume this is not a binary system. Just for whatever reason Gaia did not observe the parent star\n",
    "    4. If the star has children with source_id, then add the children to the table and remove the parent star\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    simbad_table : astropy.table.Table\n",
    "        Table from the query result\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    astropy.table.Table\n",
    "        Table with additional rows from the query result\n",
    "    \"\"\"\n",
    "    simbad_table = simbad_table.to_pandas()\n",
    "    simbad_table[\"otype\"] = simbad_table[\"otype\"].fillna(\"\")\n",
    "    simbad_table = simbad_table[simbad_table[\"otype\"].apply(lambda x: \"err\" not in x)]\n",
    "\n",
    "    # more strict on resolving the source_id in case of binary stars\n",
    "    source_id = simbad_table[\"ids\"].apply(lambda x: extract_gaia_number(x, strict=True))\n",
    "    # these stars probably should have source_id, so we want to resolve possible binary stars\n",
    "    hip_wo_source_id_idx = (\n",
    "        pd.isna(source_id)\n",
    "        & (simbad_table[\"V\"].fillna(99.99) > 3.0)\n",
    "        & (simbad_table[\"J\"].fillna(99.99) > 2.0)\n",
    "        & ~pd.isna(simbad_table[\"ids\"])  # need to at least have some id\n",
    "    )\n",
    "    simbad_table_temp = simbad_table[hip_wo_source_id_idx]\n",
    "    result = []\n",
    "    with warnings.catch_warnings():\n",
    "        warnings.simplefilter(\"ignore\")\n",
    "        for i, irow in tqdm.tqdm(\n",
    "            simbad_table_temp.iterrows(),\n",
    "            total=len(simbad_table_temp),\n",
    "            desc=\"Resolving potential binary\",\n",
    "        ):\n",
    "            # if children have no SAO/HD/HR id, then use the parent id\n",
    "            parent_sao = re_sao.search(irow[\"ids\"])\n",
    "            parent_hd = re_hd.search(irow[\"ids\"])\n",
    "            parent_hr = re_hr.search(irow[\"ids\"])\n",
    "\n",
    "            tmp_query = custom_simbad.query_hierarchy(\n",
    "                irow[\"main_id\"],\n",
    "                hierarchy=\"children\",\n",
    "                get_query_payload=True,\n",
    "                detailed_hierarchy=False,\n",
    "            )[\"QUERY\"]\n",
    "            tmp_hash = hashlib.sha1(tmp_query.encode(\"utf-8\")).hexdigest()\n",
    "            tmp_file = simbad_heriarchy_cache / f\"{tmp_hash}.dat\"\n",
    "            # if not exist then query and save to file for caching\n",
    "            if tmp_file.exists():\n",
    "                temp = pd.read_csv(tmp_file)\n",
    "            else:\n",
    "                temp = custom_simbad.query_hierarchy(\n",
    "                    irow[\"main_id\"], hierarchy=\"children\", detailed_hierarchy=False\n",
    "                )\n",
    "                temp = temp.to_pandas()\n",
    "                temp[\"user_specified_id\"] = irow[\"user_specified_id\"].strip()\n",
    "                source_id = temp[\"ids\"].apply(\n",
    "                    lambda x: extract_gaia_number(x, cone_search=True)\n",
    "                )\n",
    "                # add letter to the end of the user_specified_id, so first row in temp will end with \"A\", second row in temp will end with \"B\", etc\n",
    "                for idx, row in temp.iterrows():\n",
    "                    temp.loc[idx, \"user_specified_id\"] = (\n",
    "                        f\"{row['user_specified_id']}{chr(65 + idx)}\"\n",
    "                    )\n",
    "                    # add user_specified_id to the table ids column\n",
    "                    temp.loc[idx, \"ids\"] = (\n",
    "                        f\"{temp.loc[idx, 'user_specified_id']}|{temp.loc[idx, 'ids']}\"\n",
    "                    )\n",
    "                    # add temp_source_id for each row to the ids column if it is not NaN\n",
    "                    if not pd.isna(source_id[idx]):\n",
    "                        temp.loc[idx, \"ids\"] = (\n",
    "                            f\"{temp.loc[idx, 'ids']}|Gaia DR3 {source_id[idx]}\"\n",
    "                        )\n",
    "                temp.to_csv(tmp_file, index=False)\n",
    "            source_id = temp[\"ids\"].apply(lambda x: extract_gaia_number(x))\n",
    "            if len(temp) <= 1:\n",
    "                # in case some stars only have one component which is itself or no result\n",
    "                continue\n",
    "            # the second row and beyond should have source_id, if not then assume it is not a binary system\n",
    "            # also if all of them has the same source_id, then assume it is not a binary system\n",
    "            if pd.isna(source_id[1:]).any() or source_id.duplicated(keep=False).any():\n",
    "                continue\n",
    "            # add the parent SAO/HD/HR id to the ids column if it is not NaN\n",
    "            temp[\"ids\"] = temp[\"ids\"].apply(\n",
    "                lambda x: f\"{x}|SAO {parent_sao.group(1)}\" if parent_sao else x,\n",
    "            )\n",
    "            temp[\"ids\"] = temp[\"ids\"].apply(\n",
    "                lambda x: f\"{x}|HD {parent_hd.group(1)}\" if parent_hd else x,\n",
    "            )\n",
    "            temp[\"ids\"] = temp[\"ids\"].apply(\n",
    "                lambda x: f\"{x}|HR {parent_hr.group(1)}\" if parent_hr else x,\n",
    "            )\n",
    "            result.append(temp)\n",
    "            simbad_table.drop(i, inplace=True)\n",
    "    if len(result) > 0:\n",
    "        # append list of result to simbad_table\n",
    "        simbad_table = pd.concat([simbad_table, *result])\n",
    "\n",
    "    # drop empty rows\n",
    "    simbad_table = simbad_table.dropna(how=\"all\")\n",
    "    # reset the index\n",
    "    simbad_table.reset_index(drop=True, inplace=True)\n",
    "    # try to resolve the source_id for the remaining stars\n",
    "    source_id = simbad_table[\"ids\"].apply(lambda x: extract_gaia_number(x, strict=True))\n",
    "    wo_source_id_idx = (\n",
    "        pd.isna(source_id)\n",
    "        & (simbad_table[\"V\"].fillna(99.99) > 3.0)\n",
    "        & ~pd.isna(simbad_table[\"ids\"])  # need to at least have some id\n",
    "    )\n",
    "    # apply extract_gaia_number with cone_search=True to the remaining stars\n",
    "    temp = simbad_table[wo_source_id_idx]\n",
    "    for idx, row in tqdm.tqdm(\n",
    "        temp.iterrows(), total=len(temp), desc=\"Resolving Gaia DR3 ID xmatch\"\n",
    "    ):\n",
    "        time.sleep(1.0)  # sleep for 1.0 seconds to prevent rate limit\n",
    "        try:\n",
    "            simbad_table.loc[idx, \"ids\"] = (\n",
    "                f\"{row['ids']}|Gaia DR3 {extract_gaia_number(row['ids'], cone_search=True)}\"\n",
    "            )\n",
    "        except NameResolveError as e:\n",
    "            print(f\"Failed to resolve {row['main_id']} due to {e}\")\n",
    "            continue\n",
    "    # simbad_table.loc[wo_source_id_idx, \"ids\"] = simbad_table.loc[wo_source_id_idx, \"ids\"].apply(\n",
    "    #     lambda x: f\"{x}|Gaia DR3 {extract_gaia_number(x, cone_search=True)}\"\n",
    "    # )\n",
    "    return simbad_table\n",
    "\n",
    "\n",
    "def parse_cross_id(result):\n",
    "    \"\"\"\n",
    "    Parse the cross id from the query result\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    result : pandas.DataFrame\n",
    "        The result from the query\n",
    "    \"\"\"\n",
    "\n",
    "    def extract_id(text, sid, group_id=1):\n",
    "        if not pd.isna(text):  # check if not NaN\n",
    "            if isinstance(sid, str):\n",
    "                match = re.search(rf\"{sid}\\s+(\\d+)([A-Z]?)\", text)\n",
    "            elif isinstance(sid, re.Pattern):\n",
    "                match = sid.search(text)\n",
    "            else:\n",
    "                raise ValueError(\"sid must be either str or re.Pattern\")\n",
    "            if match and match.group(group_id) != \"\":\n",
    "                return match.group(group_id).strip()\n",
    "        return pd.NA\n",
    "\n",
    "    def parse_component(text):\n",
    "        if pd.isna(text) or text == \"\" or text == 0:\n",
    "            return pd.NA\n",
    "        else:\n",
    "            return text\n",
    "\n",
    "    df = pd.DataFrame(\n",
    "        columns=[\"hip\", \"gaia_dr3\", \"component\", \"sao\", \"hd\", \"hr\"], dtype=\"Int64\"\n",
    "    )\n",
    "    df[\"component\"] = df[\"component\"].astype(pd.StringDtype())\n",
    "    hip, sao, hd, hr = [\n",
    "        result[\"ids\"].apply(lambda x: extract_id(x, sid)).astype(\"Int64\")\n",
    "        for sid in [re_hip, re_sao, re_hd, re_hr]\n",
    "    ]\n",
    "    gaia = result[\"ids\"].apply(lambda x: extract_gaia_number(x)).astype(\"Int64\")\n",
    "    component = (\n",
    "        result[\"ids\"]\n",
    "        .apply(lambda x: parse_component(extract_id(x, \"HIP\", group_id=2)))\n",
    "        .astype(pd.StringDtype())\n",
    "    )\n",
    "\n",
    "    # fill the dataframe\n",
    "    df[\"hip\"] = hip\n",
    "    df[\"sao\"] = sao\n",
    "    df[\"hd\"] = hd\n",
    "    df[\"hr\"] = hr\n",
    "    df[\"gaia_dr3\"] = gaia\n",
    "    df[\"component\"] = component\n",
    "\n",
    "    return df\n",
    "\n",
    "\n",
    "hip_with_binary = parse_result(hip_combined_table)\n",
    "# deal with Sirius A and Sirius B, because SIMBAD said Sirius does not have children\n",
    "hip_with_binary.loc[hip_with_binary[\"main_id\"] == \"* alf CMa\", \"ids\"] = (\n",
    "    \"HIP 32349A|\"\n",
    "    + hip_with_binary.loc[hip_with_binary[\"main_id\"] == \"* alf CMa\", \"ids\"]\n",
    ")\n",
    "\n",
    "result = custom_simbad.query_object(\"Sirius B\")\n",
    "result[\"ids\"] = \"HIP 32349B|\" + result[\"ids\"]\n",
    "# cast to pandas and back to prevent potential dtype incompatibility issue\n",
    "simbad_table = pd.concat([hip_with_binary, result.to_pandas()])\n",
    "simbad_table = simbad_table.dropna(subset=[\"ids\"])\n",
    "\n",
    "# copy simbad_table to a new variable to prevent modifying the original table\n",
    "simbad_table_db = simbad_table.copy()\n",
    "cross_id_df = pd.concat([cross_id_df, parse_cross_id(simbad_table)])\n",
    "simbad_table_db[\"hip\"] = cross_id_df[\"hip\"].fillna(0)\n",
    "simbad_table_db[\"source_id\"] = cross_id_df[\"gaia_dr3\"].fillna(0)\n",
    "simbad_table_db[\"componentid\"] = cross_id_df[\"component\"]\n",
    "# replace the NaN with 0, A to 1, B to 2, etc\n",
    "simbad_table_db[\"componentid\"] = simbad_table_db[\"componentid\"].apply(\n",
    "    lambda x: 0 if pd.isna(x) else ord(x) - 64\n",
    ")\n",
    "simbad_table_db.to_csv(base_path / \"hip_processed_with_binary.dat\", index=False)\n",
    "# drop rows with NaN \"ids\" column\n",
    "cross_id_df = cross_id_df.dropna(how=\"all\")\n",
    "\n",
    "# find which SAO id is missing, should be SAO 1 to SAO max_sao_id. Only query the missing SAO id\n",
    "existing_sao_ids = set(cross_id_df[\"sao\"].dropna().astype(int))\n",
    "all_sao_ids = set(range(1, len(sao_combined_table) + 1))\n",
    "missing_sao_idx = list(all_sao_ids - existing_sao_ids)\n",
    "# get the missing SAO id with corresponding row\n",
    "missing_sao = sao_combined_table[[i - 1 for i in missing_sao_idx]]\n",
    "sao_cross_id_df = parse_cross_id(missing_sao.to_pandas())\n",
    "cross_id_df = pd.concat([cross_id_df, sao_cross_id_df])\n",
    "\n",
    "\n",
    "# find which HD id is missing, should be HD 1 to HD max_hd_id. Only query the missing HD id\n",
    "existing_hd_ids = set(cross_id_df[\"hd\"].dropna().astype(int))\n",
    "all_hd_ids = set(range(1, len(hd_combined_table) + 1))\n",
    "missing_hd_idx = list(all_hd_ids - existing_hd_ids)\n",
    "# get the missing HD id with corresponding row\n",
    "missing_hd = hd_combined_table[[i - 1 for i in missing_hd_idx]]\n",
    "cross_id_df = pd.concat([cross_id_df, parse_cross_id(missing_hd.to_pandas())])\n",
    "# save the lookup table\n",
    "gaia_id_df.to_csv(gaia_id_path, index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# drop rows with NaN \"ids\" column\n",
    "cross_id_df = cross_id_df.dropna(how=\"all\")\n",
    "cross_id_df.to_csv(base_path / \"cross_id_with_gaiaid.dat\", index=False, sep=\"\\t\")\n",
    "# if hip is missing, then set hip to the corresponding gaia id and then delete the gaia column\n",
    "cross_id_df.loc[pd.isna(cross_id_df[\"hip\"]), \"hip\"] = cross_id_df.loc[\n",
    "    pd.isna(cross_id_df[\"hip\"]), \"gaia_dr3\"\n",
    "].values\n",
    "cross_id_df = cross_id_df.drop(columns=[\"gaia_dr3\"])\n",
    "# if all sao, hd, hr is missing, then delete those rows. No point of keeping them\n",
    "cross_id_df = cross_id_df[\n",
    "    (\n",
    "        ~pd.isna(cross_id_df[\"sao\"])\n",
    "        | ~pd.isna(cross_id_df[\"hd\"])\n",
    "        | ~pd.isna(cross_id_df[\"hr\"])\n",
    "    )\n",
    "    & ~pd.isna(cross_id_df[\"hip\"])\n",
    "]\n",
    "\n",
    "cross_id_df = cross_id_df.sort_values([\"hip\", \"component\"]).astype(\n",
    "    {\"component\": \"object\"}\n",
    ")\n",
    "\n",
    "# cross_id_df.to_csv(\n",
    "#     base_path / \"cross-id.dat\",\n",
    "#     index=False,\n",
    "#     sep=\"\\t\",\n",
    "#     header=[\"#\" + col for col in cross_id_df.columns],\n",
    "# )\n",
    "\n",
    "# change component A, B, C... to 1, 2, 3...\n",
    "cross_id_df[\"component\"] = cross_id_df[\"component\"].apply(\n",
    "    lambda x: ord(x) - ord(\"A\") + 1 if x == x else 0\n",
    ")\n",
    "cross_id_df.fillna(0, inplace=True)\n",
    "cross_id_df = cross_id_df.astype(\"Int64\")\n",
    "record_size = 21\n",
    "bdata = bytearray(record_size * len(cross_id_df))\n",
    "\n",
    "for i, row in enumerate(cross_id_df.itertuples()):\n",
    "    byte_data = struct.pack(\"=QBiii\", *row[1:])\n",
    "    bdata[i * record_size : (1 + i) * record_size] = byte_data\n",
    "\n",
    "with open(base_path / \"cross-id.cat\", \"wb\") as f:\n",
    "    f.write(bdata)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}