figures.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "sns.set_context('paper', font_scale=1.6)\n",
    "import keras.backend as K\n",
    "\n",
    "SEED = 0\n",
    "SAVE_PLOTS = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "import glob\n",
    "import itertools\n",
    "import os\n",
    "from sklearn.metrics import confusion_matrix\n",
    "from sklearn.model_selection import GridSearchCV, ParameterGrid\n",
    "\n",
    "RF_PARAM_GRID = {'n_estimators': [50, 100, 250], 'criterion': ['gini', 'entropy'],\n",
    "                 'max_features': [0.05, 0.1, 0.2, 0.3], 'min_samples_leaf': [1, 2, 3]}\n",
    "#RF_PARAM_GRID = {'bootstrap': [True], 'class_weight': [None], 'criterion': ['entropy'], 'max_depth': [None], 'max_features': [0.3], 'max_leaf_nodes': [None], 'min_impurity_split': [1e-07], 'min_samples_leaf': [1], 'min_samples_split': [2], 'min_weight_fraction_leaf': [0.0], 'n_estimators': [250], 'n_jobs': [1], 'oob_score': [False], 'random_state': [0], 'verbose': [0], 'warm_start': [False]}\n",
    "\n",
    "def parse_logs(log_files, max_epoch=250):\n",
    "    logs = [pd.read_csv(f, index_col='epoch', parse_dates=[1]) for f in log_files]\n",
    "    for log, f in zip(logs, log_files):\n",
    "        run = f.split('/')[-2]\n",
    "        log.drop(log.index[log.index > max_epoch], axis=0, inplace=True)\n",
    "        if 'time' not in log:\n",
    "            raise ValueError(\"Missing times from {}\".format(f))\n",
    "        log.columns = ['time', run + ' Train', run + ' Valid']\n",
    "        log['time'] = (log['time'] - log['time'].min()) / np.timedelta64(1, 's')\n",
    "    step_logs = pd.concat([l.drop('time', axis=1, inplace=False) for l in logs], axis=1)\n",
    "    time_logs = pd.concat([l.set_index('time') for l in logs], axis=1)\n",
    "    \n",
    "    return step_logs, time_logs\n",
    "\n",
    "\n",
    "def run_to_args(run):\n",
    "    values = run.split(' ')[0].split('_')\n",
    "    args = {'model_type': values[0].upper(), 'size': int(values[1]), \n",
    "            'num_layers': int(values[2].strip('x')),\n",
    "            'drop_frac': float(values[4].strip('drop')),\n",
    "            'bidirectional': 'bidir' in run}\n",
    "    if 'emb' in run:\n",
    "        args['embedding'] = int([v for v in values if 'emb' in v][0][3:])\n",
    "    return args\n",
    "    \n",
    "\n",
    "def training_plot(logs, loss_type='Valid', ylim=(1e-2, 1e0)):\n",
    "    LABEL_MAP = {'epoch': 'Epoch', 'time': 'Time (s)'}\n",
    "    sns.set_style(\"darkgrid\")\n",
    "    fig = plt.figure()\n",
    "    colors = sns.color_palette('deep', n_colors=int(len(logs.columns) / 2))\n",
    "    for i, c in enumerate(logs.columns):\n",
    "        if loss_type is not None and loss_type not in c:\n",
    "            continue\n",
    "        to_plot = logs.iloc[:, i]#logs[c]\n",
    "        to_plot = to_plot.dropna(inplace=False)\n",
    "        args = run_to_args(c)\n",
    "        to_plot.name = \"{model_type}({size}) x {num_layers}\".format(**args)\n",
    "        if 'embedding' in args:\n",
    "            to_plot.name += \", Embedding={}\".format(args['embedding'])\n",
    "        to_plot.plot(color=colors[int(i / 2)], legend=True, logy=True)\n",
    "    plt.xlabel(LABEL_MAP[logs.index.name])\n",
    "    plt.ylabel('Loss')\n",
    "    plt.ylim(ylim)\n",
    "    return fig\n",
    "\n",
    "\n",
    "def plot_confusion_matrix(y, y_pred, classnames, filename=None):\n",
    "    classnames = sorted(classnames)\n",
    "    sns.set_style(\"whitegrid\", {'axes.grid' : False})\n",
    "    cm = confusion_matrix(y, y_pred, classnames)\n",
    "    plt.imshow(cm, interpolation='nearest', cmap=plt.cm.Blues, vmin=0,\n",
    "               vmax=np.unique(y, return_counts=True)[1].max())\n",
    "    #plt.title(title)\n",
    "    #plt.colorbar()\n",
    "    tick_marks = np.arange(len(classnames))\n",
    "    plt.xticks(tick_marks, classnames, rotation=45)\n",
    "    plt.yticks(tick_marks, classnames)\n",
    "\n",
    "    thresh = cm.max() / 2.\n",
    "    for i, j in itertools.product(range(cm.shape[0]), range(cm.shape[1])):\n",
    "        plt.text(j, i, cm[i, j],\n",
    "                 horizontalalignment=\"center\",\n",
    "                 color=\"white\" if cm[i, j] > thresh else \"black\")\n",
    "\n",
    "    plt.tight_layout()\n",
    "    plt.ylabel('True label')\n",
    "    plt.xlabel('Predicted label');\n",
    "    if filename:\n",
    "        plt.savefig(filename)\n",
    "        \n",
    "        from copy import deepcopy\n",
    "\n",
    "def oob_search(model, param_lists, X, y):\n",
    "    best_model = None\n",
    "    best_params = None\n",
    "    best_score = -np.inf\n",
    "    \n",
    "    for g in ParameterGrid(param_lists):\n",
    "        model.set_params(**g)\n",
    "        model.fit(X, y)\n",
    "        if model.oob_score_ > best_score:\n",
    "            best_score = model.oob_score_\n",
    "            best_params = g\n",
    "            best_model = deepcopy(model)\n",
    "            \n",
    "    return best_model, best_params, best_score"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "# Write times from TensorBoard CSVs if missing from training.csv\n",
    "import datetime\n",
    "import glob\n",
    "\n",
    "dirnames = glob.glob('keras_logs/autoencoder/uneven/noise0.5/gru_064_x2*')\n",
    "#dirnames = [d for d in dirnames if 'gru' in d or 'lstm' in d]\n",
    "\n",
    "for d in dirnames:\n",
    "    training = d + '/training.csv'\n",
    "    training_df = pd.read_csv(training)\n",
    "    if 'time' not in training_df:\n",
    "        try:\n",
    "            tag_loss = glob.glob(d + '/*tag_loss.csv')[0]\n",
    "            tag_loss_df = pd.read_csv(tag_loss)\n",
    "            training_df.insert(1, 'time', [str(datetime.datetime.fromtimestamp(v)) for v in tag_loss_df['Wall time'].values])\n",
    "            training_df.to_csv(training, index=False)\n",
    "        except:\n",
    "            pass"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Period estimator training plots"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "log_files = glob.glob('keras_logs/period/uneven/noise0.5/gru_*x2*_bidir/training.csv')\n",
    "\n",
    "step_logs, time_logs = parse_logs(log_files)\n",
    "step_plot = training_plot(step_logs)\n",
    "time_plot = training_plot(time_logs)\n",
    "if SAVE_PLOTS:\n",
    "    step_plot.savefig('paper/figures/gru_period_step.pdf')\n",
    "    time_plot.savefig('paper/figures/gru_period_time.pdf')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "log_files = (glob.glob('keras_logs/period/uneven/noise0.5/gru*64*bidir/training.csv')\n",
    "             + glob.glob('keras_logs/period/uneven/noise0.5/lstm*64*bidir/training.csv'))\n",
    "\n",
    "step_logs, time_logs = parse_logs(log_files)\n",
    "step_plot = training_plot(step_logs)\n",
    "time_plot = training_plot(time_logs)\n",
    "if SAVE_PLOTS:\n",
    "    step_plot.savefig('figures/gru_lstm_period_step.pdf')\n",
    "    time_plot.savefig('figures/gru_lstm_period_time.pdf')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "K.set_learning_phase(0)\n",
    "%run period.py --size 96 --num_layers 2 --drop_frac 0.25 --no_train --model_type gru --sigma 0.5 --lr 5e-4 --sim_type period/uneven/noise0.5 --bidirectional\n",
    "scaler.inverse_transform(Y, copy=False);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "train = np.arange(args.N_train)\n",
    "test = args.N_train + np.arange(args.N_test)\n",
    "pred_gru = model.predict(X)\n",
    "scaler.inverse_transform(pred_gru)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "import gatspy\n",
    "from distutils.version import LooseVersion\n",
    "assert LooseVersion(gatspy.__version__) < '0.3.0'\n",
    "\n",
    "pred_gat = np.zeros(pred_gru.shape)\n",
    "for i in range(pred_gat.shape[0]):\n",
    "    t = np.cumsum(X[i, :, 0])\n",
    "    x = X[i, :, 1]\n",
    "    opt_args = {'period_range': (1.0, 10.0), 'quiet': True}\n",
    "#    opt_args = {'period_range': (0.05, 0.95 * (t.max() - t.min())), 'quiet': True}\n",
    "    model_gat = gatspy.periodic.LombScargleFast(fit_period=True, optimizer_kwds=opt_args)#, silence_warnings=True)\n",
    "    model_gat.fit(t, x)\n",
    "    omega = 2 * np.pi / model_gat.best_period\n",
    "    off, A2, A1 = model_gat._best_params(omega)\n",
    "#    pred_gat[i] = np.array([model_gat.best_period, np.sqrt(A1 ** 2 + A2 ** 2), np.arctan2(A2, A1), off + model_gat.ymean_])\n",
    "    pred_gat[i] = np.array([model_gat.best_period, A1, A2, off + model_gat.ymean_])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "inds = test\n",
    "fig = sns.jointplot(pred_gru[inds, 0] - Y[inds, 0], pred_gat[inds, 0] - Y[inds, 0], xlim=(-1., 1.), ylim=(-1., 1.), stat_func=None)\n",
    "fig.ax_joint.annotate(\"GRU MSE: {:1.4f}\\nLomb-Scargle MSE: {:1.4f}\".format(np.mean((pred_gru[inds, 0] - Y[inds, 0]) ** 2), np.mean((pred_gat[inds, 0] - Y[inds, 0]) ** 2)), (-0.95, 0.78))\n",
    "fig.ax_joint.set_xlabel('GRU error')\n",
    "fig.ax_joint.set_ylabel('Lomb-Scargle error');\n",
    "#plt.title(\"GRU MAE: {}\\nGatspy MAE: {}\".format(np.median(np.abs(pred_gru[:, 0] - Y[inds, 0])), np.median(np.abs(pred_gat[:, 0] - Y[inds, 0]))))\n",
    "#sns.jointplot(pred_gru[inds, 0], Y[inds, 0])\n",
    "#sns.jointplot(pred_gat[inds, 0], Y[inds, 0])\n",
    "if SAVE_PLOTS:\n",
    "    plt.savefig('figures/ls_period.pdf')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Autoencoder training plots"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from natsort import natsorted\n",
    "\n",
    "log_files = natsorted(glob.glob('keras_logs/autoencoder/uneven/noise0.5/gru_064_x2*/training.csv'))\n",
    "log_files = [l for l in log_files if 'emb64' not in l]\n",
    "\n",
    "step_logs, time_logs = parse_logs(log_files)\n",
    "step_plot = training_plot(step_logs, ylim=(0.01, 2.0))\n",
    "time_plot = training_plot(time_logs, ylim=(0.01, 2.0))\n",
    "if SAVE_PLOTS:\n",
    "    step_plot.savefig('paper/figures/gru_autoencoder_step.pdf')\n",
    "    time_plot.savefig('paper/figures/gru_autoencoder_time.pdf')"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "%%time\n",
    "K.set_learning_phase(0)\n",
    "%run autoencoder.py --size 96 --num_layers 2 --drop_frac 0.25 --no_train --model_type gru --sigma 0.5 --lr 5e-4 --sim_type autoencoder/uneven/noise0.5 --embedding 32 --bidirectional"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "def sinusoid(p, A1, A2, b):\n",
    "    return lambda t: A1 * np.cos(2 * np.pi / p * t) + A2 * np.sin(2 * np.pi / p * t) + b\n",
    "\n",
    "N = 3\n",
    "inds = [16, 10, 25]\n",
    "colors = sns.color_palette(n_colors=len(inds))\n",
    "for j, i in enumerate(inds):\n",
    "    t = X_raw[i, :, 0]\n",
    "    m = X[i, :, 1]\n",
    "    T = np.linspace(0, t.max(), 501)\n",
    "    plt.figure()\n",
    "    plt.plot(T, sinusoid(*sample_data.phase_to_sin_cos(Y)[i])(T), color=colors[j])\n",
    "    plt.plot(t, m, 'o', color=colors[j])\n",
    "    plt.xlabel('t')\n",
    "    plt.ylabel('f(t)')\n",
    "    plt.title(f'ω={Y[i, 0]:1.3}, A={Y[i, 1]:1.3}, ϕ={Y[i, 2]:1.3}, b={Y[i, 3]:1.3}')\n",
    "    if SAVE_PLOTS:\n",
    "        plt.savefig(f'figures/sinusoid_{j}.pdf')"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "%%time\n",
    "train = np.arange(args.N_train)\n",
    "test = np.arange(args.N_train, args.N_train + args.N_test)\n",
    "\n",
    "# TODO re-run without this; why so slow...?\n",
    "#train = train[:5000]\n",
    "#X = np.r_[X[train], X[test]]; Y = np.r_[Y[train], Y[test]]; test = np.arange(len(train), len(train) + args.N_test)\n",
    "\n",
    "encode_model = Model(input=model.input, output=model.get_layer('encoding').output)\n",
    "encoding = encode_model.predict({'main_input': X, 'aux_input': np.delete(X, 1, axis=2)})"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "%%time\n",
    "df = pd.DataFrame(np.c_[Y[train], encoding[train]], \n",
    "                  columns=[f'Y_{i}' for i in range(4)] + [f'Embedding {i}' for i in range(encoding.shape[1])])\n",
    "sns.pairplot(df, kind='reg', x_vars=df.columns[:4], y_vars=df.columns[4:])\n",
    "\n",
    "df['Y_0'] **= -1\n",
    "df.rename(columns={'Y_0': 'Period', 'Y_2': 'Phase (radian)'}, inplace=True)\n",
    "R = df.corr(method='spearman')\n",
    "\n",
    "emb_period = np.argmax(np.abs(R.iloc[Y.shape[1]:, 0]))\n",
    "sns.jointplot(df['Period'], df[emb_period], kind='hex', stat_func=None)\n",
    "if SAVE_PLOTS:\n",
    "    plt.savefig('figures/period_hex.pdf')\n",
    "\n",
    "emb_phase  = np.argmax(np.abs(R.iloc[Y.shape[1]:, 2]))\n",
    "sns.jointplot(df['Phase (radian)'], df[emb_phase], kind='hex', stat_func=None)\n",
    "if SAVE_PLOTS:\n",
    "    plt.savefig('figures/phase_hex.pdf')"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "from sklearn.svm import SVR\n",
    "from sklearn.ensemble import RandomForestRegressor\n",
    "\n",
    "model = RandomForestRegressor(n_estimators=1000)\n",
    "model.fit(encoding[train], 1. / Y[train, 0])\n",
    "sns.regplot(1. / Y[test, 0], model.predict(encoding[test]), line_kws={'color': 'black', 'linestyle': 'dotted'})\n",
    "plt.annotate(f\"Mean squared error: {np.mean((1. / Y[test, 0] - model.predict(encoding[test])) ** 2):1.4f}\", (0.8, 9.8))\n",
    "plt.xlabel('True period')\n",
    "plt.ylabel('Estimated period')\n",
    "if SAVE_PLOTS:\n",
    "    plt.savefig('figures/period_rf.pdf')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## ASAS reconstructions"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "from natsort import natsorted\n",
    "\n",
    "log_files = natsorted(glob.glob('keras_logs/asas_full/n200_ls0.2/gru_064_x2_5m04_drop25_emb*_bidir/training.csv'))\n",
    "\n",
    "step_logs, time_logs = parse_logs(log_files)\n",
    "step_plot = training_plot(step_logs, ylim=(3e-2, 1e-1))\n",
    "time_plot = training_plot(time_logs, ylim=(3e-2, 1e-1))\n",
    "if SAVE_PLOTS:\n",
    "    step_plot.savefig('figures/asas_autoencoder_step.pdf')\n",
    "    time_plot.savefig('figures/asas_autoencoder_time.pdf')"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "arg_grid = ParameterGrid({'size': [64, 96], 'num_layers': [2, ], 'model_type': ['lstm', 'gru'],\n",
    "    'sim_type': ['asas_full/n200_ls0.2'], 'n_min': [200], 'n_max': [200], 'lomb_score': [0.2],\n",
    "    'lr': [1e-3, 5e-4], 'embedding': [16, 32, 48, 64, 96], 'bidirectional': [True, False],\n",
    "    'drop_frac': [0.25], 'survey_files': [['data/asas/full.pkl']], 'no_train': [True]})\n",
    "results = {}\n",
    "for args in arg_grid:\n",
    "    ...\n",
    "    args.__dict__.pop('survey_files')  # can't hash list of fnames\n",
    "    results[tuple(args.__dict__.items())] = (oob_score, asas_model.score(encoding[valid], y[valid]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "import os\n",
    "import joblib\n",
    "from keras_util import get_run_id\n",
    "from survey_autoencoder import main as survey_autoencoder\n",
    "\n",
    "arg_dict = {'size': 64, 'embedding': 64, 'num_layers': 1, 'model_type': 'gru',\n",
    "            'sim_type': 'asas_full/n200_ss0.7', 'n_min': 200, 'n_max': 200,\n",
    "            'lr': 1e-3, 'bidirectional': True, 'ss_resid': 0.7,\n",
    "            'drop_frac': 0.25, 'survey_files': ['data/asas/full.pkl'],\n",
    "            'period_fold': False, 'no_train': True}\n",
    "run = get_run_id(**arg_dict)\n",
    "log_dir = os.path.join(os.getcwd(), 'keras_logs', arg_dict['sim_type'], run)\n",
    "weights_path = os.path.join(log_dir, 'weights.h5')\n",
    "if not os.path.exists(weights_path):\n",
    "    raise FileNotFoundError(weights_path)\n",
    "X, X_raw, model, means, scales, wrong_units, args = survey_autoencoder(arg_dict)\n",
    "full = joblib.load('data/asas/full.pkl')\n",
    "split = [el for lc in full for el in lc.split(args.n_min, args.n_max)]\n",
    "if args.ss_resid:\n",
    "    split = [el for el in split if el.ss_resid <= args.ss_resid]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "(len(X), 0.8 * len(X), 0.2 * len(X))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "pred = model.predict({'main_input': X[:], 'aux_input': X[:, :, [0,]]})\n",
    "ords = np.argsort(np.sum((pred.squeeze() - X[:, :, 1]) ** 2, axis=1))\n",
    "inds, labels = zip(*[(8320, '25'), (23994, '75')])\n",
    "for ord_i, label in zip(inds, labels):\n",
    "    i = ords[ord_i]\n",
    "    print(i)\n",
    "    t = X_raw[~wrong_units][i, :, 0]\n",
    "    m = X_raw[~wrong_units][i, :, 1]\n",
    "    e = X_raw[~wrong_units][i, :, 2]\n",
    "    m_est = pred[i] * scales[i, 0] + means[i]\n",
    "    plt.figure()\n",
    "    plt.errorbar(t, m, e, None, 'o');\n",
    "#    plt.errorbar(t, m_est, None, None, 'o');\n",
    "    plt.xlabel('Time [day]')\n",
    "    plt.ylabel('Magnitude')\n",
    "    plt.legend(['Original', 'Reconstructed'])\n",
    "    plt.title(f'{split[i].survey} {split[i].name} ({split[i].label})')\n",
    "    plt.gca().invert_yaxis()\n",
    "    if SAVE_PLOTS:\n",
    "        plt.savefig(f'paper/figures/asas_reconstruct_{label}.pdf')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "import os\n",
    "from keras_util import get_run_id\n",
    "from survey_autoencoder import main as survey_autoencoder\n",
    "\n",
    "arg_dict = {'size': 64, 'embedding': 64, 'num_layers': 1, 'model_type': 'gru',\n",
    "            'sim_type': 'asas_fold/n200', 'n_min': 200, 'n_max': 200,\n",
    "            'lr': 1e-3, 'bidirectional': True, 'ss_resid': 0.7,\n",
    "            'drop_frac': 0.25, 'survey_files': ['data/asas/full.pkl'],\n",
    "            'period_fold': True, 'no_train': True}\n",
    "run = get_run_id(**arg_dict)\n",
    "log_dir = os.path.join(os.getcwd(), 'keras_logs', arg_dict['sim_type'], run)\n",
    "weights_path = os.path.join(log_dir, 'weights.h5')\n",
    "if not os.path.exists(weights_path):\n",
    "    raise FileNotFoundError(weights_path)\n",
    "X_fold, X_raw_fold, model, means, scales, wrong_units, args = survey_autoencoder(arg_dict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "pred_fold = model.predict({'main_input': X_fold[:], 'aux_input': X_fold[:, :, [0,]]})\n",
    "ords_fold = np.argsort(np.sum((pred_fold.squeeze() - X_fold[:, :, 1]) ** 2, axis=1))\n",
    "for ord_i, label in zip(inds, labels):\n",
    "    i = ords[ord_i]\n",
    "    t = X_raw_fold[~wrong_units][i, :, 0]\n",
    "    m = X_raw_fold[~wrong_units][i, :, 1]\n",
    "    e = X_raw_fold[~wrong_units][i, :, 2]\n",
    "    m_est = (pred[i] * scales[i, 0] + means[i])[np.argsort(split[i].times % split[i].p)]\n",
    "    m_fold = pred_fold[i] * scales[i, 0] + means[i]\n",
    "    plt.figure()\n",
    "    plt.errorbar(t, m, e, None, 'o');\n",
    "    plt.errorbar(t, m_est, None, None, 'o', alpha=0.6);\n",
    "    plt.errorbar(t, m_fold, None, None, 'o');\n",
    "    plt.xlabel('Phase [day]')\n",
    "    plt.ylabel('Magnitude')\n",
    "    plt.title(f'{split[i].survey} {split[i].name} ({split[i].label})')\n",
    "    plt.legend(['Original', 'Reconstructed (non-folded)', 'Reconstructed (folded)'])\n",
    "    plt.gca().invert_yaxis()\n",
    "    if SAVE_PLOTS:\n",
    "        plt.savefig(f'figures/asas_folded_reconstruct_{label}.pdf')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## ASAS autoencoder random forests"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "import os\n",
    "import joblib\n",
    "from sklearn.model_selection import StratifiedKFold, GridSearchCV\n",
    "from sklearn.ensemble import RandomForestClassifier\n",
    "from light_curve import LightCurve\n",
    "from keras.preprocessing.sequence import pad_sequences\n",
    "from keras.models import Model\n",
    "from keras_util import get_run_id\n",
    "from survey_autoencoder import main as survey_autoencoder, preprocess\n",
    "\n",
    "arg_dict = {'size': 96, 'embedding': 64, 'num_layers': 2, 'model_type': 'gru',\n",
    "            'sim_type': 'asas_fold_early/n200_ss0.7', 'n_min': 200, 'n_max': 200,\n",
    "            'lr': 5e-4, 'bidirectional': True, 'ss_resid': 0.7,\n",
    "            'drop_frac': 0.25, 'survey_files': ['data/asas/full.pkl'],\n",
    "            'period_fold': True, 'no_train': True}\n",
    "run = get_run_id(**arg_dict)\n",
    "log_dir = os.path.join(os.getcwd(), 'keras_logs', arg_dict['sim_type'], run)\n",
    "weights_path = os.path.join(log_dir, 'weights.h5')\n",
    "if not os.path.exists(weights_path):\n",
    "    raise FileNotFoundError(weights_path)\n",
    "X, X_raw, model, means, scales, wrong_units, args = survey_autoencoder(arg_dict)\n",
    "\n",
    "top_classes = ['RR_Lyrae_FM', 'W_Ursae_Maj', 'Classical_Cepheid', 'Beta_Persei', 'Semireg_PV']\n",
    "full = joblib.load('data/asas/full.pkl')\n",
    "full = [lc for lc in full if lc.label in top_classes]\n",
    "split = [el for lc in full for el in lc.split(args.n_min, args.n_max)]\n",
    "#if args.ss_resid:\n",
    "#    split = [el for el in split if el.ss_resid <= args.ss_resid]\n",
    "if args.period_fold:\n",
    "    for lc in split:\n",
    "        lc.period_fold()\n",
    "X_list = [np.c_[lc.times, lc.measurements, lc.errors] for lc in split]\n",
    "classnames, indices = np.unique([lc.label for lc in split], return_inverse=True)\n",
    "periods = [lc.p for lc in split]\n",
    "y = classnames[indices]\n",
    "train, valid = list(StratifiedKFold(n_splits=5, shuffle=True, random_state=SEED).split(X_list, y))[0]\n",
    "\n",
    "X_raw = pad_sequences(X_list, value=0., dtype='float', padding='post')\n",
    "X, means, scales, wrong_units = preprocess(X_raw, args.m_max)\n",
    "encode_model = Model(input=model.input, output=model.get_layer('encoding').output)\n",
    "encoding = encode_model.predict({'main_input': X, 'aux_input': np.delete(X, 1, axis=2)})\n",
    "encoding = np.c_[encoding, means, scales, periods]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "len(fset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "#asas_model = RandomForestClassifier(oob_score=True)\n",
    "#asas_model, best_params, oob_score = oob_search(asas_model, RF_PARAM_GRID, encoding[train], y[train])\n",
    "asas_model = GridSearchCV(RandomForestClassifier(random_state=0), RF_PARAM_GRID)\n",
    "asas_model.fit(encoding[train], y[train])\n",
    "print(\"Training score: {}%\".format(100 * asas_model.score(encoding[train], y[train])))\n",
    "print(\"Validation score: {}%\".format(100 * asas_model.score(encoding[valid], y[valid])))\n",
    "plot_confusion_matrix(y[valid], asas_model.predict(encoding[valid]), classnames,\n",
    "                      'figures/asas_confusion.pdf' if SAVE_PLOTS else None)"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "from sklearn.metrics import roc_curve, auc\n",
    "from sklearn.preprocessing import label_binarize\n",
    "from sklearn.multiclass import OneVsRestClassifier\n",
    "\n",
    "Y = label_binarize(y, np.unique(y))\n",
    "binarized = OneVsRestClassifier(asas_model.best_estimator_)\n",
    "binarized.fit(encoding[train], Y[train])\n",
    "Y_score = binarized.predict_proba(encoding[valid])\n",
    "fpr = dict()\n",
    "tpr = dict()\n",
    "for i in range(Y.shape[-1]):\n",
    "    fpr[i], tpr[i], _ = roc_curve(Y[valid, i], Y_score[:, i])\n",
    "    plt.plot(fpr[i], tpr[i], label=top_classes[i])\n",
    "plt.legend(loc=\"lower right\")\n",
    "plt.xlabel('False positive rate')\n",
    "plt.ylabel('True positive rate');\n",
    "if SAVE_PLOTS:\n",
    "    plt.savefig('figures/asas_roc.pdf')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## ASAS Richards random forests"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "from sklearn.model_selection import StratifiedKFold\n",
    "from light_curve import LightCurve\n",
    "import joblib\n",
    "from cesium.features import LOMB_SCARGLE_FEATS, CADENCE_FEATS, GENERAL_FEATS\n",
    "from cesium.featurize import featurize_time_series\n",
    "\n",
    "from argparse import Namespace; args = Namespace(n_min=200, n_max=200)\n",
    "\n",
    "top_classes = ['RR_Lyrae_FM', 'W_Ursae_Maj', 'Classical_Cepheid', 'Beta_Persei', 'Semireg_PV']\n",
    "full = joblib.load('data/asas/full.pkl')\n",
    "full = [lc for lc in full if lc.label in top_classes]\n",
    "split = [el for lc in full for el in lc.split(args.n_min, args.n_max)]\n",
    "X_list = [np.c_[lc.times, lc.measurements, lc.errors] for lc in split]\n",
    "classnames, indices = np.unique([lc.label for lc in split], return_inverse=True)\n",
    "y = classnames[indices]\n",
    "train, valid = list(StratifiedKFold(n_splits=5, shuffle=True, random_state=SEED).split(X_list, y))[0]\n",
    "\n",
    "times, measurements, errors = zip(*[x.T for x in X_list])\n",
    "\n",
    "features_to_use = LOMB_SCARGLE_FEATS + GENERAL_FEATS# + CADENCE_FEATS\n",
    "\n",
    "#fset = featurize_time_series(times, measurements, errors, features_to_use)\n",
    "#joblib.dump(fset, 'asas_richards.pkl', compress=True)\n",
    "fset = joblib.load('asas_richards.pkl')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "from sklearn.ensemble import RandomForestClassifier\n",
    "\n",
    "#cs_model = RandomForestClassifier(oob_score=True)\n",
    "#cs_model, best_params, oob_score = oob_search(cs_model, RF_PARAM_GRID, feat_df.iloc[train], y[train])\n",
    "cs_model = GridSearchCV(RandomForestClassifier(random_state=0), RF_PARAM_GRID)\n",
    "#cs_model = deepcopy(asas_model.best_estimator_)\n",
    "cs_model.fit(fset.iloc[train], y[train])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "cs_pred = cs_model.predict(fset)\n",
    "print(f\"Training score: {100 * np.mean(cs_pred[train] == y[train])}%\")\n",
    "print(f\"Validation score: {100 * np.mean(cs_pred[valid] == y[valid])}%\")\n",
    "plot_confusion_matrix(y[valid], cs_model.predict(fset.iloc[valid]), classnames,\n",
    "                      'figures/asas_richards_confusion.pdf' if SAVE_PLOTS else None)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## LINEAR reconstructions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import joblib\n",
    "from sklearn.model_selection import StratifiedKFold, GridSearchCV\n",
    "from sklearn.ensemble import RandomForestClassifier\n",
    "from light_curve import LightCurve\n",
    "from keras.preprocessing.sequence import pad_sequences\n",
    "from keras.models import Model\n",
    "from keras_util import parse_model_args, get_run_id\n",
    "from survey_autoencoder import main as survey_autoencoder\n",
    "\n",
    "arg_dict = {'size': 96, 'embedding': 64, 'num_layers': 2, 'model_type': 'gru',\n",
    "            'sim_type': 'linear/n200', 'n_min': 200, 'n_max': 200,\n",
    "            'lr': 1e-3, 'bidirectional': True, 'drop_frac': 0.25, 'period_fold': False,\n",
    "            'survey_files': ['data/linear/full.pkl'], 'no_train': True}\n",
    "run = get_run_id(**arg_dict)\n",
    "log_dir = os.path.join(os.getcwd(), 'keras_logs', arg_dict['sim_type'], run)\n",
    "weights_path = os.path.join(log_dir, 'weights.h5')\n",
    "if not os.path.exists(weights_path):\n",
    "    raise FileNotFoundError(weights_path)\n",
    "X, X_raw, model, means, scales, wrong_units, args = survey_autoencoder(arg_dict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "(len(X), 0.8 * len(X), 0.2 * len(X))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "pred  = model.predict({'main_input': X[:], 'aux_input': X[:, :, 0:1]})\n",
    "ords = np.argsort(np.mean((pred.squeeze() - X[:, :, 1]) ** 2 / (X_raw[:, :, 2] / scales), axis=1))\n",
    "inds, labels = zip(*[(1142, '25'), (3425, '75')])\n",
    "#for i, label in zip(inds, labels):\n",
    "#    t = X_raw[~wrong_units][i, :, 0]\n",
    "#    m = X_raw[~wrong_units][i, :, 1]\n",
    "#    e = X_raw[~wrong_units][i, :, 2]\n",
    "#    m_est = pred[i] * scales[i, 0] + means[i]\n",
    "#    plt.figure()\n",
    "#    plt.errorbar(t, m, e, None, 'o');\n",
    "#    plt.errorbar(t, m_est, None, None, 'o');\n",
    "#    plt.xlabel('Time [day]')\n",
    "#    plt.ylabel('Magnitude')\n",
    "#    plt.legend(['Original', 'Reconstructed'])\n",
    "#    plt.gca().invert_yaxis()\n",
    "#    if SAVE_PLOTS:\n",
    "#        plt.savefig(f'figures/linear_reconstruct_{label}.pdf')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "import os\n",
    "import joblib\n",
    "from keras_util import get_run_id\n",
    "from survey_autoencoder import main as survey_autoencoder\n",
    "\n",
    "arg_dict = {'size': 96, 'embedding': 64, 'num_layers': 2, 'model_type': 'gru',\n",
    "            'sim_type': 'asas_linear_fold/n200_ss0.7', 'n_min': 200, 'n_max': 200,\n",
    "            'lr': 5e-4, 'bidirectional': True,\n",
    "            'drop_frac': 0.25, 'survey_files': ['data/linear/full.pkl'],\n",
    "            'period_fold': True, 'no_train': True}\n",
    "run = get_run_id(**arg_dict)\n",
    "log_dir = os.path.join(os.getcwd(), 'keras_logs', arg_dict['sim_type'], run)\n",
    "weights_path = os.path.join(log_dir, 'weights.h5')\n",
    "if not os.path.exists(weights_path):\n",
    "    raise FileNotFoundError(weights_path)\n",
    "X, X_raw, model, means, scales, wrong_units, args = survey_autoencoder(arg_dict)\n",
    "pred_fold = model.predict({'main_input': X[:], 'aux_input': X[:, :, [0,]]}, batch_size=2000)\n",
    "ords_fold = np.argsort(np.sum((pred_fold.squeeze() - X[:, :, 1]) ** 2, axis=1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "full = joblib.load('data/linear/full.pkl')\n",
    "split = [el for lc in full for el in lc.split(args.n_min, args.n_max)]\n",
    "for ord_i, label in zip(inds, labels):\n",
    "    i = ords[ord_i]\n",
    "    print(i)\n",
    "    t = X_raw[~wrong_units][i, :, 0]\n",
    "    m = X_raw[~wrong_units][i, :, 1]\n",
    "    e = X_raw[~wrong_units][i, :, 2]\n",
    "    m_est = (pred[i] * scales[i, 0] + means[i])[np.argsort(split[i].times % split[i].p)]\n",
    "    m_fold = pred_fold[i] * scales[i, 0] + means[i]\n",
    "    plt.figure()\n",
    "    plt.errorbar(t, m, e, None, 'o');\n",
    "    plt.errorbar(t, m_est, None, None, 'o', alpha=0.6);\n",
    "    plt.errorbar(t, m_fold, None, None, 'o');\n",
    "    plt.xlabel('Phase [day]')\n",
    "    plt.ylabel('Magnitude')\n",
    "    plt.legend(['Original', 'Reconstructed (non-folded)', 'Reconstructed (folded)'])\n",
    "    plt.title(f'{split[i].survey} {split[i].name} ({split[i].label})')\n",
    "    plt.gca().invert_yaxis()\n",
    "    if SAVE_PLOTS:\n",
    "        plt.savefig(f'figures/linear_folded_reconstruct_{label}.pdf')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## LINEAR autoencoder random forests"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "import os\n",
    "import joblib\n",
    "from sklearn.model_selection import StratifiedKFold, GridSearchCV\n",
    "from sklearn.ensemble import RandomForestClassifier\n",
    "from light_curve import LightCurve\n",
    "from keras.preprocessing.sequence import pad_sequences\n",
    "from keras.models import Model\n",
    "from keras_util import get_run_id\n",
    "from survey_autoencoder import main as survey_autoencoder, preprocess\n",
    "\n",
    "#arg_dict = {'size': 96, 'embedding': 64, 'num_layers': 2, 'model_type': 'gru',\n",
    "#        'sim_type': 'asas_linear_fold/n200_ss0.7', 'n_min': 200, 'n_max': 200,\n",
    "#        'lr': 5e-4, 'bidirectional': True, 'drop_frac': 0.25, 'period_fold': True,\n",
    "#        'survey_files': ['data/linear/full.pkl'], 'no_train': True}\n",
    "arg_dict = args.__dict__\n",
    "run = get_run_id(**arg_dict)\n",
    "log_dir = os.path.join(os.getcwd(), 'keras_logs', arg_dict['sim_type'], run)\n",
    "weights_path = os.path.join(log_dir, 'weights.h5')\n",
    "if not os.path.exists(weights_path):\n",
    "    raise FileNotFoundError(weights_path)\n",
    "X, X_raw, model, means, scales, wrong_units, args = survey_autoencoder(arg_dict)\n",
    "\n",
    "full = joblib.load('data/linear/full.pkl')\n",
    "split = [el for lc in full for el in lc.split(args.n_min, args.n_max)]\n",
    "if args.period_fold:\n",
    "    for lc in split:\n",
    "        lc.period_fold()\n",
    "X_list = [np.c_[lc.times, lc.measurements, lc.errors] for lc in split]\n",
    "classnames, indices = np.unique([lc.label for lc in split], return_inverse=True)\n",
    "periods = [lc.p for lc in split]\n",
    "y = classnames[indices]\n",
    "#train, valid = list(StratifiedKFold(n_splits=5, shuffle=True, random_state=SEED).split(X_list, y))[0]\n",
    "\n",
    "encode_model = Model(input=model.input, output=model.get_layer('encoding').output)\n",
    "encoding = encode_model.predict({'main_input': X, 'aux_input': np.delete(X, 1, axis=2)})\n",
    "encoding = np.c_[encoding, means, scales, periods]\n",
    "#linear_model = GridSearchCV(RandomForestClassifier(random_state=0), RF_PARAM_GRID)\n",
    "#linear_model.fit(encoding[train], y[train])\n",
    "#print(\"Training score: {}%\".format(100 * linear_model.score(encoding[train], y[train])))\n",
    "#print(\"Validation score: {}%\".format(100 * linear_model.score(encoding[valid], y[valid])))\n",
    "###print(f\"OOB score: {100 * oob_score}\")\n",
    "##plot_confusion_matrix(y[valid], cs_model.predict(encoding[valid]), classnames,\n",
    "##                      'figures/linear_confusion.pdf' if SAVE_PLOTS else None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "print(\"LINEAR autoencoder\")\n",
    "for train, valid in StratifiedKFold(n_splits=5, shuffle=True, random_state=SEED).split(X_list, y):\n",
    "    linear_model = GridSearchCV(RandomForestClassifier(random_state=0), RF_PARAM_GRID)\n",
    "    linear_model.fit(encoding[train], y[train])\n",
    "    print(\"Training score: {}%\".format(100 * linear_model.score(encoding[train], y[train])))\n",
    "    print(\"Validation score: {}%\".format(100 * linear_model.score(encoding[valid], y[valid])))"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "from sklearn.metrics import roc_curve, auc\n",
    "from sklearn.preprocessing import label_binarize\n",
    "from sklearn.multiclass import OneVsRestClassifier\n",
    "\n",
    "Y = label_binarize(y, np.unique(y))\n",
    "binarized = OneVsRestClassifier(linear_model.best_estimator_)\n",
    "binarized.fit(encoding[train], Y[train])\n",
    "Y_score = binarized.predict_proba(encoding[valid])\n",
    "fpr = dict()\n",
    "tpr = dict()\n",
    "for i in range(Y.shape[-1]):\n",
    "    fpr[i], tpr[i], _ = roc_curve(Y[valid, i], Y_score[:, i])\n",
    "    plt.plot(fpr[i], tpr[i], label=classnames[i])\n",
    "plt.legend(loc=\"lower right\")\n",
    "plt.xlabel('False positive rate')\n",
    "plt.ylabel('True positive rate');\n",
    "if SAVE_PLOTS:\n",
    "    plt.savefig('figures/linear_roc.pdf')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## LINEAR Richards random forests"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "from sklearn.model_selection import StratifiedKFold\n",
    "from light_curve import LightCurve\n",
    "import joblib\n",
    "from cesium.features import LOMB_SCARGLE_FEATS, CADENCE_FEATS, GENERAL_FEATS\n",
    "from cesium.featurize import featurize_time_series\n",
    "\n",
    "from argparse import Namespace; args = Namespace(n_min=200, n_max=200)\n",
    "\n",
    "full = joblib.load('data/linear/full.pkl')\n",
    "split = [el for lc in full for el in lc.split(args.n_min, args.n_max)]\n",
    "X_list = [np.c_[lc.times, lc.measurements, lc.errors] for lc in split]\n",
    "classnames, indices = np.unique([lc.label for lc in split], return_inverse=True)\n",
    "y = classnames[indices]\n",
    "train, valid = list(StratifiedKFold(n_splits=5, shuffle=True, random_state=SEED).split(X_list, y))[0]\n",
    "\n",
    "times, measurements, errors = zip(*[x.T for x in X_list])\n",
    "\n",
    "features_to_use = LOMB_SCARGLE_FEATS + GENERAL_FEATS# + CADENCE_FEATS\n",
    "\n",
    "#fset = featurize_time_series(times, measurements, errors, features_to_use)\n",
    "#joblib.dump(fset, 'linear_richards.pkl', compress=True)\n",
    "fset = joblib.load('linear_richards.pkl')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "from sklearn.ensemble import RandomForestClassifier\n",
    "\n",
    "cs_model = GridSearchCV(RandomForestClassifier(random_state=0), RF_PARAM_GRID)\n",
    "#cs_model = deepcopy(asas_model.best_estimator_)\n",
    "cs_model.fit(fset.iloc[train], y[train])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "cs_pred = cs_model.predict(fset)\n",
    "print(f\"Training score: {100 * np.mean(cs_pred[train] == y[train])}%\")\n",
    "print(f\"Validation score: {100 * np.mean(cs_pred[valid] == y[valid])}%\")\n",
    "plot_confusion_matrix(y[valid], cs_model.predict(fset.iloc[valid]), classnames,\n",
    "                      'figures/linear_richards_confusion.pdf' if SAVE_PLOTS else None)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## LINEAR noisy examples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "%run survey_autoencoder.py --size 96 --num_layers 2 --drop_frac 0.25 --no_train \\\n",
    "                           --model_type gru --lr 1e-3 --sim_type linear/n200 \\\n",
    "                           --n_min 200 --n_max 200 --embedding 64 \\\n",
    "                           --survey_files data/linear/full.pkl --bidirectional --period_fold\n",
    "            \n",
    "full = joblib.load('data/linear/full.pkl')\n",
    "split = [el for lc in full for el in lc.split(args.n_min, args.n_max)]\n",
    "if args.period_fold:\n",
    "    for lc in split:\n",
    "        lc.period_fold()\n",
    "\n",
    "#lc = [lc for lc in split if lc.name == '3182815'][0]\n",
    "lc = [lc for lc in split if lc.name == '1018967'][0]\n",
    "\n",
    "X = np.c_[lc.times, lc.measurements, lc.errors][np.newaxis, :]\n",
    "errors = X[:, :, 2]\n",
    "sample_gen = ku.noisify_samples({'main_input': X, 'aux_input': np.delete(X, 1, axis=2)},\n",
    "                                X[:, :, [1]], errors, batch_size=1, sample_weight=None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "colors = sns.color_palette('deep', n_colors=2)\n",
    "t = X_raw[0, :, 0]\n",
    "m = X[0, :, 1]\n",
    "e = errors[0]\n",
    "plt.figure()\n",
    "sns.set_style(\"darkgrid\")\n",
    "plt.errorbar(t, m, e, None, 'o', color=colors[0]);\n",
    "for i in range(3):\n",
    "    (X_i, _), y_i, w_i = next(sample_gen)\n",
    "    plt.plot(t, X[:, :, 1].mean() + X[:, :, 1].std() * X_i[0, :, 1], 'o', alpha=0.5, color=colors[1])\n",
    "    if i == 0:\n",
    "        plt.legend(['Sampled', 'Original'])\n",
    "plt.xlabel('Time [day]')\n",
    "plt.ylabel('Magnitude')\n",
    "plt.gca().invert_yaxis()\n",
    "plt.title(f'{lc.survey} {lc.name} ({lc.label})')\n",
    "\n",
    "if 1:\n",
    "    plt.savefig('paper/figures/linear_noisy.pdf')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Transfer learning"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "# LINEAR autoencoder -> ASAS RF -> ASAS\n",
    "import os\n",
    "import joblib\n",
    "from sklearn.model_selection import StratifiedKFold, GridSearchCV\n",
    "from sklearn.ensemble import RandomForestClassifier\n",
    "from light_curve import LightCurve\n",
    "from keras.preprocessing.sequence import pad_sequences\n",
    "from keras.models import Model\n",
    "from keras_util import get_run_id\n",
    "from survey_autoencoder import main as survey_autoencoder, preprocess\n",
    "\n",
    "arg_dict = {'size': 96, 'embedding': 64, 'num_layers': 2, 'model_type': 'gru',\n",
    "            'sim_type': 'linear_fold/n200', 'n_min': 200, 'n_max': 200,\n",
    "            'lr': 1e-3, 'bidirectional': True,# 'ss_resid': 0.7,\n",
    "            'drop_frac': 0.25, 'survey_files': ['data/linear/full.pkl'],\n",
    "            'period_fold': True, 'no_train': True}\n",
    "run = get_run_id(**arg_dict)\n",
    "log_dir = os.path.join(os.getcwd(), 'keras_logs', arg_dict['sim_type'], run)\n",
    "weights_path = os.path.join(log_dir, 'weights.h5')\n",
    "if not os.path.exists(weights_path):\n",
    "    raise FileNotFoundError(weights_path)\n",
    "X, X_raw, model, means, scales, wrong_units, args = survey_autoencoder(arg_dict)\n",
    "\n",
    "top_classes = ['RR_Lyrae_FM', 'W_Ursae_Maj', 'Classical_Cepheid', 'Beta_Persei', 'Semireg_PV']\n",
    "full = joblib.load('data/asas/full.pkl')\n",
    "full = [lc for lc in full if lc.label in top_classes]\n",
    "split = [el for lc in full for el in lc.split(args.n_min, args.n_max)]\n",
    "#if args.ss_resid:\n",
    "#    split = [el for el in split if el.ss_resid <= args.ss_resid]\n",
    "if args.period_fold:\n",
    "    for lc in split:\n",
    "        lc.period_fold()\n",
    "X_list = [np.c_[lc.times, lc.measurements, lc.errors] for lc in split]\n",
    "classnames, indices = np.unique([lc.label for lc in split], return_inverse=True)\n",
    "periods = [lc.p for lc in split]\n",
    "y = classnames[indices]\n",
    "train, valid = list(StratifiedKFold(n_splits=5, shuffle=True, random_state=SEED).split(X_list, y))[0]\n",
    "\n",
    "X_raw = pad_sequences(X_list, value=0., dtype='float', padding='post')\n",
    "X, means, scales, wrong_units = preprocess(X_raw, args.m_max)\n",
    "encode_model = Model(input=model.input, output=model.get_layer('encoding').output)\n",
    "encoding = encode_model.predict({'main_input': X, 'aux_input': np.delete(X, 1, axis=2)})\n",
    "encoding = np.c_[encoding, means, scales, periods]\n",
    "\n",
    "asas_model = GridSearchCV(RandomForestClassifier(random_state=0), RF_PARAM_GRID)\n",
    "asas_model.fit(encoding[train], y[train])\n",
    "print(\"Training score: {}%\".format(100 * asas_model.score(encoding[train], y[train])))\n",
    "print(\"Validation score: {}%\".format(100 * asas_model.score(encoding[valid], y[valid])))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "# ASAS autoencoder -> LINEAR RF -> LINEAR\n",
    "import os\n",
    "import joblib\n",
    "from sklearn.model_selection import StratifiedKFold, GridSearchCV\n",
    "from sklearn.ensemble import RandomForestClassifier\n",
    "from light_curve import LightCurve\n",
    "from keras.preprocessing.sequence import pad_sequences\n",
    "from keras.models import Model\n",
    "from keras_util import get_run_id\n",
    "from survey_autoencoder import main as survey_autoencoder, preprocess\n",
    "\n",
    "arg_dict = {'size': 64, 'embedding': 64, 'num_layers': 1, 'model_type': 'gru',\n",
    "            'sim_type': 'asas_fold/n200_ss0.7', 'lomb_score': None, 'n_min': 200, 'n_max': 200,\n",
    "            'lr': 1e-3, 'bidirectional': True, 'drop_frac': 0.25, 'period_fold': True,\n",
    "            'ss_resid': 0.7, 'survey_files': ['data/asas/full.pkl'], 'no_train': True}\n",
    "run = get_run_id(**arg_dict)\n",
    "log_dir = os.path.join(os.getcwd(), 'keras_logs', arg_dict['sim_type'], run)\n",
    "weights_path = os.path.join(log_dir, 'weights.h5')\n",
    "if not os.path.exists(weights_path):\n",
    "    raise FileNotFoundError(weights_path)\n",
    "X, X_raw, model, means, scales, wrong_units, args = survey_autoencoder(arg_dict)\n",
    "\n",
    "full = joblib.load('data/linear/full.pkl')\n",
    "split = [el for lc in full for el in lc.split(args.n_min, args.n_max)]\n",
    "if args.period_fold:\n",
    "    for lc in split:\n",
    "        lc.period_fold()\n",
    "X_list = [np.c_[lc.times, lc.measurements, lc.errors] for lc in split]\n",
    "classnames, indices = np.unique([lc.label for lc in split], return_inverse=True)\n",
    "periods = [lc.p for lc in split]\n",
    "y = classnames[indices]\n",
    "train, valid = list(StratifiedKFold(n_splits=5, shuffle=True, random_state=SEED).split(X_list, y))[0]\n",
    "\n",
    "X_raw = pad_sequences(X_list, value=0., dtype='float', padding='post')\n",
    "X, means, scales, wrong_units = preprocess(X_raw, args.m_max)\n",
    "encode_model = Model(input=model.input, output=model.get_layer('encoding').output)\n",
    "encoding = encode_model.predict({'main_input': X, 'aux_input': np.delete(X, 1, axis=2)})\n",
    "encoding = np.c_[encoding, means, scales, periods]\n",
    "linear_model = GridSearchCV(RandomForestClassifier(random_state=0), RF_PARAM_GRID)\n",
    "linear_model.fit(encoding[train], y[train])\n",
    "print(\"Training score: {}%\".format(100 * linear_model.score(encoding[train], y[train])))\n",
    "print(\"Validation score: {}%\".format(100 * linear_model.score(encoding[valid], y[valid])))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#%%time\n",
    "import os\n",
    "import joblib\n",
    "from sklearn.model_selection import StratifiedKFold, GridSearchCV\n",
    "from sklearn.ensemble import RandomForestClassifier\n",
    "from light_curve import LightCurve\n",
    "from keras.preprocessing.sequence import pad_sequences\n",
    "from keras.models import Model\n",
    "from keras_util import get_run_id\n",
    "from survey_autoencoder import main as survey_autoencoder, preprocess\n",
    "\n",
    "arg_dict = {'size': 96, 'embedding': 64, 'num_layers': 2, 'model_type': 'gru',\n",
    "            'lomb_score': None, 'n_min': 200, 'n_max': 200,\n",
    "            'lr': 5e-4, 'bidirectional': True, 'drop_frac': 0.25, 'period_fold': True,\n",
    "            'no_train': True}\n",
    "\n",
    "def load_and_encode(arg_dict):\n",
    "    run = get_run_id(**arg_dict)\n",
    "    log_dir = os.path.join(os.getcwd(), 'keras_logs', arg_dict['sim_type'], run)\n",
    "    weights_path = os.path.join(log_dir, 'weights.h5')\n",
    "    if not os.path.exists(weights_path):\n",
    "        raise FileNotFoundError(weights_path)\n",
    "    _, _, model, _, _, _, args = survey_autoencoder(arg_dict)\n",
    "\n",
    "    full = joblib.load(arg_dict['survey_files'][0])\n",
    "    if 'asas' in arg_dict['survey_files'][0]:\n",
    "        top_classes = ['RR_Lyrae_FM', 'W_Ursae_Maj', 'Classical_Cepheid', 'Beta_Persei', 'Semireg_PV']\n",
    "        full = [lc for lc in full if lc.label in top_classes]\n",
    "    split = [el for lc in full for el in lc.split(args.n_min, args.n_max)]\n",
    "    if args.period_fold:\n",
    "        for lc in split:\n",
    "            lc.period_fold()\n",
    "    X_list = [np.c_[lc.times, lc.measurements, lc.errors] for lc in split]\n",
    "    classnames, indices = np.unique([lc.label for lc in split], return_inverse=True)\n",
    "    periods = [lc.p for lc in split]\n",
    "    y = classnames[indices]\n",
    "    train, valid = list(StratifiedKFold(n_splits=5, shuffle=True, random_state=SEED).split(X_list, y))[0]\n",
    "\n",
    "    X_raw = pad_sequences(X_list, value=0., dtype='float', padding='post')\n",
    "    X, means, scales, wrong_units = preprocess(X_raw, args.m_max)\n",
    "    encode_model = Model(input=model.input, output=model.get_layer('encoding').output)\n",
    "    encoding = encode_model.predict({'main_input': X, 'aux_input': np.delete(X, 1, axis=2)})\n",
    "    encoding = np.c_[encoding, means, scales, periods]\n",
    "    \n",
    "    return encoding, y\n",
    "\n",
    "asas_encoding_asas, y_asas = load_and_encode({**arg_dict, 'sim_type': 'asas_fold/n200_ss0.7',\n",
    "                                              'survey_files': ['data/asas/full.pkl']})\n",
    "asas_encoding_linear, y_linear = load_and_encode({**arg_dict, 'sim_type': 'asas_fold/n200_ss0.7',\n",
    "                                                  'survey_files': ['data/linear/full.pkl']})\n",
    "linear_encoding_asas, _ = load_and_encode({**arg_dict, 'sim_type': 'linear_fold/n200',\n",
    "                                           'survey_files': ['data/asas/full.pkl']})\n",
    "linear_encoding_linear, _ = load_and_encode({**arg_dict, 'sim_type': 'linear_fold/n200',\n",
    "                                             'survey_files': ['data/linear/full.pkl']})\n",
    "\n",
    "fset_asas = joblib.load('asas_richards.pkl')\n",
    "fset_linear = joblib.load('linear_richards.pkl')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "splits_asas = StratifiedKFold(n_splits=5, shuffle=True, random_state=SEED).split(y_asas, y_asas)\n",
    "splits_linear = StratifiedKFold(n_splits=5, shuffle=True, random_state=SEED).split(y_linear, y_linear)\n",
    "\n",
    "feature_dict = {\n",
    "    'asas_auto': (asas_encoding_asas, asas_encoding_linear),\n",
    "    'linear_auto': (linear_encoding_asas, linear_encoding_linear),\n",
    "    'richards': (fset_asas, fset_linear)\n",
    "}\n",
    "\n",
    "from collections import defaultdict\n",
    "asas_scores = defaultdict(list)    # key[0] features, trained on key[1], scored on y_asas\n",
    "linear_scores = defaultdict(list)  # key[0] features, trained on key[1], scored on y_linear\n",
    "for label, (X_asas, X_linear) in feature_dict.items():\n",
    "    for train, valid in splits_asas:\n",
    "        import datetime; print(label, datetime.datetime.now())\n",
    "        asas_model = GridSearchCV(RandomForestClassifier(random_state=0), RF_PARAM_GRID)\n",
    "        asas_model.fit(X_asas[train], y_asas[train])\n",
    "        asas_scores[(label, 'asas')].append(asas_model.score(X_asas[valid], y_asas[valid]))\n",
    "\n",
    "    for train, valid in splits_linear:\n",
    "        import datetime; print(label, datetime.datetime.now())\n",
    "        linear_model = GridSearchCV(RandomForestClassifier(random_state=0), RF_PARAM_GRID)\n",
    "        linear_model.fit(X_linear[train], y_linear[train])\n",
    "        linear_scores[(label, 'linear')].append(linear_model.score(X_linear[valid], y_linear[valid]))\n",
    "        \n",
    "    asas_model.fit(X_asas, y_asas)\n",
    "    linear_scores[(label, 'asas')] = [asas_model.score(X_linear, y_linear)]\n",
    "\n",
    "    linear_model.fit(X_linear, y_linear)\n",
    "    asas_scores[(label, 'linear')] = [linear_model.score(X_asas, y_asas)]\n",
    "    break  # TODO temporary"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "asas_scores"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dropout plots "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "log_files = glob.glob('keras_logs/linear_fold_valid/n200/gru_096_x2_5m04*emb64_bidir/training.csv')\n",
    "\n",
    "step_logs, time_logs = parse_logs(log_files)\n",
    "step_plot = training_plot(step_logs, ylim=(0.5, 10.0))\n",
    "step_plot.axes[0].legend(labels=[f\"Dropout={d}%\" for d in [0, 12.5, 25, 37.5]]);\n",
    "time_plot = training_plot(time_logs, ylim=(0.5, 10.0))\n",
    "time_plot.axes[0].legend(labels=[f\"Dropout={d}%\" for d in [0, 12.5, 25, 37.5]]);\n",
    "if 1:\n",
    "    step_plot.savefig('paper/figures/linear_dropout_step.pdf')\n",
    "    time_plot.savefig('paper/figures/linear_dropout_time.pdf')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Variable length plots"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "log_files = [f'keras_logs/linear_fold_valid/n{i}/gru_096_x2_5m04_drop25_emb64_bidir/training.csv'\n",
    "             for i in ['200', '150_250', '100_300']]\n",
    "\n",
    "step_logs, time_logs = parse_logs(log_files, 2000)\n",
    "step_plot = training_plot(step_logs, ylim=(0.5, 15.0))\n",
    "step_plot.axes[0].legend(labels=[\"N=200\", \"150≤N≤250\", \"100≤N≤300\"])\n",
    "time_plot = training_plot(time_logs, ylim=(0.5, 15.0))\n",
    "time_plot.axes[0].legend(labels=[\"N=200\", \"150≤N≤250\", \"100≤N≤300\"])\n",
    "if 1:\n",
    "    step_plot.savefig('paper/figures/linear_var_length_step.pdf')\n",
    "    time_plot.savefig('paper/figures/linear_var_length_time.pdf')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "df = pd.concat((pd.DataFrame({'a': [1.0, 2.0, 3.0]}), pd.DataFrame({'a': [4.0, 5.0, 6.0]})),\n",
    "               axis=1)\n",
    "print(\"Before:\\n\", df)\n",
    "c = df.iloc[:, 0]\n",
    "c.dropna(inplace=True)\n",
    "print(\"After:\\n\", df)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## MACHO reconstructions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "import os\n",
    "from light_curve import LightCurve\n",
    "import joblib\n",
    "from keras_util import get_run_id\n",
    "from survey_autoencoder import main as survey_autoencoder\n",
    "\n",
    "arg_dict = {'size': 96, 'embedding': 64, 'num_layers': 2, 'model_type': 'gru',\n",
    "            'sim_type': 'macho_fold/n200_ss0.7_long', 'n_min': 200, 'n_max': 200,\n",
    "            'lr': 5e-4, 'bidirectional': True, 'ss_resid': 0.7,\n",
    "            'drop_frac': 0.25, 'survey_files': ['data/macho/full.pkl'],\n",
    "            'period_fold': True, 'no_train': True}\n",
    "\n",
    "run = get_run_id(**arg_dict)\n",
    "log_dir = os.path.join(os.getcwd(), 'keras_logs', arg_dict['sim_type'], run)\n",
    "weights_path = os.path.join(log_dir, 'weights.h5')\n",
    "if not os.path.exists(weights_path):\n",
    "    raise FileNotFoundError(weights_path)\n",
    "X, X_raw, model, means, scales, wrong_units, args = survey_autoencoder(arg_dict)\n",
    "full = joblib.load('data/macho/full.pkl')\n",
    "split = [el for lc in full for el in lc.split(args.n_min, args.n_max)]\n",
    "#if args.ss_resid:\n",
    "#    split = [el for el in split if el.ss_resid <= args.ss_resid]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X.shape[0] * np.array([0.25, 0.75])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "pred = model.predict({'main_input': X[:], 'aux_input': X[:, :, [0,]]})\n",
    "ords = np.argsort(np.sum((pred.squeeze() - X[:, :, 1]) ** 2, axis=1))\n",
    "inds, labels = zip(*[(5195, '25'), (15585, '75')])\n",
    "for ord_i, label in zip(inds, labels):\n",
    "    i = ords[ord_i]\n",
    "    print(i)\n",
    "    t = X_raw[~wrong_units][i, :, 0]\n",
    "    m = X_raw[~wrong_units][i, :, 1]\n",
    "    e = X_raw[~wrong_units][i, :, 2]\n",
    "    m_est = pred[i] * scales[i, 0] + means[i]\n",
    "    plt.figure()\n",
    "    plt.errorbar(t, m, e, None, 'o');\n",
    "    plt.errorbar(t, m_est, None, None, 'o');\n",
    "    plt.xlabel('Time [day]')\n",
    "    plt.ylabel('Magnitude')\n",
    "    plt.legend(['Original', 'Reconstructed'])\n",
    "    plt.title(f'{split[i].survey} {split[i].name} ({split[i].label})')\n",
    "    plt.gca().invert_yaxis()\n",
    "    if SAVE_PLOTS:\n",
    "        plt.savefig(f'paper/figures/macho_reconstruct_{label}.pdf')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "from sklearn.model_selection import StratifiedKFold, GridSearchCV\n",
    "from keras.preprocessing.sequence import pad_sequences\n",
    "from survey_autoencoder import main as survey_autoencoder, preprocess\n",
    "from keras.models import Model\n",
    "from sklearn.ensemble import RandomForestClassifier\n",
    "\n",
    "for lc in split:\n",
    "    lc.period_fold()\n",
    "X_list = [np.c_[lc.times, lc.measurements, lc.errors] for lc in split]\n",
    "classnames, indices = np.unique([lc.label for lc in split], return_inverse=True)\n",
    "periods = [lc.p for lc in split]\n",
    "y = classnames[indices]\n",
    "train, valid = list(StratifiedKFold(n_splits=5, shuffle=True, random_state=SEED).split(X_list, y))[0]\n",
    "\n",
    "X_raw = pad_sequences(X_list, value=0., dtype='float', padding='post')\n",
    "X, means, scales, wrong_units = preprocess(X_raw, args.m_max)\n",
    "encode_model = Model(input=model.input, output=model.get_layer('encoding').output)\n",
    "encoding = encode_model.predict({'main_input': X, 'aux_input': np.delete(X, 1, axis=2)})\n",
    "encoding = np.c_[encoding, means, scales, periods]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "(len(full), len(split))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "encoding.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "train.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#from distributed import Client\n",
    "import dask.multiprocessing\n",
    "import dask_searchcv as dcv\n",
    "\n",
    "train = np.arange(500)\n",
    "\n",
    "macho_model = dcv.GridSearchCV(RandomForestClassifier(random_state=0), RF_PARAM_GRID,\n",
    "                               scheduler=dask.multiprocessing.get)\n",
    "%time macho_model.fit(encoding[train], y[train])"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}