clstm_lorenz_demo.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# cLSTM Lorenz-96 Demo\n",
    "- In this notebook, we train a cLSTM model on data simulated from a Lorenz-96 system."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import torch\n",
    "from models.clstm import cLSTM, train_model_gista\n",
    "from synthetic import simulate_lorenz_96\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# For GPU acceleration\n",
    "device = torch.device('cuda')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Simulate data\n",
    "p = 10\n",
    "X_np, GC = simulate_lorenz_96(p, T=1000)\n",
    "X = torch.tensor(X_np[np.newaxis], dtype=torch.float32, device=device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot data\n",
    "fig, axarr = plt.subplots(1, 2, figsize=(16, 5))\n",
    "axarr[0].plot(X_np)\n",
    "axarr[1].plot(X_np[:100])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set up model\n",
    "clstm = cLSTM(p, hidden=10).cuda(device=device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Train with GISTA\n",
    "check_every = 100\n",
    "train_loss_list, train_mse_list = train_model_gista(\n",
    "    clstm, X, lam=6.6, lam_ridge=1e-4, lr=0.005, max_iter=20000, check_every=check_every, truncation=5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Loss function plot\n",
    "fig, axarr = plt.subplots(1, 2, figsize=(16, 5))\n",
    "\n",
    "axarr[0].plot(check_every * np.arange(len(train_loss_list)), train_loss_list)\n",
    "axarr[0].set_title('Train loss')\n",
    "\n",
    "axarr[1].plot(check_every * np.arange(len(train_mse_list)), train_mse_list)\n",
    "axarr[1].set_title('Train MSE')\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Verify learned Granger causality\n",
    "GC_est = clstm.GC().cpu().data.numpy()\n",
    "\n",
    "print('True variable usage = %.2f%%' % (100 * np.mean(GC)))\n",
    "print('Estimated variable usage = %.2f%%' % (100 * np.mean(GC_est)))\n",
    "print('Accuracy = %.2f%%' % (100 * np.mean(GC == GC_est)))\n",
    "\n",
    "# Make figures\n",
    "fig, axarr = plt.subplots(1, 2, figsize=(10, 5))\n",
    "axarr[0].imshow(GC, cmap='Blues')\n",
    "axarr[0].set_title('GC actual')\n",
    "axarr[0].set_ylabel('Affected series')\n",
    "axarr[0].set_xlabel('Causal series')\n",
    "axarr[0].set_xticks([])\n",
    "axarr[0].set_yticks([])\n",
    "\n",
    "axarr[1].imshow(GC_est, cmap='Blues', vmin=0, vmax=1, extent=(0, p, p, 0))\n",
    "axarr[1].set_ylabel('Affected series')\n",
    "axarr[1].set_xlabel('Causal series')\n",
    "axarr[1].set_xticks([])\n",
    "axarr[1].set_yticks([])\n",
    "\n",
    "# Mark disagreements\n",
    "for i in range(p):\n",
    "    for j in range(p):\n",
    "        if GC[i, j] != GC_est[i, j]:\n",
    "            rect = plt.Rectangle((j, i-0.05), 1, 1, facecolor='none', edgecolor='red', linewidth=1)\n",
    "            axarr[1].add_patch(rect)\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}