analysis_2d.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "41481747-0049-4511-b512-df8489041673",
   "metadata": {},
   "outputs": [],
   "source": [
    "import argparse\n",
    "from collections import OrderedDict\n",
    "import datetime\n",
    "get_ipython().run_line_magic('matplotlib', 'inline')\n",
    "import matplotlib.pylab as plt\n",
    "from numbers import Number\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import gc\n",
    "pd.options.display.max_rows = 1500\n",
    "pd.options.display.max_columns = 200\n",
    "pd.options.display.width = 1000\n",
    "pd.set_option('max_colwidth', 400)\n",
    "import pdb\n",
    "import pickle\n",
    "import pprint as pp\n",
    "import time\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from torch import optim\n",
    "from torch.utils.data import DataLoader\n",
    "from deepsnap.batch import Batch as deepsnap_Batch\n",
    "\n",
    "import sys, os\n",
    "sys.path.append(os.path.join(os.path.dirname(\"__file__\"), '..'))\n",
    "sys.path.append(os.path.join(os.path.dirname(\"__file__\"), '..', '..'))\n",
    "from le_pde.argparser import arg_parse\n",
    "from le_pde.datasets.load_dataset import load_data\n",
    "from le_pde.models import load_model\n",
    "from le_pde.pytorch_net.util import groupby_add_keys, filter_df, get_unique_keys_df, Attr_Dict, Printer, get_num_params, get_machine_name, pload, pdump, to_np_array, get_pdict, reshape_weight_to_matrix, ddeepcopy as deepcopy, plot_vectors, record_data, filter_filename, Early_Stopping, str2bool, get_filename_short, print_banner, plot_matrices, get_num_params, init_args, filter_kwargs, to_string, COLOR_LIST\n",
    "from le_pde.utils import update_legacy_default_hyperparam, EXP_PATH\n",
    "from le_pde.utils import deepsnap_to_pyg, LpLoss, to_cpu, to_tuple_shape, parse_multi_step, loss_op, get_device\n",
    "\n",
    "device = torch.device(\"cuda:0\")\n",
    "p = Printer()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "34edf99f-9265-45b0-906d-650e4261dc77",
   "metadata": {},
   "source": [
    "## 0. Functions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "67193e7b-5777-4aac-9426-f91e6837d848",
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_learning_curve(data_record):\n",
    "    x_axis = np.arange(len(data_record[\"train_loss\"]))\n",
    "    plt.figure(figsize=(16,6))\n",
    "    plt.subplot(1,2,1)\n",
    "    plt.plot(x_axis, data_record[\"train_loss\"], label=\"train\")\n",
    "    plt.plot(x_axis, data_record[\"val_loss\"], label=\"val\")\n",
    "    plt.plot(x_axis, data_record[\"test_loss\"], label=\"test\")\n",
    "    plt.legend()\n",
    "    plt.subplot(1,2,2)\n",
    "    plt.semilogy(x_axis, data_record[\"train_loss\"], label=\"train\")\n",
    "    plt.semilogy(x_axis, data_record[\"val_loss\"], label=\"val\")\n",
    "    plt.semilogy(x_axis, data_record[\"test_loss\"], label=\"test\")\n",
    "    plt.legend()\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "5c90d44c-3295-4a89-8c95-d3c5d2bf892e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_results_2d(all_hash, dirname, mode=-1, suffix=\"\"):\n",
    "    all_dict = {}\n",
    "    df_dict_list = []\n",
    "    isplot = True\n",
    "    for hash_str in all_hash:\n",
    "        # Load model:\n",
    "        dirname = EXP_PATH + dirname\n",
    "        df_dict = {}\n",
    "\n",
    "        filename = filter_filename(dirname, include=hash_str)\n",
    "        if len(filename) == 0:\n",
    "            dirname = dirname = EXP_PATH + \"tailin-2d_2022-7-27/\"\n",
    "            filename = filter_filename(dirname, include=hash_str)\n",
    "            if len(filename) == 0:\n",
    "                print(f\"hash {hash_str} does not exist!\")\n",
    "                continue\n",
    "\n",
    "        try:\n",
    "            data_record = pload(dirname + filename[0])\n",
    "        except Exception as e:\n",
    "            print(f\"error {e}\")\n",
    "            continue\n",
    "        if \"train_loss\" not in data_record:\n",
    "            continue\n",
    "        if isplot:\n",
    "            plot_learning_curve(data_record)\n",
    "        args = init_args(update_legacy_default_hyperparam(data_record[\"args\"]))\n",
    "        args.filename = filename\n",
    "        if mode == \"best\":\n",
    "            model = load_model(data_record[\"best_model_dict\"], device=device)\n",
    "            print(\"Load the model with best validation loss.\")\n",
    "        else:\n",
    "            assert isinstance(mode, int)\n",
    "            print(f'Load the model at epoch {data_record[\"epoch\"][mode]}')\n",
    "            model = load_model(data_record[\"model_dict\"][mode], device=device)\n",
    "        model.eval()\n",
    "        p.print(f\"Hash {hash_str}, best model at epoch {data_record['best_epoch']}:\", banner_size=100)\n",
    "\n",
    "        # Load test dataset:\n",
    "        args_test = deepcopy(args)\n",
    "        if args.temporal_bundle_steps == 1:\n",
    "            if args.dataset in [\"fno\", \"fno-2\", \"fno-3\"]:\n",
    "                args_test.multi_step = \"20\"\n",
    "            elif args.dataset in [\"fno-1\"]:\n",
    "                args_test.multi_step = \"40\"\n",
    "            elif args.dataset in [\"fno-4\"]:\n",
    "                args_test.multi_step = \"10\"\n",
    "            else:\n",
    "                raise\n",
    "        else:\n",
    "            pass\n",
    "        args_test.batch_size = 20\n",
    "        (dataset_train_val, dataset_test), (train_loader, val_loader, test_loader) = load_data(args_test)\n",
    "        myloss = LpLoss(size_average=False)\n",
    "\n",
    "        # Compute loss:\n",
    "        loss_list = []\n",
    "        pred_list = []\n",
    "        y_list = []\n",
    "        for data in dataset_test:\n",
    "            data.to(device)\n",
    "            preds, info = model(\n",
    "                data,\n",
    "                pred_steps=np.arange(1,int(args_test.multi_step)+1),\n",
    "                latent_pred_steps=None,\n",
    "                is_recons=False,\n",
    "                use_grads=False,\n",
    "                is_y_diff=False,\n",
    "                is_rollout=False,\n",
    "                use_pos=args.use_pos,\n",
    "            )\n",
    "            pred_reshape = preds[\"n0\"].reshape(1,-1)\n",
    "            y_reshape = data.node_label[\"n0\"].reshape(1,-1)\n",
    "            pred_list.append(preds[\"n0\"].detach())\n",
    "            y_list.append(data.node_label[\"n0\"].detach())\n",
    "            loss_ele = myloss(pred_reshape, y_reshape)\n",
    "            loss_list.append(loss_ele.item())\n",
    "        loss_mean = np.mean(loss_list)\n",
    "        pred_list = torch.stack(pred_list).squeeze(-1)\n",
    "        y_list = torch.stack(y_list).squeeze(-1)\n",
    "        all_dict[hash_str] = (data_record['best_epoch'], loss_mean, len(data_record[\"train_loss\"]), args.epochs)\n",
    "        print(all_dict[hash_str])\n",
    "        print(\"Test for {} is: {:.6e} at epoch {}\".format(hash_str, loss_mean, data_record['best_epoch']))\n",
    "\n",
    "        # df_dict:\n",
    "        df_dict[\"loss\"] = loss_mean\n",
    "        df_dict.update(args.__dict__)\n",
    "        df_dict[\"hash_str\"] = hash_str\n",
    "        df_dict[\"train_loss\"] = data_record[\"train_loss\"][-1]\n",
    "        df_dict[\"val_loss\"] = data_record[\"val_loss\"][-1]\n",
    "        df_dict[\"test_loss\"] = data_record[\"test_loss\"][-1]\n",
    "        df_dict[\"epoch\"] = data_record[\"epoch\"][-1]\n",
    "        df_dict_list.append(df_dict)\n",
    "\n",
    "        # Plotting:\n",
    "        pred_reshape = preds[\"n0\"].permute(1,0,2).squeeze(-1)\n",
    "        y_reshape = data.node_label[\"n0\"].permute(1,0,2).squeeze(-1)\n",
    "        loss_ele = myloss(pred_reshape, y_reshape)\n",
    "        num = 10\n",
    "        print(\"loss_cumu:\")\n",
    "        mse_full = nn.MSELoss(reduction=\"none\")(pred_list, y_list)\n",
    "        mse_time = to_np_array(mse_full.mean((0,1)))\n",
    "        plt.figure(figsize=(12,5))\n",
    "        plt.subplot(1,2,1)\n",
    "        plt.plot(mse_time)\n",
    "        plt.xlabel(\"rollout step\")\n",
    "        plt.ylabel(\"MSE\")\n",
    "        plt.title(\"MSE vs. rollout step (linear scale)\")\n",
    "        plt.subplot(1,2,2)\n",
    "        plt.semilogy(mse_time)\n",
    "        plt.xlabel(\"rollout step\")\n",
    "        plt.ylabel(\"MSE\")\n",
    "        plt.title(\"MSE vs. rollout step (log scale)\")\n",
    "        plt.show()\n",
    "\n",
    "        print(\"individual:\")\n",
    "        for i in range(len(pred_reshape) // num):\n",
    "            print(\"pred:\")\n",
    "            plot_matrices(pred_reshape[num*i:num*(i+1)].reshape(num, 64, 64), images_per_row=num, scale_limit=\"auto\")\n",
    "            print(\"y:\")\n",
    "            plot_matrices(y_reshape[num*i:num*(i+1)].reshape(num, 64, 64), images_per_row=num, scale_limit=\"auto\")\n",
    "            print(\"diff:\")\n",
    "            plot_matrices((pred_reshape - y_reshape)[num*i:num*(i+1)].reshape(num, 64, 64), images_per_row=num, scale_limit=\"auto\")\n",
    "    pdump(all_dict, f\"all_dict_2d{suffix}.p\")\n",
    "    df = pd.DataFrame(df_dict_list)\n",
    "    # dff = filter_df(df, {\"decoder_type\": \"neural-basis-mulr\", \"decoder_act_name\": \"gelu\"})\n",
    "    # dff[[\"hash_str\", \"epoch\", \"loss\"]]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ed718678-b0cc-4895-ac2e-d87defe82ed0",
   "metadata": {
    "tags": []
   },
   "source": [
    "## 1. Analysis:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "29c7b349-c958-4e7e-b38f-01dbd01a30c3",
   "metadata": {},
   "outputs": [],
   "source": [
    "EXP_PATH = \"/dfs/project/plasma/results/\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2220eb8c-6f30-44fd-a911-23ef144d0ba3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# all_hash is a list of hashes, each of which corresponds to one experiment.\n",
    "# For example, if one experiment is saved under ./results/tailin-1d_2022-5-14/fno-4_train_-1_algo_contrast_ebm_False_ebmt_cd_enc_cnn-s_evo_cnn_act_elu_hid_256_lo_mse_recef_1.0_conef_1.0_nconv_4_nlat_1_clat_3_lf_True_reg_None_id_0_Hash_Lvdn+dA9_turing2.p\n",
    "# Then, the \"Lvdn+dA9_turing2\" (located at the end of the filename) is the {hash}_{machine-name} of this file.\n",
    "# The \"tailin_2022-5-8\" is the \"{--exp_id}_{--date_time}\" of the training command.\n",
    "# all_hash can contain multiple hashes, and analyze them sequentially.\n",
    "all_hash = [\"Lvdn+dA9\"]\n",
    "get_results_2d(all_hash, dirname=\"tailin_2022-5-8/\", mode=-1, suffix=\"\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4b4494e7-afaf-4d55-a1fc-b32168ca62b5",
   "metadata": {},
   "source": [
    "## 2. Timing:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "74c7348a-9562-4cc5-8b16-892bd1869427",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "isplot = True\n",
    "all_hash = [\n",
    "    \"nG2NBRJH_turing2\", # latent_size: 512\n",
    "    \"Lvdn+dA9_turing2\", # latent_size: 256\n",
    "    \"ECshE1x1_turing2\", # latent_size: 128\n",
    "    \"2o9axpaz_turing2\", # latent_size: 64\n",
    "    \"QLkuuTf9_turing2\", # latent_size: 32\n",
    "    \"hXzEAQeu_turing2\", # latent_size: 16\n",
    "    \"YkILg8wB_turing2\", # latent_size: 8\n",
    "    \"ck3kAARM_turing2\", # latent_size: 4\n",
    "]\n",
    "hash_str = all_hash[0]\n",
    "dirname = EXP_PATH + \"tailin_2022-5-8/\"\n",
    "filename = filter_filename(dirname, include=hash_str)\n",
    "if len(filename) == 0:\n",
    "    filename = filter_filename(EXP_PATH + \"takashi_2022-5-8/\", include=hash_str)\n",
    "    if len(filename) == 0:\n",
    "        print(f\"hash {hash_str} does not exist!\")\n",
    "        raise\n",
    "\n",
    "try:\n",
    "    data_record = pload(dirname + filename[0])\n",
    "except Exception as e:\n",
    "    print(f\"error {e}\")\n",
    "    # continue\n",
    "    raise\n",
    "if isplot:\n",
    "    plot_learning_curve(data_record)\n",
    "args = init_args(update_legacy_default_hyperparam(data_record[\"args\"]))\n",
    "args.filename = filename\n",
    "# model = load_model(data_record[\"best_model_dict\"], device=device)\n",
    "model = load_model(data_record[\"model_dict\"][-1], device=device)\n",
    "model.eval()\n",
    "p.print(filename, banner_size=100)\n",
    "\n",
    "# Load test dataset:\n",
    "args_test = deepcopy(args)\n",
    "if args.temporal_bundle_steps == 1:\n",
    "    if args.dataset in [\"fno\", \"fno-2\", \"fno-3\"]:\n",
    "        args_test.multi_step = \"20\"\n",
    "    elif args.dataset in [\"fno-1\"]:\n",
    "        args_test.multi_step = \"40\"\n",
    "    elif args.dataset in [\"fno-4\"]:\n",
    "        args_test.multi_step = \"10\"\n",
    "    else:\n",
    "        raise\n",
    "else:\n",
    "    pass\n",
    "args_test.batch_size = 20\n",
    "\n",
    "\n",
    "(dataset_train_val, dataset_test), (train_loader, val_loader, test_loader) = load_data(args_test)\n",
    "test_loader = DataLoader(dataset_test, num_workers=0, collate_fn=deepsnap_Batch.collate(),\n",
    "                         batch_size=20, shuffle=False, drop_last=False)\n",
    "for data in test_loader:\n",
    "    break\n",
    "data.to(device)\n",
    "\n",
    "all_list_dict = {}\n",
    "for hash_str in all_hash:\n",
    "    t_list = []\n",
    "    for i in range(100):\n",
    "        t_start = time.time()\n",
    "        preds, info = model(\n",
    "            data,\n",
    "            pred_steps=np.arange(1,max(parse_multi_step(args_test.multi_step).keys())+1),\n",
    "            latent_pred_steps=None,\n",
    "            is_recons=False,\n",
    "            use_grads=False,\n",
    "            is_y_diff=False,\n",
    "            is_rollout=False,\n",
    "        )\n",
    "        t_end = time.time()\n",
    "        t_list.append(t_end - t_start)\n",
    "        del preds\n",
    "        gc.collect()\n",
    "    full_time = np.mean(t_list)\n",
    "\n",
    "    t_list_evo = []\n",
    "    for i in range(100):\n",
    "        t_start = time.time()\n",
    "        preds, info = model(\n",
    "            data,\n",
    "            pred_steps=[],\n",
    "            latent_pred_steps=np.arange(1,max(parse_multi_step(args_test.multi_step).keys())+1),\n",
    "            is_recons=False,\n",
    "            use_grads=False,\n",
    "            is_y_diff=False,\n",
    "            is_rollout=False,\n",
    "        )\n",
    "        t_end = time.time()\n",
    "        t_list_evo.append(t_end - t_start)\n",
    "        del preds\n",
    "        gc.collect()\n",
    "    evo_time = np.mean(t_list_evo)\n",
    "    print(\"hash {}, full time: {:.6f} +- {:.6f}  evo time: {:.6f} += {:.6f}\".format(\n",
    "        hash_str, full_time, np.std(t_list),\n",
    "        evo_time, np.std(t_list_evo)))\n",
    "    all_list_dict[hash_str] = {\n",
    "        \"evo\": t_list_evo,\n",
    "        \"full\": t_list,\n",
    "    }"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f2f4215a-05a1-4bdd-aa02-4b25b6d79a57",
   "metadata": {},
   "source": [
    "## 3. Plotting:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b00ef8c2-d465-454a-ac75-d2c117cfe4c2",
   "metadata": {},
   "outputs": [],
   "source": [
    "k = 4\n",
    "cmap = \"jet\"\n",
    "import matplotlib.colors as clr\n",
    "import numpy as np\n",
    "\n",
    "\n",
    "data = dataset_test[k]\n",
    "data.to(device)\n",
    "preds, info = model(\n",
    "    data,\n",
    "    pred_steps=np.arange(1,int(args_test.multi_step)+1),\n",
    "    latent_pred_steps=None,\n",
    "    is_recons=False,\n",
    "    use_grads=False,\n",
    "    is_y_diff=False,\n",
    "    is_rollout=False,\n",
    ")\n",
    "pred_reshape = preds[\"n0\"].permute(1,0,2).squeeze(-1)\n",
    "y_reshape = data.node_label[\"n0\"].permute(1,0,2).squeeze(-1)\n",
    "pred_reshape = pred_reshape.reshape(pred_reshape.shape[0], 64, 64)\n",
    "y_reshape = y_reshape.reshape(-1, 64, 64)\n",
    "\n",
    "fig = plt.figure(figsize=(10,4))\n",
    "fig.set_canvas(plt.gcf().canvas)\n",
    "\n",
    "vmax = pred_reshape.max().item()\n",
    "vmin = pred_reshape.min().item()\n",
    "\n",
    "ax = fig.add_subplot(2, 5, 1)\n",
    "ax.matshow(to_np_array(y_reshape[0]), cmap=cmap)\n",
    "ax.set_title(\"Initial vorticity\")\n",
    "plt.xticks(np.array([]))\n",
    "plt.yticks(np.array([]))\n",
    "\n",
    "for i, idx in enumerate([5,10,15,20]):\n",
    "    ax = fig.add_subplot(2, 5, i + 2)\n",
    "    ax.matshow(to_np_array(y_reshape[idx-1]), cmap=cmap)\n",
    "    ax.set_title(f\"t={idx}\")\n",
    "    plt.xticks(np.array([]))\n",
    "    plt.yticks(np.array([]))\n",
    "for i, idx in enumerate([5,10,15,20]):\n",
    "    ax = fig.add_subplot(2, 5, i + 3+4)\n",
    "    ax.matshow(to_np_array(pred_reshape[idx-1]), cmap=cmap)\n",
    "    \n",
    "    plt.xticks(np.array([]))\n",
    "    plt.yticks(np.array([]))\n",
    "plt.savefig(\"2d_nv.pdf\", bbox_inches='tight')\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "base"
  },
  "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.9.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}