Example of the Aurora model

docs/notebooks/Aurora/1. Download and prepare RGB stars catalogue.ipynb

@@ -0,0 +1,347 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "6e6c8d8e-b253-4307-b4fc-9a7b9b3516cd",
+   "metadata": {},
+   "source": [
+    "# 1. Download and prepare RGB stars catalogue"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "33aa7797-5f2e-4ecd-9f36-dc00e5cd9b57",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "hpx_order=7 --> (hpx_nside=128, hpx_npix=196608)\n",
+      "model_hpx_order=5 --> (model_hpx_nside=32, model_hpx_npix=12288)\n"
+     ]
+    }
+   ],
+   "source": [
+    "import pathlib\n",
+    "\n",
+    "import numpy as np\n",
+    "import pandas as pd\n",
+    "\n",
+    "import astropy\n",
+    "import astropy.units as u\n",
+    "from astropy.coordinates import SkyCoord, Galactocentric\n",
+    "\n",
+    "import config\n",
+    "import utils"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "03c9fbe6-e69d-40d2-af97-097428b0e4ca",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "cache_path = pathlib.Path(config.cache_path)\n",
+    "cache_path.mkdir(exist_ok=True)\n",
+    "\n",
+    "fig_path = pathlib.Path(config.fig_path)\n",
+    "fig_path.mkdir(exist_ok=True)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "76c418ff-cd08-4c5b-ad88-adb1281351f1",
+   "metadata": {},
+   "source": [
+    "## Download Andrae et al. (2023) vetted sample of RGB stars\n",
+    "\n",
+    "Paper: https://ui.adsabs.harvard.edu/abs/2023ApJS..267....8A\n",
+    "\n",
+    "Data: https://zenodo.org/records/7945154"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "bbae7a80-44e2-4d7e-9f39-5d381ef9844b",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Starting download from https://zenodo.org/records/7945154/files/table_2_catwise.csv.gz?download=1\n",
+      "File 'cache/Andrae2023_table_2_catwise.csv.gz' already downloaded\n",
+      "Done\n"
+     ]
+    }
+   ],
+   "source": [
+    "# Download the Andrae et al. (2023) vetted sample of RGB stars\n",
+    "\n",
+    "url = 'https://zenodo.org/records/7945154/files/table_2_catwise.csv.gz?download=1'\n",
+    "file_name = cache_path / 'Andrae2023_table_2_catwise.csv.gz'\n",
+    "utils.download_file(url, file_name)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "0f48b759-9162-409d-b0a3-39f02d922c2a",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Num. of RGB stars: 17558141\n"
+     ]
+    }
+   ],
+   "source": [
+    "# Read in the downloaded catalogue\n",
+    "# Use only necessary columns\n",
+    "usecols = ['source_id', 'ra', 'dec', 'l', 'b', 'parallax', 'parallax_error', \\\n",
+    "           'pmra', 'pmdec', 'radial_velocity', \\\n",
+    "           'phot_g_mean_mag', 'phot_bp_mean_mag', 'phot_rp_mean_mag', \\\n",
+    "           'mh_xgboost', 'teff_xgboost', 'logg_xgboost']\n",
+    "rgb = pd.read_csv(file_name, usecols=usecols)\n",
+    "print(\"Num. of RGB stars:\", len(rgb))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "fed5ce33-a1dc-4542-882d-b8c411aab266",
+   "metadata": {},
+   "source": [
+    "## Make the transformation from Galactic to Galactocentric frame"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "e3c5b2c5-9946-4935-ab97-72052fa980e5",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Use inverse parallax as a distance measure\n",
+    "rgb['dist'] = 1 / rgb['parallax']"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "abb07d6d-49ca-475a-b099-bbc66540b806",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Sun Cartesian pos: -8.1219733661223 0.0 0.020800000000000003\n",
+      "Sun Cartesian vel: 12.899999999999999 245.6 7.779999999999999\n"
+     ]
+    }
+   ],
+   "source": [
+    "# The Galactocentric frame\n",
+    "gc_frame = Galactocentric()\n",
+    "\n",
+    "# Use the default position and velocity of the Sun\n",
+    "# At the September 2024, the coordinates are taken from GRAVITY Collaboration et al. (2018) and Bennett & Bovy (2019);\n",
+    "# velocities are taken brom Drimmel & Poggio (2018), GRAVITY Collaboration et al. (2018), Reid & Brunthaler (2004)\n",
+    "c_sun_icrs = SkyCoord(0*u.deg, 0*u.deg, 0*u.kpc, pm_ra_cosdec=0*u.mas/u.yr, pm_dec=0*u.mas/u.yr, radial_velocity=0*u.km/u.s, frame='icrs')\n",
+    "c_sun_gc = c_sun_icrs.transform_to(gc_frame)\n",
+    "\n",
+    "x_sun = c_sun_gc.x.value\n",
+    "y_sun = c_sun_gc.y.value\n",
+    "z_sun = c_sun_gc.z.value\n",
+    "print(\"Sun Cartesian pos:\", x_sun, y_sun, z_sun)\n",
+    "\n",
+    "vx_sun = c_sun_gc.v_x.value\n",
+    "vy_sun = c_sun_gc.v_y.value\n",
+    "vz_sun = c_sun_gc.v_z.value\n",
+    "print(\"Sun Cartesian vel:\", vx_sun, vy_sun, vz_sun)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "id": "7a7ae180-4d78-4f06-9e9e-be7a41555276",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "{'parameters': {'galcen_coord': <ICRS Coordinate: (ra, dec) in deg\n",
+       "      (266.4051, -28.936175)>,\n",
+       "  'galcen_distance': <Quantity 8.122 kpc>,\n",
+       "  'galcen_v_sun': <CartesianDifferential (d_x, d_y, d_z) in km / s\n",
+       "      (12.9, 245.6, 7.78)>,\n",
+       "  'z_sun': <Quantity 20.8 pc>,\n",
+       "  'roll': <Quantity 0. deg>},\n",
+       " 'references': {'galcen_coord': 'https://ui.adsabs.harvard.edu/abs/2004ApJ...616..872R',\n",
+       "  'galcen_distance': 'https://ui.adsabs.harvard.edu/abs/2018A%26A...615L..15G',\n",
+       "  'galcen_v_sun': ['https://ui.adsabs.harvard.edu/abs/2018RNAAS...2..210D',\n",
+       "   'https://ui.adsabs.harvard.edu/abs/2018A%26A...615L..15G',\n",
+       "   'https://ui.adsabs.harvard.edu/abs/2004ApJ...616..872R'],\n",
+       "  'z_sun': 'https://ui.adsabs.harvard.edu/abs/2019MNRAS.482.1417B',\n",
+       "  'roll': None}}"
+      ]
+     },
+     "execution_count": 7,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "# Clickable references:\n",
+    "astropy.coordinates.galactocentric_frame_defaults.get_from_registry('latest')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "id": "823eb226-ae7e-4158-a86c-ab55cdd2f141",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Another way to define the Galactic frame when we like someone's estimate for the Sun position and velocity\n",
+    "#gc_frame = Galactocentric(galcen_distance=abs(x_sun)*u.kpc, z_sun=z_sun*u.kpc, galcen_v_sun=[vx_sun, vy_sun, vz_sun]*u.km/u.s)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "id": "93a28e9d-d9cc-4fbc-b865-30e375bf0a11",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Transfrom\n",
+    "# Actually we don't need velocities for our model but let them be\n",
+    "\n",
+    "c_icrs = SkyCoord(rgb['ra'].values * u.deg,\n",
+    "                  rgb['dec'].values * u.deg,\n",
+    "                  rgb['dist'].values * u.kpc,\n",
+    "                  pm_ra_cosdec=rgb['pmra'].values * u.mas/u.yr,\n",
+    "                  pm_dec=rgb['pmdec'].values * u.mas/u.yr,\n",
+    "                  radial_velocity=rgb['radial_velocity'].values * u.km/u.s,\n",
+    "                  frame='icrs')\n",
+    "\n",
+    "c_gc = c_icrs.transform_to(gc_frame)\n",
+    "\n",
+    "x = c_gc.x.to(u.kpc).value\n",
+    "y = c_gc.y.to(u.kpc).value\n",
+    "z = c_gc.z.to(u.kpc).value\n",
+    "\n",
+    "vx = c_gc.v_x.to(u.km/u.s).value\n",
+    "vy = c_gc.v_y.to(u.km/u.s).value\n",
+    "vz = c_gc.v_z.to(u.km/u.s).value\n",
+    "\n",
+    "rgb['x'], rgb['y'], rgb['z']    =  x,  y,  z\n",
+    "rgb['vx'], rgb['vy'], rgb['vz'] = vx, vy, vz"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "cce90191-a005-4750-a882-5736916ecde7",
+   "metadata": {},
+   "source": [
+    "## Save RGB stars catalogue"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "id": "657ded99-992b-4b17-9b77-ec1cb6d3c284",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "NaNs:\n",
+      "source_id                 0\n",
+      "l                         0\n",
+      "b                         0\n",
+      "ra                        0\n",
+      "dec                       0\n",
+      "parallax                  0\n",
+      "parallax_error            0\n",
+      "pmra                      0\n",
+      "pmdec                     0\n",
+      "radial_velocity     4982977\n",
+      "phot_g_mean_mag           0\n",
+      "phot_bp_mean_mag          0\n",
+      "phot_rp_mean_mag          0\n",
+      "mh_xgboost                0\n",
+      "teff_xgboost              0\n",
+      "logg_xgboost              0\n",
+      "dist                      0\n",
+      "x                         0\n",
+      "y                         0\n",
+      "z                         0\n",
+      "vx                  4982977\n",
+      "vy                  4982977\n",
+      "vz                  4982977\n",
+      "dtype: int64\n",
+      "Columns: Index(['source_id', 'l', 'b', 'ra', 'dec', 'parallax', 'parallax_error',\n",
+      "       'pmra', 'pmdec', 'radial_velocity', 'phot_g_mean_mag',\n",
+      "       'phot_bp_mean_mag', 'phot_rp_mean_mag', 'mh_xgboost', 'teff_xgboost',\n",
+      "       'logg_xgboost', 'dist', 'x', 'y', 'z', 'vx', 'vy', 'vz'],\n",
+      "      dtype='object')\n"
+     ]
+    }
+   ],
+   "source": [
+    "# We don't care of NaNs in velocities or extinctions\n",
+    "print(\"NaNs:\")\n",
+    "print(rgb.isna().sum())\n",
+    "\n",
+    "print(\"Columns:\", rgb.columns)\n",
+    "\n",
+    "# The HDF5 format is fast!\n",
+    "# https://pandas.pydata.org/docs/reference/api/pandas.DataFrame.to_hdf.html#pandas.DataFrame.to_hdf\n",
+    "# https://pandas.pydata.org/docs/reference/api/pandas.read_hdf.html#pandas.read_hdf\n",
+    "columns = ['source_id', 'l', 'b', 'parallax', 'parallax_error', 'dist', 'x', 'y', 'z', 'vx', 'vy', 'vz', \\\n",
+    "           'phot_g_mean_mag', 'phot_bp_mean_mag', 'phot_rp_mean_mag', \\\n",
+    "           'teff_xgboost', 'mh_xgboost']\n",
+    "data_columns = ['source_id', 'l', 'b', 'parallax', 'phot_g_mean_mag', 'phot_bp_mean_mag', 'phot_rp_mean_mag', 'teff_xgboost', 'mh_xgboost']\n",
+    "complevel = 0  # 0..9\n",
+    "rgb[columns].to_hdf(cache_path / 'rgb.hdf5', key='rgb', format='table', data_columns=data_columns, complevel=complevel, mode='w', index=False)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "618fcef7-8da6-461f-88ff-3fe1ca024e0f",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "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.11.9"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}