{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Example Notebook to fit a microlensing event\n",
"\n",
"The procedure is:\n",
"1. Estimate the point lens parameters (`t_0`, `u_0`, `t_E`) from the light curve.\n",
"2. Fit a point lens model (excluding the planetary perturbation).\n",
"3. Estimate the planet parameters (`s`, `alpha`) from the light curve.\n",
"4. Search for the best planetary model using a grid of `s`, `q` (fixed parameters); fits for `rho` and `alpha` but holds the PSPL parameters (`t_0`, `u_0`, `t_E`) fixed.\n",
"\n",
"This notebook is setup to run on a simulated data file: `WFIRST_1827.dat`. To run it on a different data file requires some limited user interaction indicated by `***`:\n",
"1. Set `filename` and `path`\n",
"2. Set the plot limits: `t_min`, `t_max`\n",
"3. Define the time window of the planetary perturbation: `t_planet_start`, `t_planet_stop`\n",
"\n",
"The goal of the notebook is to demonstrate the procedure and run quickly, so it is not robust. The fitting procedure is very simplistic and relies on assuming this is a Gould & Loeb (1996) type planet, which means\n",
"- It is a planetary caustic perturbation\n",
"- It is well described as a _perturbation_, which means\n",
" - The perturbation data can be isolated from the underlying PSPL event \n",
" - The location of the planet can be estimated from the location of the images at the time of the perturbation\n",
" \n",
"This notebook will **not** perform well for:\n",
"- central caustic planets (u_0 << 1 and/or u_0 < rho)\n",
"- resonant caustic planets (s ~ 1)\n",
"- binaries (i.e. the assumption that q << 1 is false)\n",
"\n",
"Simple modifications that could improve performance indicated by `*`:\n",
"- Change the `magnification_methods`, i.e. the method used and the time range it is used for.\n",
"- Change the minimization routine in `fit_model()` (This notebook is set up to use a 'Nelder-Mead' simplex algorithm. Simplex algorithms are known to perform poorly on microlensing events.)\n",
"- Change the size of the grid or the grid steps: `delta_log_s`, `delta_log_q`, `grid_log_s`, `grid_log_q`"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"collapsed": true
},
"outputs": [],
"source": [
"# Import packages\n",
"from datetime import datetime\n",
"start_time = datetime.now() # initialize timer\n",
"import MulensModel as mm\n",
"import matplotlib.pyplot as pl\n",
"import numpy as np\n",
"import scipy.optimize as op\n",
"import os\n",
"import astropy.units as u"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"collapsed": true
},
"outputs": [],
"source": [
"# Define fitting functions\n",
"\n",
"def chi2_fun(theta, event, parameters_to_fit):\n",
" \"\"\" \n",
" Chi2 function. Changes values of the parameters and recalculates chi2.\n",
" \n",
" event = a MulensModel.Event\n",
" parameters_to_fit = list of names of parameters to be changed\n",
" theta = values of the corresponding parameters\n",
" \"\"\"\n",
" # key = the name of the MulensModel parameter\n",
" for (index, key) in enumerate(parameters_to_fit):\n",
" if (key == 't_E' or key =='rho') and theta[index] < 0.:\n",
" return np.inf\n",
" setattr(event.model.parameters, key, theta[index])\n",
" return event.get_chi2()\n",
"\n",
"def fit_model(event, parameters_to_fit):\n",
" \"\"\"\n",
" Fit an \"event\" with \"parameters_to_fit\" as free parameters.\n",
" \n",
" event = a MulensModel event\n",
" parameters_to_fit = list of parameters to fit\n",
" \"\"\"\n",
" # Take the initial starting point from the event.\n",
" x0 = []\n",
" for key in parameters_to_fit:\n",
" value = getattr(event.model.parameters, key)\n",
" if isinstance(value, u.Quantity):\n",
" x0.append(value.value)\n",
" else:\n",
" x0.append(value)\n",
"\n",
" # *Execute fit using a 'Nelder-Mead' algorithm*\n",
" result = op.minimize(\n",
" chi2_fun, x0=x0, args=(event, parameters_to_fit),\n",
" method='Nelder-Mead')\n",
"\n",
" return result"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"collapsed": true
},
"outputs": [],
"source": [
"# ***Read in data file***\n",
"path = os.path.join(mm.MODULE_PATH, \"data\")\n",
"filename = 'WFIRST_1827.dat' # Planet file\n",
"file = os.path.join(path, filename)\n",
"data = mm.MulensData(file_name=file)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Plot the data\n",
"pl.errorbar(data.time, data.mag, yerr=data.err_mag, fmt='o')\n",
"pl.gca().invert_yaxis()\n",
"pl.show()\n",
"\n",
"# ***Define plot limits for a zoom (of the planetary perturbation)***\n",
"(t_min, t_max) = (2460980, 2460990)\n",
"\n",
"# Plot a zoom of the data\n",
"pl.errorbar(data.time - 2460000., data.mag, yerr=data.err_mag, fmt='o')\n",
"pl.xlim(t_min - 2460000., t_max - 2460000.)\n",
"pl.xlabel('t - 2460000')\n",
"pl.gca().invert_yaxis()\n",
"pl.show()"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"collapsed": true
},
"outputs": [],
"source": [
"# ***Set time range of planetary perturbation (including 2460000).***\n",
"(t_planet_start, t_planet_stop) = (2460982., 2460985.)\n",
"\n",
"# *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_planet_start, 'VBBL', t_planet_stop, \n",
" 'point_source_point_lens', 2470000.]"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"collapsed": true
},
"outputs": [],
"source": [
"# Flag data related to the planet\n",
"flag_planet = (data.time > t_planet_start) & (data.time < t_planet_stop) | np.isnan(data.err_mag)\n",
"\n",
"# Exclude those data from the fitting (for now)\n",
"data.bad = flag_planet"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Estimate point lens parameters assuming zero blending\n",
"# \n",
"# Equation for point lens magnification:\n",
"#\n",
"# A(u) = (u^2 + 2) / (u * sqrt(u^2 + 4))\n",
"#\n",
"# where\n",
"# \n",
"# u = sqrt(u_0^2 + tau^2) and tau = (t - t_0) / t_E\n",
"#\n",
"# Thus, the light curve is defined by 3 variables: t_0, u_0, t_E\n",
"#\n",
"\n",
"# Estimate t_0 (time of peak magnification)\n",
"index_t_0 = np.argmin(data.mag[np.invert(flag_planet)])\n",
"t_0 = data.time[index_t_0]\n",
"\n",
"# Estimate u_0\n",
"baseline_mag = np.min([data.mag[0], data.mag[-1]]) # A crude estimate\n",
"A_max = 10.**((data.mag[index_t_0] - baseline_mag) / -2.5)\n",
"u_0 = 1. / A_max # True in the high-magnification limit\n",
"\n",
"# Estimate t_E by determining when the light curve is A(t) = 1.3 (i.e. delta_mag = 0.3)\n",
"t_1 = np.interp(\n",
" baseline_mag - 0.3, data.mag[index_t_0:0:-1], data.time[index_t_0:0:-1])\n",
"t_E = np.abs((t_0 - t_1) / np.sqrt(1. - u_0**2))"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Define the Point Lens Model\n",
"point_lens_model = mm.Model({'t_0': t_0, 'u_0': u_0, 't_E': t_E})\n",
"point_lens_event = mm.Event(datasets=data, model=point_lens_model)\n",
"print('Initial Guess')\n",
"print(point_lens_model)\n",
"\n",
"# Plot (excluded data shown as 'X')\n",
"point_lens_event.plot_model(color='black')\n",
"point_lens_event.plot_data(show_bad=True)\n",
"pl.show()\n",
"print(point_lens_event.get_ref_fluxes())\n",
"print(point_lens_event.model.magnification(t_0))\n",
"\n",
"point_lens_event.plot_model(subtract_2460000=True, color='black', zorder=10)\n",
"point_lens_event.plot_data(show_bad=True, subtract_2460000=True)\n",
"pl.xlim(t_min - 2460000., t_max - 2460000.)\n",
"pl.show()"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Fit the Point Lens Model\n",
"result = fit_model(point_lens_event, parameters_to_fit=['t_0', 'u_0', 't_E'])\n",
"print('Best-fit Point Lens')\n",
"print(point_lens_event.model)\n",
"\n",
"# Plot \n",
"point_lens_event.plot_model(\n",
" t_range=[point_lens_event.model.parameters.t_0 - 5. * point_lens_event.model.parameters.t_E,\n",
" point_lens_event.model.parameters.t_0 + 5. * point_lens_event.model.parameters.t_E],\n",
" color='black', zorder=10)\n",
"point_lens_event.plot_data(show_bad=True)\n",
"pl.show()\n",
"\n",
"point_lens_event.plot_model(color='black', zorder=10)\n",
"point_lens_event.plot_data(show_bad=True)\n",
"pl.xlim(t_min, t_max)\n",
"pl.show()"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"collapsed": true
},
"outputs": [],
"source": [
"# Un-flag planet data (include it in future fits)\n",
"data.bad = np.isnan(data.err_mag)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"collapsed": true
},
"outputs": [],
"source": [
"# Estimate s (projected separation) of the planet, alpha (angle of source trajectory)\n",
"\n",
"# Approximate time of the planetary perturbation\n",
"t_planet = (t_planet_stop + t_planet_start) / 2.\n",
"\n",
"# Position of the source at the time of the planetary perturbation\n",
"tau_planet = ((t_planet - point_lens_event.model.parameters.t_0) /\n",
" point_lens_event.model.parameters.t_E)\n",
"u_planet = np.sqrt(\n",
" point_lens_event.model.parameters.u_0**2 + tau_planet**2)\n",
"\n",
"# Position of the lens images at the time of the planetary perturbation\n",
"# --> Estimate of the planet location\n",
"s_minus = 0.5 * (np.sqrt(u_planet**2 + 4.) - u_planet)\n",
"s_plus = 0.5 * (np.sqrt(u_planet**2 + 4.) + u_planet)\n",
"\n",
"# Angle between the source trajectory and the binary axis\n",
"alpha_planet = np.rad2deg(-np.arctan2(\n",
" point_lens_event.model.parameters.u_0, tau_planet))"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Check the estimated model\n",
"# Note that there are two possibilities for s: s_plus and s_minus. \n",
"# Only s_plus is tested here, but both are considered in the grid search below.\n",
"\n",
"# Define the model\n",
"test_model = mm.Model({\n",
" 't_0': point_lens_event.model.parameters.t_0, \n",
" 'u_0': point_lens_event.model.parameters.u_0,\n",
" 't_E': point_lens_event.model.parameters.t_E,\n",
" 'rho': 0.001,\n",
" 's': s_plus,\n",
" 'q': 10.**(-4),\n",
" 'alpha': alpha_planet})\n",
"test_model.set_magnification_methods(magnification_methods)\n",
"test_event = mm.Event(datasets=data, model=test_model)\n",
"print(test_event.model)\n",
"\n",
"# Plot the model light curve\n",
"test_event.plot_data()\n",
"test_event.plot_model(t_range=[t_min, t_max], color='black', zorder=10)\n",
"pl.xlim(t_min, t_max)\n",
"pl.show()\n",
"\n",
"# Plot the trajectory of the source relative to the caustics\n",
"test_event.model.plot_trajectory(color='black', caustics=True)\n",
"pl.xlim(-1,1)\n",
"pl.ylim(-1,1)\n",
"pl.show()\n",
"# It doesn't have to be perfect, but there should be a planetary perturbation\n",
"# at around the time of the perturbation in the data. If there is no perturbation\n",
"# and/or the source trajectory doesn't pass very near/through the caustics, there is some \n",
"# problem with the model and the fit will likely fail."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Using the Point Lens fit as input, search for a planetary solution\n",
"#\n",
"# Grid parameters: s (log), q (log)\n",
"# Fit parameters: rho, alpha\n",
"# PSPL parameters: t_0, u_0, t_E\n",
"#\n",
"\n",
"# *Define the grid*\n",
"delta_log_s = 0.01\n",
"delta_log_q = 0.25\n",
"grid_log_s = np.hstack(\n",
" (np.arange(\n",
" np.log10(s_minus) - 0.02, np.log10(s_minus) + 0.02, delta_log_s),\n",
" np.arange(\n",
" np.log10(s_plus) - 0.02, np.log10(s_plus) + 0.02, delta_log_s)))\n",
"grid_log_q = np.arange(-5, -2, delta_log_q)\n",
"\n",
"# For each grid point, fit for rho, alpha\n",
"grid = np.empty((5, len(grid_log_s) * len(grid_log_q)))\n",
"i = 0\n",
"print('{0:>12} {1:>6} {2:>7} {3:>7} {4:>7}'.format('chi2', 's', 'q', 'alpha', 'rho'))\n",
"for log_s in grid_log_s:\n",
" for log_q in grid_log_q:\n",
" # The major and minor images are on opposite sides of the lens:\n",
" if log_s < 0.:\n",
" alpha = alpha_planet + 180.\n",
" else:\n",
" alpha = alpha_planet\n",
" \n",
" # Define the Model and Event\n",
" planet_model = mm.Model({\n",
" 't_0': point_lens_event.model.parameters.t_0, \n",
" 'u_0': point_lens_event.model.parameters.u_0,\n",
" 't_E': point_lens_event.model.parameters.t_E,\n",
" 'rho': 0.001,\n",
" 's': 10.**log_s,\n",
" 'q': 10.**log_q,\n",
" 'alpha': alpha})\n",
" planet_model.set_magnification_methods(magnification_methods)\n",
" planet_event = mm.Event(datasets=[data], model=planet_model)\n",
" \n",
" # Fit the Event\n",
" result = fit_model(planet_event, parameters_to_fit=['rho', 'alpha'])\n",
" if result.success:\n",
" chi2 = planet_event.get_chi2()\n",
" else:\n",
" chi2 = np.inf\n",
" \n",
" # Print and store result of fit\n",
" print('{0:12.2f} {1:6.4f} {2:7.5f} {3:7.2f} {4:7.5f}'.format(\n",
" chi2, 10.**log_s, 10.**log_q, \n",
" planet_event.model.parameters.alpha, planet_event.model.parameters.rho))\n",
" grid[0, i] = log_s\n",
" grid[1, i] = log_q\n",
" grid[2, i] = chi2\n",
" grid[3, i] = planet_event.model.parameters.alpha.value\n",
" grid[4, i] = planet_event.model.parameters.rho\n",
" \n",
" i += 1"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Plot the grid\n",
"\n",
"# Identify the best model(s)\n",
"index_best = np.argmin(np.array(grid[2,:]))\n",
"index_sorted = np.argsort(np.array(grid[2,:]))\n",
"n_best = 5 \n",
"colors = ['magenta', 'green', 'cyan','yellow']\n",
"if len(colors) < n_best - 1:\n",
" raise ValueError('colors must have at least n_best -1 entries.')\n",
"\n",
"# Plot the grid\n",
"fig, axes = pl.subplots(nrows=1, ncols=2)\n",
"n_plot = 0\n",
"for i in np.arange(2):\n",
" if i == 0:\n",
" index_logs = np.where(grid_log_s < 0.)[0]\n",
" index_grid = np.where(grid[0, :] < 0.)[0]\n",
" else:\n",
" index_logs = np.where(grid_log_s >= 0.)[0]\n",
" index_grid = np.where(grid[0, :] >= 0.)[0]\n",
" \n",
" # Plot chi2 map\n",
" chi2 = np.transpose(\n",
" grid[2, index_grid].reshape(len(index_logs), len(grid_log_q)))\n",
"\n",
" im = axes[i].imshow(\n",
" chi2, aspect='auto', origin='lower',\n",
" extent=[\n",
" np.min(grid_log_s[index_logs]) - delta_log_s / 2., \n",
" np.max(grid_log_s[index_logs]) + delta_log_s / 2.,\n",
" np.min(grid_log_q) - delta_log_q / 2., \n",
" np.max(grid_log_q) + delta_log_q / 2.],\n",
" cmap='viridis', \n",
" vmin=np.min(grid[2,:]), vmax=np.nanmax(grid[2,np.isfinite(grid[2,:])])) \n",
" \n",
" # Mark best values: best=\"X\", other good=\"o\"\n",
" if index_best in index_grid:\n",
" axes[i].scatter(grid[0, index_best], grid[1, index_best], marker='x', color='white')\n",
" for j, index in enumerate(index_sorted[1:n_best]):\n",
" if index in index_grid:\n",
" axes[i].scatter(grid[0, index], grid[1, index], marker='o', color=colors[j - 1])\n",
" \n",
"fig.subplots_adjust(right=0.9)\n",
"cbar_ax = fig.add_axes([0.95, 0.15, 0.05, 0.7])\n",
"fig.colorbar(im, cax=cbar_ax)\n",
"\n",
"fig.text(0.5, 0.92, r'$\\chi^2$ Map', ha='center')\n",
"fig.text(0.5, 0.04, 'log s', ha='center')\n",
"fig.text(0.04, 0.5, 'log q', va='center', rotation='vertical')\n",
"fig.text(1.1, 0.5, r'$\\chi^2$', va='center', rotation='vertical')\n",
"\n",
"pl.show()\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"collapsed": true
},
"outputs": [],
"source": [
"def make_grid_model(index):\n",
" \"\"\"\n",
" Define a Model using the gridpoint information + the PSPL parameters.\n",
" \n",
" index = index of the grid point for which to generate the model\n",
" \"\"\"\n",
" model = mm.Model({\n",
" 't_0': point_lens_event.model.parameters.t_0, \n",
" 'u_0': point_lens_event.model.parameters.u_0,\n",
" 't_E': point_lens_event.model.parameters.t_E,\n",
" 'rho': grid[4, index],\n",
" 's': 10.**grid[0, index],\n",
" 'q': 10.**grid[1, index],\n",
" 'alpha': grid[3, index]})\n",
" model.set_magnification_methods(magnification_methods)\n",
" return model"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Plot the best-fit model\n",
"best_fit_model = make_grid_model(index_best)\n",
"print('Best Models')\n",
"print(best_fit_model)\n",
" \n",
"best_fit_event = mm.Event(datasets=data, model=best_fit_model)\n",
"(f_source, f_blend) = best_fit_event.get_ref_fluxes()\n",
"\n",
"# Whole model\n",
"t_range_whole = [best_fit_model.parameters.t_0 - 5. * best_fit_model.parameters.t_E,\n",
" best_fit_model.parameters.t_0 + 5. * best_fit_model.parameters.t_E]\n",
"best_fit_event.plot_model(t_range=t_range_whole, subtract_2460000=True, color='black', lw=4)\n",
"best_fit_event.plot_data(subtract_2460000=True)\n",
"pl.show()\n",
"\n",
"\n",
"# Zoom of planet\n",
"t_range_planet = [t_min, t_max]\n",
"# Best model = black\n",
"best_fit_event.plot_data(subtract_2460000=True, s=10, zorder=0)\n",
"best_fit_event.plot_model(\n",
" t_range=t_range_planet, subtract_2460000=True, color='black', lw=3, label='best',\n",
" zorder=10)\n",
"# Other models (color-coding matches grid)\n",
"for j, index in enumerate(index_sorted[1:n_best]):\n",
" model = make_grid_model(index)\n",
" model.plot_lc(\n",
" t_range=t_range_planet, source_flux=f_source, blend_flux=f_blend,\n",
" subtract_2460000=True, color=colors[j - 1], lw=2)\n",
" print(model)\n",
" \n",
"pl.title('{0} best models'.format(n_best))\n",
"pl.xlim(np.array(t_range_planet) - 2460000.)\n",
"pl.legend(loc='best')\n",
"\n",
"pl.show()"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Refine the n_best minima to get the best-fit solution\n",
"parameters_to_fit = ['t_0', 'u_0', 't_E', 'rho', 'alpha', 's', 'q']\n",
"\n",
"fits = []\n",
"for index in index_sorted[:n_best]:\n",
" model = make_grid_model(index)\n",
" event = mm.Event(datasets=data, model=model)\n",
" print(event.model)\n",
" result = fit_model( \n",
" event, parameters_to_fit=parameters_to_fit)\n",
" fits.append([result.fun, result.x])\n",
" print(result)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Plot the best-fit model and output the parameters\n",
"\n",
"# Extract best fit\n",
"chi2 = [x[0] for x in fits]\n",
"fit_parameters = [x[1] for x in fits]\n",
"index_best = np.argmin(chi2)\n",
"\n",
"# Setup the model and event\n",
"parameters = {}\n",
"for i, parameter in enumerate(parameters_to_fit):\n",
" parameters[parameter] = fit_parameters[index_best][i]\n",
" \n",
"final_model = mm.Model(parameters)\n",
"final_model.set_magnification_methods(magnification_methods)\n",
"final_event = mm.Event(datasets=data, model=final_model)\n",
"print(final_event.model)\n",
"print('chi2: {0}'.format(final_event.get_chi2()))\n",
"\n",
"# Plot the whole light curve\n",
"final_event.plot_data(subtract_2460000=True)\n",
"final_event.plot_model(t_range=t_range_whole, \n",
" subtract_2460000=True, color='black', zorder=10)\n",
"pl.show()\n",
"\n",
"# Plot zoom of the planet\n",
"final_event.plot_data(subtract_2460000=True)\n",
"final_event.plot_model(t_range=t_range_planet, subtract_2460000=True, color='black', zorder=10)\n",
"pl.xlim(t_min - 2460000., t_max - 2460000.)\n",
"pl.show()"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"end_time = datetime.now()\n",
"print('Total Runtime: {0}'.format(end_time - start_time))"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"collapsed": true
},
"outputs": [],
"source": []
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3",
"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.6.8"
}
},
"nbformat": 4,
"nbformat_minor": 2
}