{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Welcome to the microlensing data challenge tutorial\n",
    "\n",
    "This notebook is intended to provide the framework for non-microlensers to participate in the microlensing data challenge ( http://microlensing-source.org/data-challenge/ ). It requires installation of the pyLIMA, MulensModel and muLAn modules. \n",
    "\n",
    "The data presented in this tutorial is from the Sagan workshop, thus we know the \"true\" parameter values. For an introduction to microlensing data and analysis, please review the slides and hands on session materials from the 2017 Sagan Workshop ( http://nexsci.caltech.edu/workshop/2017/ ). \n",
    "\n",
    "See the data challenge website for the real data, and replace any of the parameter inputs, as necessary for the type of light curve.\n",
    "\n",
    "This notebook ingests initial guesses for the input WFIRST light curve, creates a model light curve using pyLIMA, MulensModel and muLAn, measures the goodness of fit of the model to the data, and recalculates the model for a few steps within a specified grid range. It calculates the chi-squared at all these positions and determines the minimum. \n",
    "\n",
    "To de-bug, uncomment print and plot statements.\n",
    "\n",
    "Note: The scipy.stats.chisquare() function is intended as a placeholder for the user to replace with a goodness of fit algorithm of their choosing.\n",
    "\n",
    "If you would like an introduction to python notebooks, please read this tutorial: https://medium.com/codingthesmartway-com-blog/getting-started-with-jupyter-notebook-for-python-4e7082bd5d46"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "####If you want interactive plots, uncomment this line and add plt.close() after running each plot\n",
    "#%matplotlib notebook\n",
    "\n",
    "###Import required libraries\n",
    "import numpy as np\n",
    "import scipy.optimize as op\n",
    "import matplotlib.pyplot as plt\n",
    "import astropy.units as u\n",
    "import os\n",
    "import sys\n",
    "import time\n",
    "from scipy.stats import chisquare\n",
    "from datetime import datetime\n",
    "from astropy.coordinates import SkyCoord\n",
    "import ConfigParser as cp\n",
    "import matplotlib\n",
    "import pandas as pd\n",
    "from scipy import stats\n",
    "\n",
    "##Change this to the path with your microlensing packages.\n",
    "local_path='/Users/meshkat/python/'\n",
    "\n",
    "sys.path.append(local_path+'pylima/pyLIMA-master')\n",
    "from pyLIMA import event\n",
    "from pyLIMA import telescopes\n",
    "from pyLIMA import microlmodels\n",
    "from pyLIMA import microlcaustics\n",
    "from pyLIMA import microltoolbox\n",
    "\n",
    "sys.path.append(local_path+'MulensModel/source')\n",
    "import MulensModel as mm\n",
    "from MulensModel.utils import Utils\n",
    "\n",
    "import muLAn.mulan as mulan\n",
    "import muLAn.models.ephemeris as ephemeris"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "###define path to file and filename\n",
    "filename = local_path+'data_challenge/WFIRST_1827.dat'\n",
    "\n",
    "WFIRST_data = np.loadtxt(filename)\n",
    "\n",
    "figure1 = plt.figure()\n",
    "plt.errorbar(WFIRST_data[:,0]-2460000,WFIRST_data[:,1],yerr = WFIRST_data[:,2],fmt='.k')\n",
    "plt.xlabel('HJD-2460000',fontsize=20)\n",
    "plt.ylabel('MAG',fontsize=20)\n",
    "plt.gca().invert_yaxis()\n",
    "figure1.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "##Replace any infinite or NAN values with the surrounding data\n",
    "where_inf=np.where(WFIRST_data==np.inf)\n",
    "WFIRST_data[where_inf]=np.mean((WFIRST_data[where_inf[0][0]-1,where_inf[1][0]],\n",
    "                                WFIRST_data[where_inf[0][0]+1,where_inf[1][0]]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "####CAN CHANGE THESE VALUES:\n",
    "#For changes, we recommend you copy and comment these original values, to prevent losing the correct solutions.\n",
    "#Some of these parameters can be changed by looking at the plot above\n",
    "#For definitions of the parameters, see http://nexsci.caltech.edu/workshop/2017/Tuesday_HandsOn_Yee.pptx slide 3.\n",
    "to = 2460963.9  \n",
    "uo = 0.45 \n",
    "tE = 24.3\n",
    "rho = 0.007\n",
    "s = 10**0.198\n",
    "q = 10**-3.23\n",
    "alpha = -0.5 #######in radians\n",
    "\n",
    "log_s = np.log10(s)\n",
    "log_q = np.log10(q)\n",
    "\n",
    "# Time range of planetary perturbation (including 2460000).\n",
    "(t_anomaly_start, t_anomaly_stop) = (2460982., 2460985.)\n",
    "\n",
    "### time when the peak of the planet anomaly occurs (HJD)\n",
    "t_anomaly = 2460983.0\n",
    "\n",
    "### duration of the anomaly in days\n",
    "delta_t = 1\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "####CAN CHANGE THESE VALUES:\n",
    "####How much you want to vary each parameter by, in the grid search:\n",
    "####The number of grid steps can also be varied per parameter below\n",
    "grid_steps=5\n",
    "to_grid=np.linspace(-0.1,0.1,grid_steps)+to\n",
    "uo_grid=np.linspace(0.95,1.05,grid_steps)*uo\n",
    "tE_grid=np.linspace(-1,1,grid_steps)+tE\n",
    "rho_grid=np.linspace(0.95,1.05,grid_steps)*rho\n",
    "log_s_grid=np.linspace(0.95,1.05,grid_steps)*log_s\n",
    "log_q_grid=np.linspace(0.95,1.05,grid_steps)*log_q\n",
    "alpha_grid=np.linspace(-0.1,0.1,grid_steps)+alpha\n",
    "\n",
    "\n",
    "###For the sake of this example, we will only perform a chi-squared minimization on the tE and uo parameters."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Pylima"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "start_time_pl = datetime.now()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "\n",
    "pl_event = event.Event()\n",
    "\n",
    "\n",
    "pl_event.name = 'WFIRST_1827'\n",
    "\n",
    "### We define a telescope object.\n",
    "wfirst = telescopes.Telescope(name='WFIRST', camera_filter='W149', light_curve_magnitude=WFIRST_data)\n",
    "\n",
    "\n",
    "\n",
    "pl_event.telescopes.append(wfirst)\n",
    "\n",
    "### Let's define a binary model, name in pyLIMA is Uniform Source Binary Lens.\n",
    "\n",
    "binary_model = microlmodels.create_model('USBL', pl_event)\n",
    "binary_model.define_model_parameters()\n",
    "\n",
    "full_chi_sq_pylima=np.zeros((grid_steps,grid_steps))\n",
    "full_model_in_magnitude=np.zeros((grid_steps,grid_steps,WFIRST_data[:-1,0].shape[0]))\n",
    "\n",
    "##Run loop a small grid on different values of your choice (in this case, uo and tE):\n",
    "for j,uo_i in enumerate(uo_grid):\n",
    "    uo0=uo_i\n",
    "    for k,tE_i in enumerate(tE_grid):\n",
    "        tE0=tE_i\n",
    "        \n",
    "        ### Here are your input parameter guesses:\n",
    "        estimate_parameters = [to, uo0, tE0, rho, np.log10(s), np.log10(q), alpha]\n",
    "\n",
    "        pyLIMA_parameters = binary_model.compute_pyLIMA_parameters(estimate_parameters)\n",
    "\n",
    "\n",
    "        pl_model, wfirst_source_flux, wfirst_blend_flux = binary_model.compute_the_microlensing_model(wfirst, pyLIMA_parameters)\n",
    "        model_in_magnitude = microltoolbox.flux_to_magnitude(pl_model)  \n",
    "\n",
    "#         plt.figure()\n",
    "#         plt.subplot(111)\n",
    "#         plt.errorbar(WFIRST_data[:,0]-2460000,WFIRST_data[:,1],yerr = WFIRST_data[:,2],fmt='.k')\n",
    "#         plt.gca().invert_yaxis()\n",
    "#         plt.plot(wfirst.lightcurve_flux[:,0]-2460000, model_in_magnitude,zorder=10)\n",
    "#         plt.title('Plot with parameter 1: '+str(uo0)+' and parameter 2: '+str(tE0))\n",
    "#         plt.show()\n",
    "\n",
    "        full_chi_sq_pylima[j,k],_=chisquare(WFIRST_data[:-1,1], f_exp=model_in_magnitude)\n",
    "        full_model_in_magnitude[j,k,:]=model_in_magnitude\n",
    "    \n",
    "plt.imshow(full_chi_sq_pylima,origin='bottom',extent=[tE_grid[0],tE_grid[-1],uo_grid[0],uo_grid[-1]],aspect=\"auto\")\n",
    "plt.xlabel('tE',fontsize=20)\n",
    "plt.ylabel('uo',fontsize=20)\n",
    "plt.colorbar(label='Chi-squared')\n",
    "plt.show()\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "min_index=np.unravel_index(full_chi_sq_pylima.argmin(),full_chi_sq_pylima.shape)\n",
    "uo0 = uo_grid[min_index[0]]\n",
    "tE0 = tE_grid[min_index[1]]\n",
    "chi_sq_pylima=np.min(full_chi_sq_pylima)\n",
    "model_in_magnitude_pylima=full_model_in_magnitude[min_index[0],min_index[1]]\n",
    "print \"\\n Grid search space is minimized at uo=\",uo_grid[min_index[0]],\" and tE=\",tE_grid[min_index[1]],'\\n'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "end_time_pl=datetime.now()\n",
    "print('Total Runtime: {0}'.format(end_time_pl - start_time_pl))\n",
    "t= format(end_time_pl - start_time_pl)\n",
    "(hours, minutes, seconds) = t.split(':')\n",
    "hours_to_run_pl = float(hours) + float(minutes) / 60 + float(seconds) / 3600"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "ModelID_pl='WFIRST_1827.dat_pylima' #Model ID containing both target ID and solution number e.g. [target]_[solution]\n",
    "t0_pl = to #Time of peak, t0 [days] - priority\n",
    "sig_t0_pl = to_grid[1]-to_grid[0] #Uncertainty in t0 [days] - priority\n",
    "tE_pl = tE0 #Einstein crossing time, tE [days] - priority\n",
    "sig_tE_pl = tE_grid[1]-tE_grid[0] #Uncertainty in tE [days] - priority\n",
    "u0_pl = uo0 #Minimum impact parameter, u0 [normalised by θE] - priority\n",
    "sig_u0_pl =  uo_grid[1]-uo_grid[0] #Uncertainty in u0 - priority\n",
    "rho_pl = rho #Angular source size parameter, rho\n",
    "sig_rho_pl = rho_grid[1]-rho_grid[0] #Uncertainty on rho\n",
    "piEE_pl = 'None'  #Parallax parameter πE,E\n",
    "sig_piEE_pl = 'None'  #Uncertainty on πE,E\n",
    "piEN_pl = 'None'  #Parallax parameter πE,N\n",
    "sig_piEN_pl = 'None'  #Uncertainty on πE,N\n",
    "fs_pl = 0  #Source flux, fs [counts/s] - priority\n",
    "sig_fs_pl = 0  #Uncertainty in fs [counts/s] - priority\n",
    "fb_pl = 0 #Blend flux, fb [counts/s] - priority\n",
    "sig_fb_pl = 0 #Uncertainty in fb [counts/s] - priority\n",
    "s_pl = 10**(log_s) #Binary separation, s, [normalised by θE] - priority\n",
    "sig_s_pl = 10**(log_s_grid[1]-log_s_grid[0]) #Uncertainty on s - priority\n",
    "q_pl = 10**(log_q) #Mass ratio, q = M2/M1 - priority\n",
    "sig_q_pl = 10**(log_q_grid[1]-log_q_grid[0]) #Uncertainty on q - priority\n",
    "alpha_pl = alpha #Angle of lens motion, alpha - priority\n",
    "sig_alpha_pl = alpha_grid[1]-alpha_grid[0] #Uncertainty on alpha - priority\n",
    "dsdt_pl = 'None'  #Rate of change of s, ds/dt\n",
    "sig_dsdt_pl = 'None'  #Uncertainty on ds/dt\n",
    "dadt_pl =  'None' #Rate of change of alpha, dalpha/dt\n",
    "sig_dadt_pl = 'None'  #Uncertainty on dalpha/dt\n",
    "t0_par_pl = 'None'  #t0_par [days]\n",
    "chisq_pl = chi_sq_pylima #Chi squared of the fitted model\n",
    "M1_pl = 'None'  #Primary lens mass, M1 [Msolar]\n",
    "sig_M1_pl = 'None'  #Uncertainty on M1 [Msolar]\n",
    "M2_pl = 'None'  #Secondary lens mass, M2 [MJupiter] \n",
    "sig_M2_pl = 'None'  #Uncertainty on M2 [MJupiter] \n",
    "DL_pl = 'None'  #Distance to the lens, DL [kpc]\n",
    "sig_DL_pl = 'None'  #Uncertainty on DL [kpc]\n",
    "DS_pl = 'None'  #Distance to the source, DS [kpc]\n",
    "sig_DS_pl = 'None'  #Uncertainty on DS [kpc]\n",
    "aperp_pl = 'None'  #Projected separation of lens masses, aperp [AU]\n",
    "sig_aperp_pl = 'None'  #Uncertainty on aperp [AU]\n",
    "t_fit_pl = hours_to_run_pl #Time taken to fit the lightcurve from data ingest to final output [hrs]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# MulensModel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "start_time_mlm = datetime.now()\n",
    "\n",
    "data = mm.MulensData(file_name=filename)\n",
    "\n",
    "### Uncomment this to view the lightcurve: \n",
    "# plt.errorbar(data.time, data.mag, yerr=data.err_mag, fmt='o')\n",
    "# plt.gca().invert_yaxis()\n",
    "# plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# *Set the magnification methods for the planet model*\n",
    "# VBBL method will be used between t_planet_start and t_planet_stop, \n",
    "# and point_source_point_lens will be used everywhere else.\n",
    "magnification_methods = [\n",
    "    0., 'point_source_point_lens', \n",
    "    t_anomaly_start, 'VBBL', t_anomaly_stop, \n",
    "    'point_source_point_lens', 2470000.]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "data.bad = np.isnan(data.err_mag)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Check the estimated model\n",
    "# Note that there are two possibilities for s: s_plus and s_minus. \n",
    "\n",
    "full_chi_sq_mulens=np.zeros((grid_steps,grid_steps))\n",
    "full_model_in_magnitude=np.zeros((grid_steps,grid_steps,WFIRST_data[:,0].shape[0]))\n",
    "\n",
    "#Run loop a small grid on different values of your choice (in this case, uo and tE):\n",
    "for j,uo_i in enumerate(uo_grid):\n",
    "    uo0=uo_i\n",
    "    for k,tE_i in enumerate(tE_grid):\n",
    "        tE0=tE_i\n",
    "\n",
    "        # Define the model\n",
    "        mlm_model = mm.Model({\n",
    "            't_0': to, \n",
    "            'u_0': uo0,\n",
    "            't_E': tE0,\n",
    "            'rho': rho,\n",
    "            's': s,\n",
    "            'q': q,\n",
    "            'alpha': np.rad2deg(alpha)})\n",
    "\n",
    "        mlm_model.set_magnification_methods(magnification_methods)\n",
    "        mlm_event = mm.Event(datasets=data, model=mlm_model)\n",
    "\n",
    "        times=WFIRST_data[:,0]\n",
    "\n",
    "        mlm_event.get_chi2()\n",
    "        flux = mlm_event.fit.get_flux(data)\n",
    "\n",
    "        model_in_magnitude0=Utils.get_mag_from_flux(flux)\n",
    "        \n",
    "#         plt.figure()\n",
    "#         plt.subplot(111)\n",
    "#         plt.errorbar(WFIRST_data[:,0]-2460000,WFIRST_data[:,1],yerr = WFIRST_data[:,2],fmt='.k')\n",
    "#         plt.gca().invert_yaxis()\n",
    "#         plt.plot(times-2460000, model_in_magnitude0,zorder=10)\n",
    "#         plt.title('Plot with parameter 1: '+str(uo0)+' and parameter 2: '+str(tE0))\n",
    "#         plt.show()\n",
    "\n",
    "        full_chi_sq_mulens[j,k],_=chisquare(WFIRST_data[:,1],f_exp=model_in_magnitude0)\n",
    "        full_model_in_magnitude[j,k,:]=model_in_magnitude0\n",
    "        \n",
    "plt.imshow(full_chi_sq_mulens,origin='bottom',extent=[tE_grid[0],tE_grid[-1],uo_grid[0],uo_grid[-1]],aspect=\"auto\")\n",
    "plt.xlabel('tE',fontsize=20)\n",
    "plt.ylabel('uo',fontsize=20)\n",
    "plt.colorbar(label='Chi-squared')\n",
    "plt.show()\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "min_index=np.unravel_index(full_chi_sq_mulens.argmin(),full_chi_sq_mulens.shape)\n",
    "uo0 = uo_grid[min_index[0]]\n",
    "tE0 = tE_grid[min_index[1]]\n",
    "chi_sq_mulens=np.min(full_chi_sq_mulens)\n",
    "model_in_magnitude_mulens=full_model_in_magnitude[min_index[0],min_index[1]]\n",
    "print \"\\n Grid search space is minimized at uo=\",uo_grid[min_index[0]],\" and tE=\",tE_grid[min_index[1]],\"\\n\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "end_time_mlm = datetime.now()\n",
    "print('Total Runtime: {0}'.format(end_time_mlm - start_time_mlm))\n",
    "t= format(end_time_mlm - start_time_mlm)\n",
    "(hours, minutes, seconds) = t.split(':')\n",
    "hours_to_run_mlm = float(hours) + float(minutes) / 60 + float(seconds) / 3600"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "ModelID_mlm= 'WFIRST_1827.dat_MulensModel' #Model ID containing both target ID and solution number e.g. [target]_[solution]\n",
    "t0_mlm = to #Time of peak, t0 [days] - priority\n",
    "sig_t0_mlm = to_grid[1]-to_grid[0] #Uncertainty in t0 [days] - priority\n",
    "tE_mlm = tE0  #Einstein crossing time, tE [days] - priority\n",
    "sig_tE_mlm = tE_grid[1]-tE_grid[0] #Uncertainty in tE [days] - priority\n",
    "u0_mlm = uo0  #Minimum impact parameter, u0 [normalised by θE] - priority\n",
    "sig_u0_mlm = uo_grid[1]-uo_grid[0] #Uncertainty in u0 - priority\n",
    "rho_mlm = rho #Angular source size parameter, rho\n",
    "sig_rho_mlm = rho_grid[1]-rho_grid[0] #Uncertainty on rho\n",
    "piEE_mlm = 'None'  #Parallax parameter πE,E\n",
    "sig_piEE_mlm = 'None'  #Uncertainty on πE,E\n",
    "piEN_mlm = 'None'  #Parallax parameter πE,N\n",
    "sig_piEN_mlm = 'None'  #Uncertainty on πE,N\n",
    "fs_mlm = 'None' #Source flux, fs [counts/s] - priority\n",
    "sig_fs_mlm = 'None' #Uncertainty in fs [counts/s] - priority\n",
    "fb_mlm = 0 #Blend flux, fb [counts/s] - priority\n",
    "sig_fb_mlm = 0 #Uncertainty in fb [counts/s] - priority\n",
    "s_mlm = 10**(log_s) #Binary separation, s, [normalised by θE] - priority\n",
    "sig_s_mlm = 10**(log_s_grid[1]-log_s_grid[0]) #Uncertainty on s - priority\n",
    "q_mlm = 10**(log_q) #mlm_event.model.parameters.q #Mass ratio, q = M2/M1 - priority\n",
    "sig_q_mlm = 10**(log_q_grid[1]-log_q_grid[0]) #Uncertainty on q - priority\n",
    "alpha_mlm = alpha #mlm_event.model.parameters.alpha #Angle of lens motion, alpha - priority\n",
    "sig_alpha_mlm = alpha_grid[1]-alpha_grid[0] #Uncertainty on alpha - priority\n",
    "dsdt_mlm = 'None'  #Rate of change of s, ds/dt\n",
    "sig_dsdt_mlm = 'None'  #Uncertainty on ds/dt\n",
    "dadt_mlm =  'None' #Rate of change of alpha, dalpha/dt\n",
    "sig_dadt_mlm = 'None'  #Uncertainty on dalpha/dt\n",
    "t0_par_mlm = 'None'  #t0_par [days]\n",
    "chisq_mlm = chi_sq_mulens #Chi squared of the fitted model\n",
    "M1_mlm = 'None'  #Primary lens mass, M1 [Msolar]\n",
    "sig_M1_mlm = 'None'  #Uncertainty on M1 [Msolar]\n",
    "M2_mlm = 'None'  #Secondary lens mass, M2 [MJupiter] \n",
    "sig_M2_mlm = 'None'  #Uncertainty on M2 [MJupiter] \n",
    "DL_mlm = 'None'  #Distance to the lens, DL [kpc]\n",
    "sig_DL_mlm = 'None'  #Uncertainty on DL [kpc]\n",
    "DS_mlm = 'None'  #Distance to the source, DS [kpc]\n",
    "sig_DS_mlm = 'None'  #Uncertainty on DS [kpc]\n",
    "aperp_mlm = 'None'  #Projected separation of lens masses, aperp [AU]\n",
    "sig_aperp_mlm = 'None'  #Uncertainty on aperp [AU]\n",
    "t_fit_mlm = hours_to_run_mlm #Time taken to fit the lightcurve from data ingest to final output [hrs]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# MuLAN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "start_time_mln = datetime.now()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "cfgsetup = cp.SafeConfigParser()\n",
    "cfgsetup.add_section('EventDescription')\n",
    "cfgsetup.set('EventDescription', 'RA', '17h57m11.52s')\n",
    "cfgsetup.set('EventDescription', 'DEC', '-29d07m03.36s')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Import the model\n",
    "import muLAn.models.BLcontU as esbl_vbb\n",
    "\n",
    "# Define the dates we want to compute the model (e.g. from 8350 to 9122)\n",
    "n_dates = WFIRST_data[:,0].shape[0]\n",
    "timeserie = np.linspace(to-100-2460000, to+100-2460000, n_dates)\n",
    "\n",
    "\n",
    "full_chi_sq_mulan=np.zeros((grid_steps,grid_steps))\n",
    "full_model_in_magnitude=np.zeros((grid_steps,grid_steps,WFIRST_data[:,0].shape[0]))\n",
    "\n",
    "# Prepare dataset before fit: compute the flux\n",
    "photometry = pd.DataFrame({'date': WFIRST_data[:,0] - 2460000,\n",
    "                           'flux': 10**(-0.4*WFIRST_data[:,1])})\n",
    "photometry['magnification'] = 1.0  ## initialization of the column magnification\n",
    "\n",
    "#Run loop a small grid on different values of your choice (in this case, uo and tE):\n",
    "for j,uo_i in enumerate(uo_grid):\n",
    "    uo0=uo_i\n",
    "    for k,tE_i in enumerate(tE_grid):\n",
    "        tE0=tE_i\n",
    "        \n",
    "        # Define the parameters (depends on each model)\n",
    "        # --> in this example, no parallax\n",
    "        lens_params = dict({'u0': -uo0,\n",
    "                            'tE': tE0,\n",
    "                            't0': to-2460000,\n",
    "                            'piEN': 0.0,\n",
    "                            'piEE': 0.0,\n",
    "                            'rho': rho,\n",
    "                            's': s,\n",
    "                            'q': q,\n",
    "                            'alpha': np.pi-alpha,\n",
    "                            'dalpha': 0.0,\n",
    "                            'ds': 0.0\n",
    "                           })\n",
    "        \n",
    "        # No lens orbital motion (dalpha=0, ds=0)\n",
    "        tb = lens_params['t0']  # we choose t_binary = t0 here (see, e.g., Skowron et al. 2011)\n",
    "\n",
    "        # We don't want to include microlens parallax in this fit\n",
    "        Ds = dict({'N':np.zeros(n_dates), 'E':np.zeros(n_dates)})\n",
    "\n",
    "        # Compute magnification at the date of the observations\n",
    "        photometry['magnification'] = esbl_vbb.magnifcalc(photometry['date'].values, \n",
    "                                                          lens_params, Ds=Ds, tb=tb)\n",
    "\n",
    "        # Determine the flux of the source star (fs) and the blend flux (fb)\n",
    "        # By definition, flux_total(t) = fs * A(t) + fb\n",
    "        # with A=magnification\n",
    "        table_sorted = photometry.sort_values('magnification', inplace=False)\n",
    "        x = table_sorted['magnification'].values\n",
    "        y = table_sorted['flux'].values\n",
    "        fs, fb, r_value, p_value, std_err = stats.linregress(x, y)\n",
    "\n",
    "        # Derive the model in magnitudes\n",
    "        model_in_magnitude1 = - 2.5 * np.log10(fs * photometry['magnification'].values + fb)\n",
    "        \n",
    "        full_chi_sq_mulan[j,k],_=chisquare(WFIRST_data[:,1],f_exp=model_in_magnitude1)\n",
    "        full_model_in_magnitude[j,k,:]=model_in_magnitude1\n",
    "        \n",
    "        \n",
    "        # Plot the magnification\n",
    "        x = photometry['date']\n",
    "        y = model_in_magnitude1\n",
    "#         plt.figure()\n",
    "#         plt.plot(WFIRST_data[:,0]-2460000,WFIRST_data[:,1],color='black')\n",
    "#         plt.plot(x, y)\n",
    "#         plt.xlabel('HJD-2460000')\n",
    "#         plt.ylabel('Magnification')\n",
    "#         plt.xlim(910,1000)\n",
    "#         plt.gca().invert_yaxis()\n",
    "#         plt.title('Plot with parameter 1: '+str(uo0)+' and parameter 2: '+str(tE0))\n",
    "#         plt.show()\n",
    "\n",
    "\n",
    "plt.imshow(full_chi_sq_mulan,origin='bottom',extent=[tE_grid[0],tE_grid[-1],uo_grid[0],uo_grid[-1]],aspect=\"auto\")\n",
    "plt.xlabel('tE',fontsize=20)\n",
    "plt.ylabel('uo',fontsize=20)\n",
    "plt.colorbar(label='Chi-squared')\n",
    "plt.show()\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "min_index=np.unravel_index(full_chi_sq_mulan.argmin(),full_chi_sq_mulan.shape)\n",
    "uo0 = uo_grid[min_index[0]]\n",
    "tE0 = tE_grid[min_index[1]]\n",
    "chi_sq_mulan=np.min(full_chi_sq_mulan)\n",
    "model_in_magnitude_mulan=full_model_in_magnitude[min_index[0],min_index[1]]\n",
    "print \"\\n Grid search space is minimized at uo=\",uo_grid[min_index[0]],\" and tE=\",tE_grid[min_index[1]],\"\\n\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "end_time_mln = datetime.now()\n",
    "print('Total Runtime: {0}'.format(end_time_mln - start_time_mln))\n",
    "t= format(end_time_mln - start_time_mln)\n",
    "(hours, minutes, seconds) = t.split(':')\n",
    "hours_to_run_mln = float(hours) + float(minutes) / 60 + float(seconds) / 3600"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "ModelID_mln= 'WFIRST_1827.dat_Mulan' #Model ID containing both target ID and solution number e.g. [target]_[solution]\n",
    "t0_mln = to #Time of peak, t0 [days] - priority\n",
    "sig_t0_mln = to_grid[1]-to_grid[0] #Uncertainty in t0 [days] - priority\n",
    "tE_mln = tE0  #Einstein crossing time, tE [days] - priority\n",
    "sig_tE_mln = tE_grid[1]-tE_grid[0] #Uncertainty in tE [days] - priority\n",
    "u0_mln = uo0  #Minimum impact parameter, u0 [normalised by θE] - priority\n",
    "sig_u0_mln = uo_grid[1]-uo_grid[0] #Uncertainty in u0 - priority\n",
    "rho_mln = rho #Angular source size parameter, rho\n",
    "sig_rho_mln = rho_grid[1]-rho_grid[0] #Uncertainty on rho\n",
    "piEE_mln = 'None'  #Parallax parameter πE,E\n",
    "sig_piEE_mln = 'None'  #Uncertainty on πE,E\n",
    "piEN_mln = 'None'  #Parallax parameter πE,N\n",
    "sig_piEN_mln = 'None'  #Uncertainty on πE,N\n",
    "fs_mln = 'None' #Source flux, fs [counts/s] - priority\n",
    "sig_fs_mln = 'None' #Uncertainty in fs [counts/s] - priority\n",
    "fb_mln = 0 #Blend flux, fb [counts/s] - priority\n",
    "sig_fb_mln = 0 #Uncertainty in fb [counts/s] - priority\n",
    "s_mln = 10**(log_s) #Binary separation, s, [normalised by θE] - priority\n",
    "sig_s_mln = 10**(log_s_grid[1]-log_s_grid[0]) #Uncertainty on s - priority\n",
    "q_mln = 10**(log_q) #mlm_event.model.parameters.q #Mass ratio, q = M2/M1 - priority\n",
    "sig_q_mln = 10**(log_q_grid[1]-log_q_grid[0]) #Uncertainty on q - priority\n",
    "alpha_mln = alpha #mlm_event.model.parameters.alpha #Angle of lens motion, alpha - priority\n",
    "sig_alpha_mln = alpha_grid[1]-alpha_grid[0] #Uncertainty on alpha - priority\n",
    "dsdt_mln = 'None'  #Rate of change of s, ds/dt\n",
    "sig_dsdt_mln = 'None'  #Uncertainty on ds/dt\n",
    "dadt_mln =  'None' #Rate of change of alpha, dalpha/dt\n",
    "sig_dadt_mln = 'None'  #Uncertainty on dalpha/dt\n",
    "t0_par_mln = 'None'  #t0_par [days]\n",
    "chisq_mln = chi_sq_mulan #Chi squared of the fitted model\n",
    "M1_mln = 'None'  #Primary lens mass, M1 [Msolar]\n",
    "sig_M1_mln = 'None'  #Uncertainty on M1 [Msolar]\n",
    "M2_mln = 'None'  #Secondary lens mass, M2 [MJupiter] \n",
    "sig_M2_mln = 'None'  #Uncertainty on M2 [MJupiter] \n",
    "DL_mln = 'None'  #Distance to the lens, DL [kpc]\n",
    "sig_DL_mln = 'None'  #Uncertainty on DL [kpc]\n",
    "DS_mln = 'None'  #Distance to the source, DS [kpc]\n",
    "sig_DS_mln = 'None'  #Uncertainty on DS [kpc]\n",
    "aperp_mln = 'None'  #Projected separation of lens masses, aperp [AU]\n",
    "sig_aperp_mln = 'None'  #Uncertainty on aperp [AU]\n",
    "t_fit_mln = hours_to_run_mln #Time taken to fit the lightcurve from data ingest to final output [hrs]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "##Just to compare the three models, if you want to see the difference. Not doing a perfect fit, for now.\n",
    "plt.figure()\n",
    "plt.plot(WFIRST_data[:,0]-2460000,WFIRST_data[:,1],color='black',label='WFIRST simulated data')\n",
    "plt.plot(times[1:]-2460000, model_in_magnitude_pylima,color='red',label='pyLIMA, chi-sq='+str(np.round(chi_sq_pylima,3)))\n",
    "plt.plot(times-2460000, model_in_magnitude_mulens,color='blue',label='MuLensModel, chi-sq='+str(np.round(chi_sq_mulens,3)))\n",
    "plt.plot(times-2460000, model_in_magnitude_mulan,color='green',label='muLAn, chi-sq='+str(np.round(chi_sq_mulens,3)))\n",
    "#plt.xlim(2460880,2461000)\n",
    "plt.xlim(920,990)\n",
    "plt.gca().invert_yaxis()\n",
    "plt.ylabel('Magnitude')\n",
    "plt.xlabel('HJD-2460000')\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "##Zooming in\n",
    "plt.figure()\n",
    "plt.plot(WFIRST_data[:,0]-2460000,WFIRST_data[:,1],color='black',label='WFIRST simulated data')\n",
    "plt.plot(times[1:]-2460000, model_in_magnitude_pylima,color='red',label='pyLIMA, chi-sq='+str(np.round(chi_sq_pylima,3)))\n",
    "plt.plot(times-2460000, model_in_magnitude_mulens,color='blue',label='MuLensModel, chi-sq='+str(np.round(chi_sq_mulens,3)))\n",
    "plt.plot(times-2460000, model_in_magnitude_mulan,color='green',label='muLAn, chi-sq='+str(np.round(chi_sq_mulens,3)))\n",
    "#plt.xlim(920,990)\n",
    "plt.xlim(982,985)\n",
    "plt.gca().invert_yaxis()\n",
    "plt.ylabel('Magnitude')\n",
    "plt.xlabel('HJD-2460000')\n",
    "plt.legend()\n",
    "plt.savefig(local_path+'data_challenge/models.pdf')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "Class = 'GLplanet' #Classification\n",
    "f = open(local_path+'data_challenge/table1.dat', 'a')\n",
    "f.write('\\n%s %s %s %s %s %s %s %s %s %s %s %s %s %s %s %s %s %s %s %s '\n",
    "          '%s %s %s %s %s %s %s %s %s %s %s %s %s %s %s %s %s %s %s %s %s' %\n",
    "        (ModelID_pl, Class, t0_pl, sig_t0_pl, tE_pl, sig_tE_pl, u0_pl, sig_u0_pl, rho_pl, sig_rho_pl, piEE_pl,\n",
    "         sig_piEE_pl, piEN_pl, sig_piEN_pl, fs_pl, sig_fs_pl, fb_pl, sig_fb_pl, s_pl, sig_s_pl, q_pl, sig_q_pl, \n",
    "         alpha_pl, sig_alpha_pl, dsdt_pl, sig_dsdt_pl, dadt_pl, sig_dadt_pl, t0_par_pl, chisq_pl, M1_pl, \n",
    "         sig_M1_pl, M2_pl, sig_M2_pl, DL_pl, sig_DL_pl, DS_pl, sig_DS_pl, aperp_pl, sig_aperp_pl, t_fit_pl))\n",
    "f.write('\\n%s %s %s %s %s %s %s %s %s %s %s %s %s %s %s %s %s %s %s %s '\n",
    "          '%s %s %s %s %s %s %s %s %s %s %s %s %s %s %s %s %s %s %s %s %s' %\n",
    "        (ModelID_mlm, Class, t0_mlm, sig_t0_mlm, tE_mlm, sig_tE_mlm, u0_mlm, sig_u0_mlm, rho_mlm, sig_rho_mlm, piEE_mlm,\n",
    "         sig_piEE_mlm, piEN_mlm, sig_piEN_mlm, fs_mlm, sig_fs_mlm, fb_mlm, sig_fb_mlm, s_mlm, sig_s_mlm, q_mlm, sig_q_mlm, \n",
    "         alpha_mlm, sig_alpha_mlm, dsdt_mlm, sig_dsdt_mlm, dadt_mlm, sig_dadt_mlm, t0_par_mlm, chisq_mlm, M1_mlm, \n",
    "         sig_M1_mlm, M2_mlm, sig_M2_mlm, DL_mlm, sig_DL_mlm, DS_mlm, sig_DS_mlm, aperp_mlm, sig_aperp_mlm, t_fit_mlm))\n",
    "f.write('\\n%s %s %s %s %s %s %s %s %s %s %s %s %s %s %s %s %s %s %s %s '\n",
    "          '%s %s %s %s %s %s %s %s %s %s %s %s %s %s %s %s %s %s %s %s %s' %\n",
    "        (ModelID_mln, Class, t0_mln, sig_t0_mln, tE_mln, sig_tE_mln, u0_mln, sig_u0_mln, rho_mln, sig_rho_mln, piEE_mln,\n",
    "         sig_piEE_mln, piEN_mln, sig_piEN_mln, fs_mln, sig_fs_mln, fb_mln, sig_fb_mln, s_mln, sig_s_mln, q_mln, sig_q_mln, \n",
    "         alpha_mln, sig_alpha_mln, dsdt_mln, sig_dsdt_mln, dadt_mln, sig_dadt_mln, t0_par_mln, chisq_mln, M1_mln, \n",
    "         sig_M1_mln, M2_mln, sig_M2_mln, DL_mln, sig_DL_mln, DS_mln, sig_DS_mln, aperp_mln, sig_aperp_mln, t_fit_mln))\n",
    "f.close()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}