{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "ba4740d1",
   "metadata": {},
   "source": [
    "# Grab Gaia data and format it for Firefly"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c8cbf98d",
   "metadata": {},
   "source": [
    "Create a conda env\n",
    "\n",
    "```\n",
    "conda create --name Gaia python=3.10 astropy astroquery jupyter scipy numpy matplotlib pandas h5py eventlet flask flask-socketio requests\n",
    "```\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "id": "c124f0dd",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The autoreload extension is already loaded. To reload it, use:\n",
      "  %reload_ext autoreload\n"
     ]
    }
   ],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "from astroquery.gaia import Gaia\n",
    "from astropy.table import Table, Column\n",
    "import astropy.units as u\n",
    "from astropy.coordinates import SkyCoord, ICRS\n",
    "from astropy.io import ascii"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "id": "4ad3051f",
   "metadata": {},
   "outputs": [],
   "source": [
    "GaiaCatalog = \"gaiaedr3.gaia_source\" "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12119c2c-ef64-458b-87d1-e2aa7235c459",
   "metadata": {},
   "outputs": [],
   "source": [
    "# def queryMagnitudeBin(gmin,gmax,N):\n",
    "    columns = 'ra, dec, parallax, pmra, pmdec, dr2_radial_velocity, phot_g_mean_mag, bp_rp'\n",
    "\n",
    "    cmd = f\"SELECT TOP {int(N)} {columns} \\\n",
    "    FROM {GaiaCatalog} \\\n",
    "    WHERE phot_g_mean_mag>{gmin} \\\n",
    "    and phot_g_mean_mag<={gmax} \\\n",
    "    AND parallax IS NOT NULL \\\n",
    "    AND dr2_radial_velocity IS NOT NULL\"\n",
    "    print(cmd)\n",
    "\n",
    "    #synchronous commands are OK for jobs with < 2000 output rows\n",
    "    job = Gaia.launch_job(cmd, dump_to_file=False) \n",
    "\n",
    "    #asynchronous commands for larger files\n",
    "    job = Gaia.launch_job_async(cmd, dump_to_file=False)\n",
    "    print(job)\n",
    "    GaiaData = job.get_results()\n",
    "    return GaiaData\n",
    "\n",
    "magnitudeEdges = [0,10,11,12,13,14]\n",
    "dfs = []\n",
    "for i,gmax in enumerate(magnitudeEdges[1:]):\n",
    "    gmin = magnitudeEdges[i]\n",
    "    dfs+=[queryMagnitudeBin(gmin,gmax,1e7)]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "16e249d2-65b1-4d17-afe1-dd580954ec8a",
   "metadata": {},
   "source": [
    "## Convert these ra, dec, parallax coordinates to 3D cartesian\n",
    "\n",
    "A better way may be to use the Bailer-Jones distances..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 374,
   "id": "e124af92-aa2e-443c-bfda-2b9415a1b568",
   "metadata": {},
   "outputs": [],
   "source": [
    "def saveToDisk(GaiaData,gmin,gmax):\n",
    "    mask = ( (GaiaData['parallax'] > 0) &\n",
    "         ~np.isnan((GaiaData['parallax'])) &\n",
    "         ~np.isnan((GaiaData['phot_g_mean_mag'])) &\n",
    "         ~np.isnan((GaiaData['bp_rp'])) &\n",
    "         ~np.isnan((GaiaData['dr2_radial_velocity'])) \n",
    "       )\n",
    "\n",
    "    useGaiaData = GaiaData[mask]\n",
    "    print(np.sum(mask),'many particles')\n",
    "    \n",
    "    coordsRV = ICRS(ra = useGaiaData['ra'], \n",
    "              dec = useGaiaData['dec'], \n",
    "              distance = (useGaiaData['parallax']).to(u.parsec, equivalencies=u.parallax()),\n",
    "              pm_ra_cosdec = useGaiaData['pmra'],\n",
    "              pm_dec = useGaiaData['pmdec'], \n",
    "              radial_velocity = useGaiaData['dr2_radial_velocity'])\n",
    "\n",
    "    outdfRV = pd.DataFrame()\n",
    "    outdfRV['x'] = coordsRV.cartesian.x.value\n",
    "    outdfRV['y'] = coordsRV.cartesian.y.value\n",
    "    outdfRV['z'] = coordsRV.cartesian.z.value\n",
    "    outdfRV['vx'] = coordsRV.velocity.d_x.value\n",
    "    outdfRV['vy'] = coordsRV.velocity.d_y.value\n",
    "    outdfRV['vz'] = coordsRV.velocity.d_z.value\n",
    "    outdfRV['bp_rp'] = np.array(useGaiaData['bp_rp'].data)\n",
    "    outdfRV['phot_g_mean_mag'] = np.array(useGaiaData['phot_g_mean_mag'].data)\n",
    "\n",
    "    # export this to a csv file for now\n",
    "    outdfRV.to_csv('GaiaedR3RV%d-%d.csv'%(gmin,gmax))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 375,
   "id": "28cbfe1e-bbe0-42ed-b34b-e88253560265",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "335938 many particles\n",
      "575690 many particles\n",
      "1432914 many particles\n",
      "2944254 many particles\n",
      "1617465 many particles\n"
     ]
    }
   ],
   "source": [
    "for i,gmax in enumerate(magnitudeEdges[1:]):\n",
    "    gmin = magnitudeEdges[i]\n",
    "    saveToDisk(dfs[i],gmin,gmax)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 239,
   "id": "46953dc9",
   "metadata": {},
   "outputs": [],
   "source": [
    "from firefly.data_reader import ArrayReader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 265,
   "id": "f27ec1ff-9ea3-49ff-bc09-9e65f6588af2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(array([ 2.,  7., 15., 13., 18., 18., 14.,  6.,  4.,  3.]),\n",
       " array([1.55622705, 2.05395579, 2.55168452, 3.04941326, 3.547142  ,\n",
       "        4.04487074, 4.54259948, 5.04032822, 5.53805696, 6.0357857 ,\n",
       "        6.53351444]),\n",
       " <BarContainer object of 10 artists>)"
      ]
     },
     "execution_count": 265,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 5496 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.hist(temps)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 277,
   "id": "e618500d-94bf-4b01-87e9-1e00d4269d90",
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.interpolate import interp1d"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 240,
   "id": "e0d708b3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[  44.74708905, -198.27386341,   20.21138422],\n",
       "       [  94.1784582 , -341.35400452,   35.64754174],\n",
       "       [  47.42737877, -201.80665677,   21.73701154],\n",
       "       ...,\n",
       "       [-480.48988524, -141.51911534, -861.66774804],\n",
       "       [-136.36987599,  -39.3087655 , -245.32609517],\n",
       "       [-498.94824416, -146.43772507, -892.23702308]])"
      ]
     },
     "execution_count": 240,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df = pd.read_csv('tmpGaia.csv')\n",
    "coordinates = np.zeros((df.shape[0],3))\n",
    "velocities = np.zeros((df.shape[0],3))\n",
    "\n",
    "for i,axis in enumerate(['x','y','z']):\n",
    "    coordinates[:,i] = df[axis]\n",
    "    velocities[:,i] = df[f'v{axis}']\n",
    "coordinates"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 357,
   "id": "e10c023b-d9c5-4794-974a-147153346cec",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "bp_rp -0.1259861 7.2997017\n",
      "phot_g_mean_mag 2.4740036 12.0\n",
      "minus_mag -12.0 -2.4740036\n",
      "log10Temperature 3.3372207063318906 4.999980577312683\n"
     ]
    }
   ],
   "source": [
    "fields = {key:df[key].to_numpy() for key in ['bp_rp','phot_g_mean_mag']}\n",
    "fields['minus_mag'] = -fields['phot_g_mean_mag']\n",
    "fields['log10Temperature'] = temps\n",
    "for key,value in fields.items():\n",
    "    print(key,np.min(value),np.max(value))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 376,
   "id": "b3f6ac2d-a4cc-415a-ad97-164a3ad603d3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 30532 characters]\n",
      "text/plain": [
       "<Figure size 720x480 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 21132 characters]\n",
      "text/plain": [
       "<Figure size 720x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "\n",
    "def remap_nonparametrically(values,old_interval,new_interval):\n",
    "    h,_ = np.histogram(values,bins=np.linspace(*old_interval))\n",
    "    new_edges = np.linspace(*new_interval)\n",
    "    probs = np.cumsum(h)/np.sum(h)\n",
    "    probs[0] = 0\n",
    "    inverter = interp1d(probs,temp_edges[1:])\n",
    "    new_values = inverter(np.random.uniform(0,1,df.shape[0]))\n",
    "    \n",
    "    sort_indices = np.argsort(values)\n",
    "    new_values = np.sort(new_values)[np.argsort(sort_indices)]\n",
    "    return new_values\n",
    "    \n",
    "    \n",
    "\n",
    "old_interval = (0,4)\n",
    "new_interval = (np.log10(2e3),5)\n",
    "\n",
    "h,edges,p = plt.hist(fields['bp_rp'],bins=np.linspace(*old_interval))\n",
    "ax1 = plt.gca()\n",
    "ax2 = ax1.twiny()\n",
    "temps = remap_nonparametrically(fields['bp_rp'],old_interval,new_interval)\n",
    "ax2.hist(temps,bins=np.linspace(*new_interval),histtype='step',color='C1',lw=3)\n",
    "\n",
    "ax1.set_xlabel('bp_rp')\n",
    "ax2.set_xlabel('log10(T)')\n",
    "plt.gcf().set_facecolor('white')\n",
    "plt.gcf().set_dpi(120)\n",
    "\n",
    "plt.figure()\n",
    "plt.plot(bp_rps,temps,'.')\n",
    "plt.gca().set_ylabel('log10(T)')\n",
    "plt.gca().set_xlabel('bp_rp')\n",
    "plt.gcf().set_dpi(120)\n",
    "plt.gcf().set_facecolor('white')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 358,
   "id": "b7f16e44",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "JSONdir: /Users/agurvich/GaiaeDR3 -- is not a sub-directory of firefly/static/data. \n",
      "This may produce confusing or inoperable results. As such, we will create a symlink for you when you writeToDisk.\n",
      "filter/colormap/radius flags correspond to: ['bp_rp', 'phot_g_mean_mag', 'minus_mag', 'log10Temperature']\n",
      "Outputting: eDR3RV - 2344542/2344542 particles - 5 tracked fields\n"
     ]
    }
   ],
   "source": [
    "reader = ArrayReader(\n",
    "    ## don't have to pass as a list with 1 element \n",
    "    ##  if only 1 particle group\n",
    "    coordinates,\n",
    "    velocities,\n",
    "    'eDR3RV',\n",
    "    fields,\n",
    "    JSONdir='GaiaeDR3', ## if not an absolute path assumes from ~\n",
    "    clean_JSONdir=True, ## delete old files in the directory\n",
    "    write_startup=True, ## overwrite the startup file (and turn off PG append)\n",
    "    write_to_disk=False, ## whether to write to disk at the end of __init__\n",
    "    field_filter_flags=[True,True,False,True], ## flags corresponding to fields\n",
    "    field_colormap_flags=[True,True,False,True],\n",
    "    field_radius_flags=[False,False,True,False],\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 360,
   "id": "d03b60df-f7e5-4574-ad01-cfe2f891d4c8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "''"
      ]
     },
     "execution_count": 360,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "reader.settings['sizeMult']['eDR3RV'] = 0.2\n",
    "#reader.settings['radiusVariable']['eDR3-RV'] = 1\n",
    "\n",
    "## pick the last colormap, blackbody\n",
    "index = (256/8-1)\n",
    "reader.settings['colormap']['eDR3RV'] = (index + 0.5) * (8/256)\n",
    "reader.settings['colormapVariable']['eDR3RV'] = 2\n",
    "reader.settings['depthTest'] = False\n",
    "reader.settings['blendingMode'] = 'additive'\n",
    "reader.settings['showColormap']['eDR3RV'] = True\n",
    "reader.settings['GUIExcludeList'] = [\n",
    "    'eDR3RV/dropdown/general/blendingModeSelectors',\n",
    "    'eDR3RV/dropdown/colormap/colormapCheckBox',\n",
    "    'eDR3RV/colorPicker',\n",
    "    'eDR3RV/dropdown/general/radiusVariableSelector']\n",
    "\n",
    "reader.writeToDisk()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 352,
   "id": "e0f8943a-c282-4b8e-87e4-1c320c4f365a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['/Users/agurvich/Gaia_test']"
      ]
     },
     "execution_count": 352,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "reader.copyFireflySourceToTarget('Gaia_test',init_gh_pages=False,overwrite=True,dump_data=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dda402be",
   "metadata": {},
   "outputs": [],
   "source": [
    "reader.createOctrees([True],npart_min_node=1e4,npart_max_node=1e5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8f120f10",
   "metadata": {},
   "outputs": [],
   "source": [
    "## symlink = False -> data will be copied to static/data\n",
    "reader.writeToDisk(symlink=False)"
   ]
  }
 ],
 "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.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
