examples/h3.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import cmdstanpy\n",
    "import astropy.units as u"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rcat = 'YOUR DATA HERE! Make sure that the necessary columns (used below) are present, or modify the code below as appropriate.'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "filtered = rcat[np.logical_and.reduce((\n",
    "    rcat['FLAG'] == 0,  \n",
    "    rcat['SNR'] > 3, \n",
    "    rcat['Vrot'] < 5,\n",
    "    rcat['Teff'] < 7000,\n",
    "    rcat['V_tan'] < 1000,\n",
    "    np.abs(rcat['V_gsr']) < 400,\n",
    "    rcat['R_gal'] > 50,\n",
    "\n",
    "    ~np.isnan(rcat['GAIAEDR3_RA']), \n",
    "    ~np.isnan(rcat['GAIAEDR3_PMRA']), \n",
    "    ~np.isnan(rcat['GAIAEDR3_PARALLAX']),\n",
    "    ~np.isnan(rcat['Vrad']),\n",
    "    ~np.isnan(rcat['X_gal']),\n",
    "))]\n",
    "\n",
    "filtered.size"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "posobs = (~np.isnan(filtered['GAIAEDR3_RA'])).astype('int')\n",
    "pmobs = (~np.isnan(filtered['GAIAEDR3_PMRA'])).astype('int') \n",
    "distobs = (~np.isnan(filtered['dist_adpt'])).astype('int')\n",
    "vlosobs = (~np.isnan(filtered['Vrad'])).astype('int') \n",
    "d = np.nan_to_num(filtered, 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# MULTIPLY kpc by 1000 to get pc\n",
    "# DIVIDE mas by 1000 to get arcsec\n",
    "\n",
    "standata = {\n",
    "    \n",
    "    'N': len(d),\n",
    "    'pos_obs': posobs,\n",
    "    'dist_obs': distobs,\n",
    "    'pm_obs': pmobs,\n",
    "    'vlos_obs': vlosobs,\n",
    "\n",
    "    'ra_measured': d['GAIAEDR3_RA'], \n",
    "    'dec_measured': d['GAIAEDR3_DEC'],\n",
    "    'dist_measured': d['dist_adpt'], # kpc\n",
    "    \n",
    "    'ra_err': d['GAIAEDR3_RA_ERROR'] * u.mas.to(u.deg), \n",
    "    'dec_err': d['GAIAEDR3_DEC_ERROR'] * u.mas.to(u.deg),\n",
    "    'dist_err': d['dist_adpt_err'], # kpc\n",
    "\n",
    "    'Xgc': d['X_gal'],\n",
    "    'Ygc': d['Y_gal'],\n",
    "    'Zgc': d['Z_gal'],\n",
    "\n",
    "    'pmra_measured': d['GAIAEDR3_PMRA'], # mas/yr\n",
    "    'pmdec_measured': d['GAIAEDR3_PMDEC'], # mas/yr\n",
    "    'vlos_measured': d['Vrad'],\n",
    "\n",
    "    'pmra_err': d['GAIAEDR3_PMRA_ERROR'], # mas/yr\n",
    "    'pmdec_err': d['GAIAEDR3_PMDEC_ERROR'], # mas/yr\n",
    "    'vlos_err': d['Vrad_err'],\n",
    "    \n",
    "    'pos_corr': d['GAIAEDR3_RA_DEC_CORR'],\n",
    "    'pm_corr': d['GAIAEDR3_PMRA_PMDEC_CORR'],\n",
    "    'ra_pmra_corr': d['GAIAEDR3_RA_PMRA_CORR'],\n",
    "    'dec_pmdec_corr': d['GAIAEDR3_DEC_PMDEC_CORR'],\n",
    "    'ra_pmdec_corr': d['GAIAEDR3_RA_PMDEC_CORR'],\n",
    "    'dec_pmra_corr': d['GAIAEDR3_DEC_PMRA_CORR'],\n",
    "    \n",
    "    # try different priors!\n",
    "    'alpha_mean': 4.,\n",
    "    'alpha_sigma': 0.1,\n",
    "    'beta_mean': 0.3,\n",
    "    'beta_sigma': 0.05,\n",
    "    \n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sm = cmdstanpy.CmdStanModel(\n",
    "    stan_file='../models/h3p-fixed.stan', \n",
    "    stanc_options={'include_paths': ['../models']},\n",
    "    cpp_options={\"STAN_THREADS\": False, \"STAN_OPENCL\": True},\n",
    ")\n",
    "sm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def init_gen():\n",
    "    return {\n",
    "\n",
    "        'p_gamma': 0.4,\n",
    "        'p_phi0': 65,\n",
    "        'p_beta': 0.5,\n",
    "        'p_alpha': 4.,\n",
    "\n",
    "        'pmra': standata['pmra_measured'] / 5,\n",
    "        'pmdec': standata['pmdec_measured'] / 5,\n",
    "        'dist': standata['dist_measured'],\n",
    "        'vlos': standata['vlos_measured'] / 5,\n",
    "    }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "fit = sm.sample(\n",
    "    data=standata, \n",
    "#     iter_warmup=1,\n",
    "#     iter_sampling=1,\n",
    "    chains=2, \n",
    "#     adapt_delta=0.,\n",
    "#     adapt_init_phase=1000,\n",
    "#     adapt_metric_window=1000,\n",
    "#     adapt_step_size=1000,\n",
    "#     threads_per_chain=4, \n",
    "    inits=init_gen(),\n",
    "    show_progress='notebook',\n",
    "    save_profile=True,\n",
    "    output_dir='../saved',\n",
    "#     max_treedepth=11,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Refer to `gme_py.ipynb` for some examples of what to do next!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}