{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import print_function, division\n",
    "from six.moves import range\n",
    "import sys\n",
    "import numpy as np\n",
    "import scipy\n",
    "import matplotlib\n",
    "from matplotlib import pyplot as plt\n",
    "import h5py\n",
    "\n",
    "import brutus\n",
    "\n",
    "# plot in-line within the notebook\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "328\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/n/home12/czucker/.conda/envs/fink/lib/python2.7/site-packages/brutus-0.7.5-py2.7.egg/brutus/utils.py:489: RuntimeWarning: overflow encountered in power\n",
      "  phot = 10**(-0.4 * mag) * zeropoints\n"
     ]
    }
   ],
   "source": [
    "from brutus.utils import inv_magnitude\n",
    "filename = '/n/fink2/czucker/Brutus_Cloud_Averages/OrionA/OrionA_l210.9_b-19.8'\n",
    "f = h5py.File(filename+'.h5','r')\n",
    "fpix = f['photometry']['pixel 0-0']\n",
    "mag, magerr = fpix['mag'], fpix['err']\n",
    "mask = np.isfinite(magerr)  # create boolean band mask\n",
    "phot, err = inv_magnitude(mag, magerr)  # convert to flux\n",
    "objid = fpix['obj_id']\n",
    "parallax, parallax_err = fpix['parallax'] * 1e3, fpix['parallax_error'] * 1e3  # convert to mas\n",
    "psel = np.isclose(parallax_err, 0.) | np.isclose(parallax, 0.) | (parallax_err > 1e6)\n",
    "parallax[psel], parallax_err[psel] = np.nan, np.nan\n",
    "coords = np.c_[fpix['l'], fpix['b']]\n",
    "\n",
    "Nobjs = len(phot)\n",
    "print(Nobjs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Filters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Current subset: ['PS_g', 'PS_r', 'PS_i', 'PS_z', 'PS_y', '2MASS_J', '2MASS_H', '2MASS_Ks']\n"
     ]
    }
   ],
   "source": [
    "from brutus import filters\n",
    "\n",
    "filt = filters.ps[:-2] + filters.tmass\n",
    "print('Current subset:', filt)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Reading filter 2MASS_Ks           \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of Bayestar models: 40896\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Reading filter 2MASS_Ks           \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of MIST models: 745055\n"
     ]
    }
   ],
   "source": [
    "# import Bayestar models\n",
    "(models_bs, labels_bs, \n",
    " lmask_bs) = brutus.load_models('../data/grid_bayestar_v5.h5',filters=filt)\n",
    "Nmodels_bs, _, _ = models_bs.shape\n",
    "\n",
    "print('Number of Bayestar models:', Nmodels_bs)\n",
    "\n",
    "# import MIST models\n",
    "(models_mist, labels_mist,\n",
    " lmask_mist) = brutus.load_models('../data/grid_mist_v8.h5', \n",
    "                                  filters=filt, include_ms=True, \n",
    "                                  include_postms=True, include_binaries=False)\n",
    "Nmodels_mist, _, _ = models_mist.shape\n",
    "\n",
    "print('Number of MIST models:', Nmodels_mist)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Zeropoints"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Bayestar:\n",
      "MIST:\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "PS_g (-0.5%)\n",
      "PS_r (-0.4%)\n",
      "PS_i (-0.7%)\n",
      "PS_z (-0.4%)\n",
      "PS_y (0.0%)\n",
      "2MASS_J (-0.3%)\n",
      "2MASS_H (0.9%)\n",
      "2MASS_Ks (0.8%)\n",
      "PS_g (1.0%)\n",
      "PS_r (-3.0%)\n",
      "PS_i (-3.0%)\n",
      "PS_z (-4.0%)\n",
      "PS_y (-3.0%)\n",
      "2MASS_J (-1.0%)\n",
      "2MASS_H (4.0%)\n",
      "2MASS_Ks (4.0%)\n"
     ]
    }
   ],
   "source": [
    "# load zeropoints\n",
    "print('Bayestar:')\n",
    "zp_bs = brutus.load_zeropoints('../data/offsets_bs_v5.txt', \n",
    "                               filters=filt)\n",
    "\n",
    "# load zeropoints\n",
    "print('MIST:')\n",
    "zp_mist = brutus.load_zeropoints('../data/offsets_mist_v8.txt', \n",
    "                                 filters=filt)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load Fitter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load in fitter\n",
    "from brutus import fitting\n",
    "BF_bs = fitting.BruteForce(models_bs, labels_bs, lmask_bs)\n",
    "BF_mist = fitting.BruteForce(models_mist, labels_mist, lmask_mist)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Case 2\n",
    "brutus + Bayestar models with R(V) fixed (tight prior on R(V)=3.3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Fitting object 328/328 [chi2/n: 7.0/9] (mean time: 2.294 s/obj, est. time remaining:      0.000 s)    \n"
     ]
    }
   ],
   "source": [
    "# fit a set of hypothetical objects\n",
    "BF_bs.fit(phot, err, mask, \n",
    "          objid,  # additional information to be written out\n",
    "          'case_two',  # filename (.h5 to be added)\n",
    "          data_coords=coords,  # array of (l, b) coordinates used for distance prior\n",
    "          parallax=parallax, parallax_err=parallax_err,  # parallax measurements\n",
    "          rv_lim=(1.,8.), #default rv range\n",
    "          rv_gauss=(3.32, 1e-6),\n",
    "          phot_offsets=zp_bs,  # photometric offsets read in earlier\n",
    "          avlim=(0., 6.), av_gauss=(3,1e+6),\n",
    "          Ndraws=2000,  # number of samples to save\n",
    "          Nmc_prior=100,  # number of Monte Carlo draws used to integrate priors\n",
    "          logl_dim_prior=True,  # use chi2 distribution instead of Gaussian\n",
    "          save_dar_draws=True,  # save (dist, Av, Rv) samples\n",
    "          running_io=True,  # write out objects as soon as they finish\n",
    "          verbose=True)  # print progress"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Case 3\n",
    "brutus + Bayestar models with R(V) free (default prior)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "BF_bs.fit(phot, err, mask, \n",
    "          objid,  # additional information to be written out\n",
    "          'case_three',  # filename (.h5 to be added)\n",
    "          data_coords=coords,  # array of (l, b) coordinates used for distance prior\n",
    "          parallax=parallax, parallax_err=parallax_err,  # parallax measurements\n",
    "          rv_lim=(1.,8.), #default rv range\n",
    "          rv_gauss=(3.32, 0.18), #default rv_gauss\n",
    "          phot_offsets=zp_bs,  # photometric offsets read in earlier\n",
    "          avlim=(0., 6.), av_gauss=(3,1e+6),\n",
    "          Ndraws=2000,  # number of samples to save\n",
    "          Nmc_prior=100,  # number of Monte Carlo draws used to integrate priors\n",
    "          logl_dim_prior=True,  # use chi2 distribution instead of Gaussian\n",
    "          save_dar_draws=True,  # save (dist, Av, Rv) samples\n",
    "          running_io=True,  # write out objects as soon as they finish\n",
    "          verbose=True)  # print progress"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Case Four\n",
    "brutus + MIST models with R(V) free (default prior), excluding binaries (`include_binaries=False`)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# MIST\n",
    "BF_mist.fit(phot, err, mask, objid, \n",
    "            'case_four',\n",
    "            parallax=parallax, parallax_err=parallax_err, \n",
    "            data_coords=coords, \n",
    "            rv_lim=(1.,8.), #default rv range\n",
    "            rv_gauss=(3.32, 0.18), #default rv_gauss\n",
    "            avlim=(0., 6.), av_gauss=(3,1e+6),\n",
    "            phot_offsets=zp_mist,\n",
    "            running_io=False,\n",
    "            include_binaries=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Case 5\n",
    "brutus + MIST models with R(V) free and default 3-D dust map prior (`Bayestar19`)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# MIST\n",
    "BF_mist.fit(phot, err, mask, objid, \n",
    "            'case_four',\n",
    "            parallax=parallax, parallax_err=parallax_err, \n",
    "            data_coords=coords, \n",
    "            rv_lim=(1.,8.), #default rv range\n",
    "            rv_gauss=(3.32, 0.18), #default rv_gauss\n",
    "            dustfile='../data/bayestar2019_v1.h5',\n",
    "            phot_offsets=zp_mist,\n",
    "            running_io=False,\n",
    "            include_binaries=True)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:.conda-fink] *",
   "language": "python",
   "name": "conda-env-.conda-fink-py"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}