Star_Catalog_lv0_6.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Generate catalog of stars using Hipparcos-2 and Gaia data for stars brighter than 15 mag in $V$-band\n",
    "\n",
    "Everything that is not J2016 are propagated forward to J2016.\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",
    "\n",
    "`B-V` color are computed from synthetic photometry with Gaia low-res spectra while $V$ simply estimated with Gaia $G$ and $G_\\mathrm{bp} - G_\\mathrm{rp}$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Parse the result from SIMBAD"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "All stars have astrometry\n"
     ]
    }
   ],
   "source": [
    "import pathlib\n",
    "import struct\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import trimesh\n",
    "from astropy import units as u\n",
    "from astropy.coordinates import SkyCoord\n",
    "from astropy.table import Table\n",
    "from astropy.time import Time\n",
    "from astroquery.gaia import Gaia\n",
    "\n",
    "from py.gaia import apply_space_motion, gbprp_to_bv\n",
    "from py.geodesic import GeodesicGrid, radec2xyz\n",
    "\n",
    "# paths to the data\n",
    "simbad_base_path = pathlib.Path(\"./simbad_query_results\")\n",
    "gaia_base_path = pathlib.Path(\"./gaia_query_results\")\n",
    "starcatalog_base_path = pathlib.Path(\"./star_catalogs\")\n",
    "synth_phot_path = gaia_base_path / \"Gaia_XP_JKC.csv\"\n",
    "gaia_data_path = gaia_base_path / \"1733546104637O-result.fits\"\n",
    "\n",
    "# Read the data\n",
    "simbad_t = Table.read(\n",
    "    simbad_base_path / \"hip_processed_with_binary.dat\", format=\"ascii\"\n",
    ")\n",
    "simbad_df = simbad_t.to_pandas()\n",
    "# insert default astrometry epoch to J2000 beacuse that is what SIMBAD uses\n",
    "# if has gaia source_id, we will use Gaia DR3 astrometry instead\n",
    "simbad_df[\"epoch\"] = 2000.0\n",
    "\n",
    "# read synthetic photometry\n",
    "synth_phot_df = pd.read_csv(gaia_base_path / \"Gaia_XP_JKC.csv\")\n",
    "\n",
    "# SIMBAD matching is conservative, we can use Gaia best neighbor to match the rest\n",
    "gaia_t = Table.read(gaia_data_path, format=\"fits\")\n",
    "gaia_t = gaia_t.to_pandas()\n",
    "\n",
    "# if there are source_id which most HIP stars do, use astrometry from Gaia DR3 even if RA/DEC is not NaN\n",
    "# because Gaia DR3 is much better\n",
    "matched_source_id, idx1, idx2 = np.intersect1d(\n",
    "    gaia_t[\"source_id\"].values, simbad_df[\"source_id\"].values, return_indices=True\n",
    ")\n",
    "# put the astrometry from Gaia DR3 to HIP table. Every stars in Gaia have RA and DEC, so we can safely put it in\n",
    "simbad_df.loc[idx2, \"RA_d2000\"] = gaia_t.loc[idx1, \"ra\"].values\n",
    "simbad_df.loc[idx2, \"DEC_d2000\"] = gaia_t.loc[idx1, \"dec\"].values\n",
    "# not every stars in Gaia have other astrometry, especially those stars dim enough to be\n",
    "# observed but bright enough to not have astrometry besides RA and DEC\n",
    "valid_parallax_idx = ~np.isnan(gaia_t.loc[idx1, \"parallax\"].values)\n",
    "simbad_df.loc[idx2[valid_parallax_idx], \"PLX_VALUE\"] = gaia_t.loc[\n",
    "    idx1[valid_parallax_idx], \"parallax\"\n",
    "].values\n",
    "simbad_df.loc[idx2[valid_parallax_idx], \"PLX_ERROR\"] = (\n",
    "    gaia_t.loc[idx1[valid_parallax_idx], \"parallax\"].values\n",
    "    / gaia_t.loc[idx1[valid_parallax_idx], \"parallax_over_error\"].values\n",
    ")\n",
    "valid_pm_idx = ~np.isnan(gaia_t.loc[idx1, \"pmra\"].values) & ~np.isnan(\n",
    "    gaia_t.loc[idx1, \"pmdec\"].values\n",
    ")\n",
    "simbad_df.loc[idx2[valid_pm_idx], \"PMRA\"] = gaia_t.loc[\n",
    "    idx1[valid_pm_idx], \"pmra\"\n",
    "].values\n",
    "simbad_df.loc[idx2[valid_pm_idx], \"PMDEC\"] = gaia_t.loc[\n",
    "    idx1[valid_pm_idx], \"pmdec\"\n",
    "].values\n",
    "simbad_df.loc[idx2, \"epoch\"] = (\n",
    "    2016.0  # gaia DR3 is at J2016.0 (because RA/DEC is replaced, even if not PMRA/PMDEC)\n",
    ")\n",
    "\n",
    "# Gaia DR3 5236791746461774080 is HIP 55987, but SIMBAD has no source_id for it (also 15ish mag, so outside of this Gaia Catalog range)\n",
    "simbad_df.loc[simbad_df[\"hip\"] == 55987, \"source_id\"] = 5236791746461774080\n",
    "simbad_df.loc[simbad_df[\"hip\"] == 55987, \"FLUX_G\"] = 15.615307\n",
    "simbad_df.loc[simbad_df[\"hip\"] == 55987, \"FLUX_V\"] = 15.86156559\n",
    "simbad_df.loc[simbad_df[\"hip\"] == 55987, \"RA_d2000\"] = 172.11765151979628\n",
    "simbad_df.loc[simbad_df[\"hip\"] == 55987, \"DEC_d2000\"] = -66.48608788896621\n",
    "simbad_df.loc[simbad_df[\"hip\"] == 55987, \"PMRA\"] = -8.85141903479891\n",
    "simbad_df.loc[simbad_df[\"hip\"] == 55987, \"PMDEC\"] = 0.3523543478864454\n",
    "simbad_df.loc[simbad_df[\"hip\"] == 55987, \"PLX_VALUE\"] = 0.6125993002975423\n",
    "simbad_df.loc[simbad_df[\"hip\"] == 55987, \"PLX_ERROR\"] = 0.0275\n",
    "simbad_df.loc[simbad_df[\"hip\"] == 55987, \"RV_VALUE\"] = 0.0\n",
    "simbad_df.loc[simbad_df[\"hip\"] == 55987, \"RVZ_ERROR\"] = 0.0\n",
    "simbad_df.loc[simbad_df[\"hip\"] == 55987, \"epoch\"] = 2016.0\n",
    "# Gaia DR3 1810448909928688768 is HIP 101769, but Gaia HIP best neighbor is 1810448909928688640 which is wrong. Correct it here\n",
    "simbad_df.loc[simbad_df[\"hip\"] == 101769, \"source_id\"] = 1810448909928688768\n",
    "simbad_df.loc[simbad_df[\"hip\"] == 101769, \"FLUX_G\"] = 3.9991\n",
    "# Gaia DR3 2968097043228517120 is HIP 25606, but Gaia HIP best neighbor is 2968097043227107840 which is wrong. Correct it here\n",
    "simbad_df.loc[simbad_df[\"hip\"] == 101769, \"source_id\"] = 2968097043228517120\n",
    "simbad_df.loc[simbad_df[\"hip\"] == 101769, \"FLUX_G\"] = 2.627176\n",
    "# HIP 81693 components are resolved but missing parallax and radial velocity, we will put them manually\n",
    "simbad_df.loc[simbad_df[\"hip\"] == 81693, \"PMRA\"] = -461.52\n",
    "simbad_df.loc[simbad_df[\"hip\"] == 81693, \"PMDEC\"] = 342.28\n",
    "simbad_df.loc[simbad_df[\"hip\"] == 81693, \"PLX_VALUE\"] = 93.32\n",
    "simbad_df.loc[simbad_df[\"hip\"] == 81693, \"PLX_ERROR\"] = 0.47\n",
    "simbad_df.loc[simbad_df[\"hip\"] == 81693, \"RV_VALUE\"] = -67.80\n",
    "simbad_df.loc[simbad_df[\"hip\"] == 81693, \"RVZ_ERROR\"] = 0.20\n",
    "simbad_df.loc[simbad_df[\"hip\"] == 81693, \"epoch\"] = 2016.0\n",
    "\n",
    "# =============================================================================\n",
    "simbad_df[\"B_V\"] = (simbad_df[\"FLUX_B\"] - simbad_df[\"FLUX_V\"]).astype(np.float64).values\n",
    "flux_v = simbad_df[\"FLUX_V\"].values\n",
    "gaia_v_mag = np.zeros(len(simbad_df)) * np.nan\n",
    "gaia_g_mag = np.zeros(len(simbad_df)) * np.nan\n",
    "gaia_v_mag[idx2] = gaia_t.loc[idx1, \"v_mag\"].values\n",
    "gaia_g_mag[idx2] = gaia_t.loc[idx1, \"phot_g_mean_mag\"].values\n",
    "simbad_df.fillna({\"FLUX_V\": pd.Series(gaia_v_mag)}, inplace=True)\n",
    "simbad_df.fillna({\"FLUX_V\": pd.Series(gaia_g_mag)}, inplace=True)\n",
    "# if still missing, fill with FLUX_J because nothing I can do\n",
    "simbad_df.fillna({\"FLUX_V\": simbad_df[\"FLUX_J\"]}, inplace=True)\n",
    "# try to get synthetic photometry\n",
    "matched_source_id, idx3, idx4 = np.intersect1d(\n",
    "    simbad_df[\"source_id\"].values,\n",
    "    synth_phot_df[\"source_id\"].values,\n",
    "    return_indices=True,\n",
    ")\n",
    "# if still missing, fill with FLUX_J + 3 because nothing I can do\n",
    "simbad_df.fillna({\"FLUX_V\": simbad_df[\"FLUX_J\"] + 3}, inplace=True)\n",
    "# if still missing B_V, try inferior B_V estimated from XP spectra\n",
    "simbad_df.loc[idx3, \"B_V\"] = (\n",
    "    simbad_df.loc[idx3, \"B_V\"]\n",
    "    .fillna(synth_phot_df.loc[idx4, \"Jkc_mag_B\"] - synth_phot_df.loc[idx4, \"Jkc_mag_V\"])\n",
    "    .values\n",
    ")\n",
    "# finally if still missing, estimates from BP-RP\n",
    "_, _bv = gbprp_to_bv(\n",
    "    gaia_t.loc[idx1, \"phot_g_mean_mag\"].values,\n",
    "    gaia_t.loc[idx1, \"bp_rp\"].values,\n",
    "    red_correction=True,\n",
    ")\n",
    "simbad_df.loc[idx2, \"Gaia_B_V\"] = _bv\n",
    "simbad_df.fillna({\"B_V\": simbad_df[\"Gaia_B_V\"]}, inplace=True)\n",
    "\n",
    "# if some stars with componentid >= 2 and still missing FLUX_V, delete them\n",
    "simbad_df = simbad_df[~((simbad_df[\"componentid\"] >= 2) & simbad_df[\"FLUX_V\"].isna())]\n",
    "\n",
    "# nothing should have 0 vmag, a sign of bad data\n",
    "assert (\n",
    "    np.sum(simbad_df[\"FLUX_V\"] == 0) + np.sum(simbad_df[\"FLUX_V\"].isna()) == 0\n",
    "), \"Some FlUX_V are still 0 or NaN\"\n",
    "\n",
    "# if negative parallax, set to np.nan\n",
    "neg_parallax_idx = simbad_df[\"PLX_VALUE\"] < 0\n",
    "simbad_df.loc[neg_parallax_idx, \"PLX_VALUE\"] = np.nan\n",
    "simbad_df.loc[neg_parallax_idx, \"PLX_ERROR\"] = np.nan\n",
    "\n",
    "df = pd.DataFrame(\n",
    "    index=range(len(simbad_df)),\n",
    "    data={\n",
    "        \"hip\": simbad_df[\"hip\"].fillna(0).astype(int).values,\n",
    "        \"componentid\": simbad_df[\"componentid\"].fillna(0).astype(int).values,\n",
    "        \"source_id\": simbad_df[\"source_id\"].fillna(0).astype(np.int64).values,\n",
    "        \"ra\": simbad_df[\"RA_d2000\"].values,\n",
    "        \"dec\": simbad_df[\"DEC_d2000\"].values,\n",
    "        \"epoch\": simbad_df[\"epoch\"].values,\n",
    "        \"parallax\": simbad_df[\"PLX_VALUE\"].values,\n",
    "        \"parallax_error\": simbad_df[\"PLX_ERROR\"].values,\n",
    "        \"pmra\": simbad_df[\"PMRA\"].values,\n",
    "        \"pmdec\": simbad_df[\"PMDEC\"].values,\n",
    "        \"b_v\": simbad_df[\"B_V\"].values,\n",
    "        \"vmag\": simbad_df[\"FLUX_V\"].values,\n",
    "        \"radial_velocity\": simbad_df[\"RV_VALUE\"].values,\n",
    "        \"radial_velocity_error\": simbad_df[\"RVZ_ERROR\"].values,\n",
    "        \"sp_type\": simbad_df[\"SP_TYPE\"].values,\n",
    "        \"otype\": simbad_df[\"OTYPE_N\"].values,\n",
    "    },\n",
    ")\n",
    "assert df[\"source_id\"].dtype == np.int64\n",
    "have_pm = ~np.isnan(df[\"pmra\"]) & ~np.isnan(df[\"pmdec\"])\n",
    "idx = (df[\"epoch\"] != 2016.0) & have_pm\n",
    "# those without proper motion, we can't propagate so call it a day\n",
    "df.loc[~have_pm, \"epoch\"] = 2016.0\n",
    "\n",
    "ra, dec, pmra_cosdec, pmdec, parallax, rv = apply_space_motion(\n",
    "    df.loc[idx, \"ra\"].to_numpy(),\n",
    "    df.loc[idx, \"dec\"].to_numpy(),\n",
    "    df.loc[idx, \"pmra\"].to_numpy(),\n",
    "    df.loc[idx, \"pmdec\"].to_numpy(),\n",
    "    parallax=df.loc[idx, \"parallax\"].fillna(1).to_numpy(),  # just assume at 1 kpc\n",
    "    rv=df.loc[idx, \"radial_velocity\"]\n",
    "    .fillna(0)\n",
    "    .to_numpy(),  # just assume no radial motion\n",
    "    t1=Time(df.loc[idx, \"epoch\"], format=\"jyear\"),\n",
    "    t2=Time([2016.0] * np.sum(idx), format=\"jyear\"),\n",
    ")\n",
    "\n",
    "# only map for those epoch not equal to 2016, because numerical error even J2016 to J2016\n",
    "df.loc[idx, \"ra\"] = ra\n",
    "df.loc[idx, \"dec\"] = dec\n",
    "df.loc[idx, \"pmra\"] = pmra_cosdec\n",
    "df.loc[idx, \"pmdec\"] = pmdec\n",
    "# the effect on parallax and radial velocity within human lifetime is negligible\n",
    "\n",
    "# =============================================================================\n",
    "# Process the Gaia DR3 catalog\n",
    "# =============================================================================\n",
    "# re-read the Gaia DR3 data, in case the table has been modified\n",
    "gaia_t = Table.read(gaia_data_path, format=\"fits\")\n",
    "gaia_t = gaia_t.to_pandas()\n",
    "_, gaia_t[\"b_v\"] = gbprp_to_bv(\n",
    "    gaia_t[\"phot_g_mean_mag\"], gaia_t[\"bp_rp\"], red_correction=True\n",
    ")  # as a base column, in case no synthetic photometry\n",
    "\n",
    "# if no v_mag, fill with g_mag\n",
    "gaia_t.fillna({\"v_mag\": gaia_t[\"phot_g_mean_mag\"]}, inplace=True)\n",
    "\n",
    "# try to get synthetic photometry\n",
    "matched_source_id, idx4, idx5 = np.intersect1d(\n",
    "    gaia_t[\"source_id\"].values, synth_phot_df[\"source_id\"].values, return_indices=True\n",
    ")\n",
    "# put the astrometry from Gaia DR3 to SIMBAD, make sure indexing done correctly\n",
    "synth_b_v = (\n",
    "    synth_phot_df.loc[idx5, \"Jkc_mag_B\"].values\n",
    "    - synth_phot_df.loc[idx5, \"Jkc_mag_V\"].values\n",
    ")\n",
    "gaia_t.loc[idx4, \"b_v\"] = synth_b_v\n",
    "\n",
    "# exclude those source_id already included in the simbad_df\n",
    "gaia_t = gaia_t[~gaia_t[\"source_id\"].isin(simbad_df[\"source_id\"].values)]\n",
    "good_astrometry_idx = (gaia_t[\"astrometric_params_solved\"].values == 31) | (\n",
    "    gaia_t[\"astrometric_params_solved\"].values == 95\n",
    ") & (gaia_t[\"ruwe\"].values < 1.4)\n",
    "\n",
    "# make df_gaia\n",
    "df_gaia = pd.DataFrame(\n",
    "    index=range(len(gaia_t)),\n",
    "    data={\n",
    "        \"hip\": np.zeros(len(gaia_t), dtype=int),\n",
    "        \"componentid\": np.zeros(len(gaia_t), dtype=int),\n",
    "        \"source_id\": gaia_t[\"source_id\"].astype(np.int64).values,\n",
    "        \"ra\": gaia_t[\"ra\"].values,\n",
    "        \"dec\": gaia_t[\"dec\"].values,\n",
    "        \"epoch\": np.zeros(len(gaia_t), dtype=float)\n",
    "        + 2016.0,  # we have propagated to J2000.0\n",
    "        \"parallax\": gaia_t[\"parallax\"].values,\n",
    "        \"parallax_error\": gaia_t[\"parallax\"].values\n",
    "        / gaia_t[\"parallax_over_error\"].values,\n",
    "        \"pmra\": gaia_t[\"pmra\"].values,\n",
    "        \"pmdec\": gaia_t[\"pmdec\"].values,\n",
    "        \"b_v\": gaia_t[\"b_v\"].values,\n",
    "        \"vmag\": gaia_t[\"v_mag\"].values,\n",
    "        \"radial_velocity\": gaia_t[\"radial_velocity\"].values,\n",
    "        \"radial_velocity_error\": gaia_t[\"radial_velocity_error\"].values,\n",
    "        \"sp_type\": [\"\"] * len(gaia_t),\n",
    "        \"otype\": [\"*\"] * len(gaia_t),\n",
    "    },\n",
    ")\n",
    "# nothing should have 0 vmag, a sign of bad data\n",
    "assert np.sum(df_gaia[\"vmag\"] == 0) == 0\n",
    "assert df_gaia[\"source_id\"].dtype == np.int64\n",
    "\n",
    "# combine the two\n",
    "df = pd.concat([df, df_gaia], ignore_index=True)\n",
    "\n",
    "# get those still without RA and DEC\n",
    "source_id_ls = [i for i in df[df[\"ra\"].isnull()][\"source_id\"].values]\n",
    "if len(source_id_ls) == 0:\n",
    "    print(\"All stars have astrometry\")\n",
    "else:\n",
    "    # join as one string separated by comma\n",
    "    source_id_str = \"(\" + \",\".join([str(i) for i in source_id_ls]) + \")\"\n",
    "    print(\n",
    "        \"These source_id is going to be included, exclude these in the next star catalog\"\n",
    "    )\n",
    "    print(source_id_str)\n",
    "    job = Gaia.launch_job(f\"\"\"SELECT *\n",
    "                        FROM gaiadr3.gaia_source\n",
    "                        WHERE source_id in {source_id_str}\"\"\")\n",
    "    jobresult_t = job.results\n",
    "    jobresult_df = jobresult_t.to_pandas()\n",
    "    # if there are source_id, use astrometry from Gaia DR3 even if RA/DEC is not NaN\n",
    "    matched_source_id, idx1, idx2 = np.intersect1d(\n",
    "        jobresult_df[\"SOURCE_ID\"].values, df[\"source_id\"].values, return_indices=True\n",
    "    )\n",
    "    # put the astrometry from Gaia DR3 to SIMBAD, make sure indexing done correctly\n",
    "    assert np.all(\n",
    "        df.loc[idx2, \"source_id\"].values == jobresult_df.loc[idx1, \"SOURCE_ID\"].values\n",
    "    )\n",
    "\n",
    "    # turn to vmag and b_v\n",
    "    jobresult_t[\"vmag\"], jobresult_t[\"b_v\"] = gbprp_to_bv(\n",
    "        jobresult_t[\"phot_g_mean_mag\"], jobresult_t[\"bp_rp\"], red_correction=True\n",
    "    )\n",
    "    df.loc[idx2, \"ra\"] = jobresult_df.loc[idx1, \"ra\"].values\n",
    "    df.loc[idx2, \"dec\"] = jobresult_df.loc[idx1, \"dec\"].values\n",
    "    df.loc[idx2, \"parallax\"] = jobresult_df.loc[idx1, \"parallax\"].values\n",
    "    df.loc[idx2, \"parallax_error\"] = (\n",
    "        jobresult_df.loc[idx1, \"parallax\"].values\n",
    "        / jobresult_df.loc[idx1, \"parallax_over_error\"].values\n",
    "    )\n",
    "    df.loc[idx2, \"pmra\"] = jobresult_df.loc[idx1, \"pmra\"].values\n",
    "    df.loc[idx2, \"pmdec\"] = jobresult_df.loc[idx1, \"pmdec\"].values\n",
    "    # df.loc[idx2, \"vmag\"] = jobresult_df.loc[idx1, \"vmag\"].values\n",
    "    # df.loc[idx2, \"b_v\"] = jobresult_df.loc[idx1, \"b_v\"].values\n",
    "    df.loc[idx2, \"epoch\"] = 2016.0\n",
    "\n",
    "# RA and DEC should have no NaN\n",
    "assert df[\"ra\"].isnull().sum() == 0\n",
    "assert df[\"dec\"].isnull().sum() == 0\n",
    "# nothing should have 0 vmag, a sign of bad data\n",
    "assert np.sum(df[\"vmag\"] == 0) + np.sum(df[\"vmag\"].isna()) == 0\n",
    "\n",
    "# Gaia DR3 1926461164913660160 is Ross 248, will be close to solar neighborhood in the future hence of interest, make sure it is included in the first 3 catalogs\n",
    "df.loc[df[\"source_id\"] == 1926461164913660160, \"hip\"] = 9999999\n",
    "df.loc[df[\"source_id\"] == 1926461164913660160, \"otype\"] = \"BY*\"\n",
    "df.loc[df[\"source_id\"] == 1926461164913660160, \"sp_type\"] = \"M5.0V\"\n",
    "# some stars are missing parallax\n",
    "df.loc[df[\"hip\"] == 36850, \"parallax\"] = 64.12\n",
    "df.loc[df[\"hip\"] == 36850, \"parallax_error\"] = 3.75\n",
    "df.loc[df[\"hip\"] == 84345, \"parallax\"] = 9.07\n",
    "df.loc[df[\"hip\"] == 84345, \"parallax_error\"] = 1.32\n",
    "df.loc[df[\"hip\"] == 110900, \"parallax\"] = 14.5943\n",
    "df.loc[df[\"hip\"] == 110900, \"parallax_error\"] = 0.1562\n",
    "df.loc[df[\"hip\"] == 54844, \"parallax\"] = 18.35\n",
    "df.loc[df[\"hip\"] == 54844, \"parallax_error\"] = 0.96\n",
    "df.loc[df[\"hip\"] == 55203, \"parallax\"] = 114.4867\n",
    "df.loc[df[\"hip\"] == 55203, \"parallax_error\"] = 0.4316\n",
    "\n",
    "# some stars missing astrometry\n",
    "df.loc[df[\"source_id\"] == 756853643637996160, \"parallax\"] = 114.4867\n",
    "df.loc[df[\"source_id\"] == 756853643637996160, \"parallax_error\"] = 0.4316\n",
    "df.loc[df[\"source_id\"] == 756853643637996160, \"pmra\"] = -339.398\n",
    "df.loc[df[\"source_id\"] == 756853643637996160, \"pmdec\"] = -607.892\n",
    "df.loc[df[\"source_id\"] == 756853643637996160, \"radial_velocity\"] = -15.9\n",
    "df.loc[df[\"source_id\"] == 756853643637996160, \"radial_velocity_error\"] = 0.9\n",
    "df.loc[df[\"source_id\"] == 756853643637996160, \"sp_type\"] = \"G2V\"\n",
    "df.loc[df[\"source_id\"] == 756853643637996160, \"b_v\"] = 5.41 - 4.77\n",
    "\n",
    "df.loc[df[\"source_id\"] == 3683687763520080256, \"parallax\"] = 78.5233\n",
    "df.loc[df[\"source_id\"] == 3683687763520080256, \"parallax_error\"] = 1.3879\n",
    "df.loc[df[\"source_id\"] == 3683687763520080256, \"pmra\"] = -534.318\n",
    "df.loc[df[\"source_id\"] == 3683687763520080256, \"pmdec\"] = -64.270\n",
    "df.loc[df[\"source_id\"] == 3683687763520080256, \"radial_velocity\"] = -19.8\n",
    "df.loc[df[\"source_id\"] == 3683687763520080256, \"radial_velocity_error\"] = 0.9\n",
    "df.loc[df[\"source_id\"] == 3683687763520080256, \"sp_type\"] = \"F0mF2V\"\n",
    "df.loc[df[\"source_id\"] == 3683687763520080256, \"b_v\"] = 3.85 - 3.49\n",
    "\n",
    "# other columns should not have NaN, if yes then something is wrong\n",
    "df = df.fillna(\n",
    "    {\n",
    "        \"b_v\": 0.65,\n",
    "        \"parallax\": 0,\n",
    "        \"parallax_error\": 0,\n",
    "        \"pmra\": 0,\n",
    "        \"pmdec\": 0,\n",
    "        \"radial_velocity\": 0,\n",
    "        \"radial_velocity_error\": 0,\n",
    "        \"sp_type\": \"\",\n",
    "        \"otype\": \"*  \",  # need the extra space because it is a three letter string\n",
    "    }\n",
    ")\n",
    "assert df[\"source_id\"].dtype == np.int64\n",
    "neg_parallax_idx = df[\"parallax\"] < 0\n",
    "df.loc[neg_parallax_idx, \"parallax\"] = 0\n",
    "df.loc[neg_parallax_idx, \"parallax_error\"] = 0\n",
    "\n",
    "# should not be a problem anymore\n",
    "# # find duplicated source_id, drop the duplicates that have hip_id = 0\n",
    "# # this is caused by SIMBAD resolving the Gaia-HIP cross match but not in official Gaia release\n",
    "# df.drop(\n",
    "#     df[df.duplicated(subset=\"source_id\", keep=False)]\n",
    "#     .query(\"source_id > 0 and hip == 0\")\n",
    "#     .index,\n",
    "#     inplace=True,\n",
    "# )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "min_vmag = [-2.0, 6.0, 7.5, 9.0, 10.5, 12.0, 13.75]\n",
    "max_vmag = [6.0, 7.5, 9.0, 10.5, 12.0, 13.75, 15.5]\n",
    "major_version = [0, 0, 0, 0, 0, 0, 0]\n",
    "minor_version = [14, 14, 14, 9, 5, 5, 3]\n",
    "data_type = [0, 0, 0, 0, 1, 1, 1]\n",
    "level = [0, 1, 2, 3, 4, 5, 6]\n",
    "# J2000.0 is 2451545.0\n",
    "# J2016.0 is 2457389.0\n",
    "catalog_epoch_jd = Time(2016, format=\"jyear\").jd\n",
    "\n",
    "# =============================================================================\n",
    "# Write to file\n",
    "# =============================================================================\n",
    "\n",
    "# count all spectral type\n",
    "spectral_type_ls = []\n",
    "new_spectral_type_flag = True  # flag to indicate if we are creating a new spectral type\n",
    "spectral_type_file_name = starcatalog_base_path / \"stars_hip_sp_0v0_6.cat\"\n",
    "if spectral_type_file_name.exists():\n",
    "    new_spectral_type_flag = False\n",
    "    with open(spectral_type_file_name, \"r\") as f:\n",
    "        for line in f:\n",
    "            spectral_type_ls.append(line.strip())\n",
    "\n",
    "otype_ls = []\n",
    "# read otype file\n",
    "with open(starcatalog_base_path / \"otype.dat\", \"r\") as f:\n",
    "    for line in f:\n",
    "        otype_ls.append(line.strip())\n",
    "\n",
    "\n",
    "def encode_star_hip(hip, letter_value):\n",
    "    # Ensure hip is between 0 and 120000 (17-bit)\n",
    "    if hip == 9999999:  # some hip stars are filled with 9999999 to be incldued in the first 3 catalogs\n",
    "        hip = 0\n",
    "    if hip < 0 or hip > 2**17:\n",
    "        raise ValueError(f\"HIP must be between 0 and {2**17}.\")\n",
    "\n",
    "    # Combine the 17-bit ID and the 5-bit letter value into a 24-bit integer\n",
    "    combined_value = (hip << 5) | letter_value\n",
    "\n",
    "    # Pack the 24-bit value into 3 bytes\n",
    "    return struct.pack(\"<I\", combined_value)[:3]\n",
    "\n",
    "\n",
    "for lv, datatype, majver, minver, min_v, max_v in zip(\n",
    "        level,\n",
    "        data_type,\n",
    "        major_version,\n",
    "        minor_version,\n",
    "        min_vmag,\n",
    "        max_vmag,\n",
    "    ):\n",
    "    print(lv)\n",
    "    if lv == 0 or lv == 1:\n",
    "        df6 = df[(df[\"vmag\"] >= min_v) & (df[\"vmag\"] < max_v)].reset_index(drop=True)\n",
    "    elif lv == 2:  # level 2 is special\n",
    "        # on top of the vmag range, include everything with a hip id\n",
    "        df6 = df[\n",
    "            ((df[\"vmag\"] > min_v) & (df[\"hip\"] > 0))\n",
    "            | (df[\"vmag\"] > min_v) & (df[\"vmag\"] <= max_v)\n",
    "        ].reset_index(drop=True)\n",
    "    else:  # other level should only include stars with NaN hip id\n",
    "        df6 = df[\n",
    "            (df[\"vmag\"] > min_v) & (df[\"vmag\"] <= max_v) & (df[\"hip\"] == 0)\n",
    "        ].reset_index(drop=True)\n",
    "\n",
    "    # =============================================================================\n",
    "    # Deal with geodesic grid\n",
    "    # =============================================================================\n",
    "    grid = GeodesicGrid(level=lv)\n",
    "\n",
    "    # setup ray-tracing\n",
    "    xyz = radec2xyz(df6[\"ra\"], df6[\"dec\"]).T\n",
    "    vertices = np.vstack(grid.vertices[-1])\n",
    "    faces = np.vstack(grid.faces[-1])\n",
    "    mesh = trimesh.Trimesh(vertices=vertices, faces=faces)\n",
    "    ray_origins = np.tile([[0, 0, 0]], (len(df6), 1))\n",
    "    ray_directions = np.array(xyz)\n",
    "    intersected_faces, ray_indices = mesh.ray.intersects_id(\n",
    "        ray_origins=ray_origins,\n",
    "        ray_directions=ray_directions,\n",
    "        multiple_hits=False,\n",
    "    )\n",
    "    df6.loc[ray_indices, f\"zone{lv}\"] = intersected_faces\n",
    "    # find which number ray_indices are missing\n",
    "    missing = np.setdiff1d(np.arange(len(df6)), ray_indices)\n",
    "    # if there are missing values, we need to do a more expensive search with algebra\n",
    "    zones = grid.search_zone(xyz[missing])\n",
    "    df6.loc[missing, f\"zone{lv}\"] = zones\n",
    "\n",
    "    # =============================================================================\n",
    "    # Deal with past and future zones\n",
    "    # =============================================================================\n",
    "    if lv == 0 or lv == 1 or lv == 2:\n",
    "        c = SkyCoord(\n",
    "            ra=df6[\"ra\"].values * u.deg,\n",
    "            dec=df6[\"dec\"].values * u.deg,\n",
    "            distance=(1.0 / df6[\"parallax\"].values) * u.kpc,\n",
    "            pm_ra_cosdec=df6[\"pmra\"].values * u.mas / u.yr,\n",
    "            pm_dec=df6[\"pmdec\"].values * u.mas / u.yr,\n",
    "            radial_velocity=df6[\"radial_velocity\"].values * u.km / u.s,\n",
    "            obstime=Time(2016.0, format=\"jyear\"),\n",
    "            frame=\"icrs\",\n",
    "        )\n",
    "        g = GeodesicGrid(level=lv)\n",
    "        zone_now = g.search_zone(c.cartesian.xyz.T)\n",
    "        zone_past = g.search_zone(c.cartesian.xyz.T)\n",
    "        zone_future = g.search_zone(c.cartesian.xyz.T)\n",
    "        good_astrometry_idx = (df6[\"parallax\"] / df6[\"parallax_error\"] > 5)\n",
    "\n",
    "        # calculate absmag now\n",
    "        absmag = df6[\"vmag\"] - 5 * np.log10(1000.0 / df6[\"parallax\"]) + 5\n",
    "\n",
    "        counter_past = np.zeros(len(c))\n",
    "        counter_future = np.zeros(len(c))\n",
    "        max_vmag_diff = np.zeros(len(c))\n",
    "\n",
    "        for iyr in range(10000, 110000, 10000):  # over 10000 years to have some buffer\n",
    "            c_past = SkyCoord(\n",
    "                ra=df6[\"ra\"].values * u.deg,\n",
    "                dec=df6[\"dec\"].values * u.deg,\n",
    "                distance=(1.0 / df6[\"parallax\"].values) * u.kpc,\n",
    "                pm_ra_cosdec=df6[\"pmra\"].values * u.mas / u.yr,\n",
    "                pm_dec=df6[\"pmdec\"].values * u.mas / u.yr,\n",
    "                radial_velocity=df6[\"radial_velocity\"].values * u.km / u.s,\n",
    "                obstime=Time(iyr, format=\"jyear\"),\n",
    "                frame=\"icrs\",\n",
    "            )\n",
    "            c_past = c_past.apply_space_motion(Time(0.0, format=\"jyear\"))\n",
    "            temp_zone = g.search_zone(c_past.cartesian.xyz.T)\n",
    "            counter_past += temp_zone != zone_past\n",
    "            zone_past = temp_zone\n",
    "\n",
    "            c_future = SkyCoord(\n",
    "                ra=df6[\"ra\"].values * u.deg,\n",
    "                dec=df6[\"dec\"].values * u.deg,\n",
    "                distance=(1.0 / df6[\"parallax\"].values) * u.kpc,\n",
    "                pm_ra_cosdec=df6[\"pmra\"].values * u.mas / u.yr,\n",
    "                pm_dec=df6[\"pmdec\"].values * u.mas / u.yr,\n",
    "                radial_velocity=df6[\"radial_velocity\"].values * u.km / u.s,\n",
    "                obstime=Time(0.0, format=\"jyear\"),\n",
    "                frame=\"icrs\",\n",
    "            )\n",
    "            c_future = c_future.apply_space_motion(Time(iyr, format=\"jyear\"))\n",
    "            temp_zone = g.search_zone(c_future.cartesian.xyz.T)\n",
    "            counter_future += temp_zone != zone_future\n",
    "            zone_future = temp_zone\n",
    "\n",
    "            # calculate vmag in the future and the past\n",
    "            # need to check in a for loop because some stars get bright and dim again\n",
    "            vmag_past = (5 * np.log10(c_future.distance.value) - 5 + absmag)\n",
    "            vmag_future = (5 * np.log10(c_past.distance.value) - 5 + absmag)\n",
    "            max_vmag_diff = np.min([max_vmag_diff, vmag_past - df6[\"vmag\"], vmag_future - df6[\"vmag\"]], axis=0)\n",
    "\n",
    "        max_vmag_diff[np.isinf(max_vmag_diff) | ~good_astrometry_idx] = 0\n",
    "        if lv == 0:  # for level 0, we are more lenient because easy to zoom enough that border triangles are not checked\n",
    "            # no need to check brightness because these stars are very bright already\n",
    "            total_counter = (counter_past + counter_future > 0) & good_astrometry_idx\n",
    "        else:\n",
    "            total_counter = np.logical_or(np.logical_or(counter_past > 1, counter_future > 1), max_vmag_diff < -0.3)\n",
    "        print(np.sum(total_counter))\n",
    "        df6.loc[total_counter, f\"zone{lv}\"] = 20 * 4 ** lv\n",
    "\n",
    "    # sort by zone, within each zone sort by vmag\n",
    "    df6 = df6.sort_values([f\"zone{lv}\", \"vmag\"]).reset_index(drop=True)\n",
    "    df6[\"pmra_wo_cosdec\"] = df6[\"pmra\"] / np.cos(np.radians(df6[\"dec\"]))\n",
    "\n",
    "    f = open(\n",
    "        starcatalog_base_path / f\"./stars_{lv}_{datatype}v{majver}_{minver}.cat\", \"w+b\"\n",
    "    )\n",
    "\n",
    "    f.write(b\"\\n\\x04_\\x83\")  # Magic Number\n",
    "    f.write(np.int32(datatype).tobytes())  # Data Type\n",
    "    f.write(np.int32(majver).tobytes())  # Major Version\n",
    "    f.write(np.int32(minver).tobytes())  # Minor Version\n",
    "    f.write(np.int32(lv).tobytes())  # Level\n",
    "    f.write(np.int32(min_v * 1000).tobytes())  # Magnitude Minimum\n",
    "    f.write(np.float32(catalog_epoch_jd).tobytes())  # Catalog Epoch\n",
    "    n_zones = 20 * 4 ** lv + 1  # plus 1 global zone\n",
    "\n",
    "    # count number of stars in each zone in df6\n",
    "    zone_info = df6[f\"zone{lv}\"].value_counts().sort_index()\n",
    "    for z in range(n_zones):\n",
    "        f.write(struct.pack(\"I\", zone_info.get(z, 0)))\n",
    "\n",
    "    max_records = sum(zone_info)\n",
    "\n",
    "    df6[\"source_id\"] = df6[\"source_id\"].astype(np.int64)\n",
    "    if datatype == 0:\n",
    "        # =============================================================================\n",
    "        # astrometry\n",
    "        # =============================================================================\n",
    "        sra = np.sin(np.radians(df6[\"ra\"]))\n",
    "        cra = np.cos(np.radians(df6[\"ra\"]))\n",
    "        sdec = np.sin(np.radians(df6[\"dec\"]))\n",
    "        cdec = np.cos(np.radians(df6[\"dec\"]))\n",
    "        # normal triad\n",
    "        p = np.array([-sra, cra, np.zeros(len(df6))])\n",
    "        q = np.array([-sdec * cra, -sdec * sra, cdec])\n",
    "        r = np.array([cdec * cra, cdec * sra, sdec])\n",
    "        # proper motion\n",
    "        pmvec0 = np.atleast_2d(df6[\"pmra\"]) * p + np.atleast_2d(df6[\"pmdec\"]) * q\n",
    "\n",
    "        r *= 2e9\n",
    "        pmvec0 *= 1000\n",
    "        # round to nearest integer\n",
    "        r = np.around(r).astype(np.int32)\n",
    "        pmvec0 = np.around(pmvec0).astype(np.int32)\n",
    "        df6[\"b_v_1000\"] = (df6[\"b_v\"] * 1000.0).round().astype(int)\n",
    "        df6[\"vmag_1000\"] = (df6[\"vmag\"] * 1000.0).round().astype(int)\n",
    "        df6[\"parallax_50\"] = (df6[\"parallax\"] * 50.0).round().astype(int)\n",
    "        df6[\"parallax_error_100\"] = (df6[\"parallax_error\"] * 100.0).round().astype(int)\n",
    "        # if radial velocity is rudiculously large, set to 0 (suspicious e.g., a quasar)\n",
    "        df6[\"radial_velocity_10\"] = (\n",
    "            np.where(\n",
    "                np.abs(df6[\"radial_velocity\"]) > np.iinfo(np.int16).max / 10.0,\n",
    "                0,\n",
    "                df6[\"radial_velocity\"] * 10,\n",
    "            )\n",
    "            .round()\n",
    "            .astype(int)\n",
    "        )\n",
    "\n",
    "        # pre-allocate\n",
    "        bdata = bytearray(48 * max_records)\n",
    "\n",
    "        for i in range(max_records):\n",
    "            # star_header = (\n",
    "            #     \"gaia_id\",\n",
    "            #     \"x0\"\n",
    "            #     \"x1\"\n",
    "            #     \"x2\",\n",
    "            #     \"dx0\",\n",
    "            #     \"dx1\",\n",
    "            #     \"dx2\",\n",
    "            #     \"b_v\",\n",
    "            #     \"mag\",\n",
    "            #     \"plx\",\n",
    "            #     \"plx_err\",\n",
    "            #     \"rv\"\n",
    "            #     \"sp_int\",\n",
    "            #     \"object_type\",\n",
    "            #     \"hip + component\",\n",
    "            # )\n",
    "            _sptype_index = 0\n",
    "            _sptype = df6.loc[i, \"sp_type\"]\n",
    "            if _sptype not in spectral_type_ls and _sptype != \"\" and new_spectral_type_flag:\n",
    "                spectral_type_ls.append(_sptype)\n",
    "            try:\n",
    "                _sptype_index = spectral_type_ls.index(_sptype) + (1 * new_spectral_type_flag)\n",
    "            except ValueError:\n",
    "                pass  # it can be the case where we are not creating a new spectral type files but the spectral type is not in the list\n",
    "\n",
    "            hip_bytes = encode_star_hip(\n",
    "                np.uint32(df6.loc[i, \"hip\"]), int(df6.loc[i, \"componentid\"])\n",
    "            )\n",
    "            byte_data = (\n",
    "                struct.pack(\n",
    "                    \"qiiiiiihhHHhHB\",\n",
    "                    df6.loc[i, \"source_id\"],\n",
    "                    r[0, i],\n",
    "                    r[1, i],\n",
    "                    r[2, i],\n",
    "                    pmvec0[0, i],\n",
    "                    pmvec0[1, i],\n",
    "                    pmvec0[2, i],\n",
    "                    df6.loc[i, \"b_v_1000\"],\n",
    "                    df6.loc[i, \"vmag_1000\"],\n",
    "                    df6.loc[i, \"parallax_50\"],\n",
    "                    df6.loc[i, \"parallax_error_100\"],\n",
    "                    df6.loc[i, \"radial_velocity_10\"],\n",
    "                    _sptype_index,  # 0 is reserved for No Information\n",
    "                    otype_ls.index(df6.loc[i, \"otype\"]),\n",
    "                )\n",
    "                + hip_bytes\n",
    "            )\n",
    "            bdata[i * 48 : (1 + i) * 48] = byte_data\n",
    "    elif datatype == 1:\n",
    "        # round to nearest integer\n",
    "        df6[\"ra_3600000\"] = (df6[\"ra\"] * 3_600_000).round().astype(int)\n",
    "        df6[\"dec_3600000\"] = (df6[\"dec\"] * 3_600_000).round().astype(int)\n",
    "        df6[\"pmra_wo_cosdec_1000\"] = (df6[\"pmra_wo_cosdec\"] * 1000).round().astype(int)\n",
    "        df6[\"pmdec_1000\"] = (df6[\"pmdec\"] * 1000).round().astype(int)\n",
    "        df6[\"b_v_1000\"] = (df6[\"b_v\"] * 1000).round().astype(int)\n",
    "        df6[\"vmag_1000\"] = (df6[\"vmag\"] * 1000).round().astype(int)\n",
    "        df6[\"parallax_100\"] = (df6[\"parallax\"] * 100).round().astype(int)\n",
    "        df6[\"parallax_error_100\"] = (df6[\"parallax_error\"] * 100).round().astype(int)\n",
    "        bdata = bytearray(32 * max_records)\n",
    "\n",
    "        for i in range(max_records):\n",
    "            # star_header = (\n",
    "            #     \"gaia_id\",\n",
    "            #     \"x0\",\n",
    "            #     \"x1\",\n",
    "            #     \"dx0\",\n",
    "            #     \"dx1\",\n",
    "            #     \"b_v\",\n",
    "            #     \"mag\",\n",
    "            #     \"plx\",\n",
    "            #     \"plx_err\"\n",
    "            # )\n",
    "            byte_data = struct.pack(\n",
    "                \"qiiiihhHH\",\n",
    "                df6.loc[i, \"source_id\"],\n",
    "                df6.loc[i, \"ra_3600000\"],\n",
    "                df6.loc[i, \"dec_3600000\"],\n",
    "                df6.loc[i, \"pmra_wo_cosdec_1000\"],\n",
    "                df6.loc[i, \"pmdec_1000\"],\n",
    "                df6.loc[i, \"b_v_1000\"],\n",
    "                df6.loc[i, \"vmag_1000\"],\n",
    "                df6.loc[\n",
    "                    i, \"parallax_100\"\n",
    "                ],  # switch form 50 to 100 beacuse proxima centauri is not here\n",
    "                df6.loc[i, \"parallax_error_100\"],\n",
    "            )\n",
    "            bdata[i * 32 : (1 + i) * 32] = byte_data\n",
    "    else:\n",
    "        raise ValueError(\"Data Type must be 0 or 1\")\n",
    "    f.write(bdata)\n",
    "    f.close()\n",
    "\n",
    "# write the spectral type to a file, first line empty, then one spectral type per line\n",
    "if new_spectral_type_flag:\n",
    "    with open(spectral_type_file_name, \"wb\") as f:\n",
    "        f.write(b\"\\n\")\n",
    "        for s in spectral_type_ls:\n",
    "            f.write(f\"{s}\\n\".encode())"
   ]
  }
 ],
 "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.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}