figure_emulators.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import os\n",
    "import sys\n",
    "import L96sim\n",
    "\n",
    "from L96_emulator.util import dtype, dtype_np, device\n",
    "\n",
    "res_dir = '/gpfs/work/nonnenma/results/emulators/L96/'\n",
    "data_dir = '/gpfs/work/nonnenma/data/emulators/L96/'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Emulator evaluation results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "import torch \n",
    "from L96_emulator.run import setup, sel_dataset_class\n",
    "from L96_emulator.eval import sortL96fromChannels, sortL96intoChannels, load_model_from_exp_conf\n",
    "from L96_emulator.networks import named_network, Model_forwarder_predictorCorrector, Model_forwarder_rk4default\n",
    "from L96_emulator.util import predictor_corrector, rk4_default, get_data, as_tensor\n",
    "from L96sim.L96_base import f1, f2, pf2\n",
    "\n",
    "# experiments to use: \n",
    "\n",
    "#exp_ids = [91, 92, 93,   77, 85, 86,  94, 95, 96, 74, 89, 90,   71, 87, 88,   68, 81, 82,   65, 83, 84]\n",
    "exp_ids = [77, 97, 90, 71]\n",
    "\n",
    "exp_id_model_sorted = [np.arange(len(exp_ids))]\n",
    "model_names = ['emulator training']\n",
    "\n",
    "\n",
    "all_lgnd = []\n",
    "all_models, all_model_forwarders, all_training_outputs = [], [], []\n",
    "\n",
    "fig = plt.figure(figsize=(8,8))\n",
    "for exp_id in exp_ids:\n",
    "\n",
    "    exp_names = os.listdir('experiments/')   \n",
    "    conf_exp = exp_names[np.where(np.array([name.split('_')[0] for name in exp_names])==str(exp_id))[0][0]][:-4]\n",
    "\n",
    "    args = setup(conf_exp=f'experiments/{conf_exp}.yml')\n",
    "    args.pop('conf_exp')\n",
    "    args['model_forwarder'] = 'rk4_default'\n",
    "\n",
    "    K,J = args['K'], args['J']\n",
    "    assert args['dt_net'] == args['dt']\n",
    "\n",
    "    if J > 0:\n",
    "        F, h, b, c = 10., 1., 10., 10.\n",
    "    else:\n",
    "        h, b, c = 1., 10., 10.\n",
    "        F = 10. if K==36 else 8.\n",
    "\n",
    "    exp_str = 'N=' + str(args['N_trials'] * int(args['T']/args['dt'] - args['spin_up_time']/args['dt']))\n",
    "        \n",
    "    all_lgnd.append(exp_str)\n",
    "\n",
    "    if args['padding_mode'] == 'valid':\n",
    "        print('switching from local training to global evaluation')\n",
    "        args['padding_mode'] = 'circular'\n",
    "    model, model_forwarder, training_outputs = load_model_from_exp_conf(res_dir, args)\n",
    "    all_models.append(model)\n",
    "    all_model_forwarders.append(model_forwarder)\n",
    "    all_training_outputs.append(training_outputs)\n",
    "\n",
    "    if not training_outputs is None:\n",
    "        seq_length = args['seq_length']\n",
    "        plt.semilogy(training_outputs['validation_loss'], label=all_lgnd[-1])\n",
    "        print('final loss', np.min(training_outputs['validation_loss']))\n",
    "\n",
    "all_lgnd = np.array(all_lgnd)\n",
    "plt.title('training')\n",
    "plt.ylabel('validation error')\n",
    "plt.legend()\n",
    "fig.patch.set_facecolor('xkcd:white')\n",
    "plt.show()\n",
    "\n",
    "dX_dt = np.empty(K*(J+1), dtype=dtype_np)\n",
    "\n",
    "train_frac = args['train_frac']\n",
    "normalize_data = bool(args['normalize_data'])\n",
    "dt = args['dt']\n",
    "\n",
    "N_trials = 1000\n",
    "spin_up_time = 50\n",
    "T = (1000)*dt + spin_up_time\n",
    "\n",
    "out, _ = get_data(K=K, J=J, T=T, dt=dt, N_trials=N_trials, F=F, h=h, b=b, c=c, \n",
    "                  resimulate=True, solver=rk4_default,\n",
    "                  save_sim=False, data_dir=data_dir)\n",
    "\n",
    "if J > 0:\n",
    "    def fun(t, x):\n",
    "        return f2(x, F, h, b, c, dX_dt, K, J)\n",
    "else:\n",
    "    def fun(t, x):\n",
    "        return f1(x, F, dX_dt, K)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "T_start = np.arange(int(spin_up_time/dt), int(T/dt)) # grab initial states for rollout from long-running simulations\n",
    "i_trial = np.random.choice(N_trials, size=T_start.shape)\n",
    "idx_show = np.arange(0,len(T_start)-1, len(T_start)//3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# state-prediction RMSES"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "from L96_emulator.networks import Model_forwarder_predictorCorrector, Model_forwarder_rk4default\n",
    "from L96_emulator.networks import Model_forwarder_forwardEuler\n",
    "import torch \n",
    "\n",
    "print('\\n')\n",
    "print('MSEs are on system state !')\n",
    "print('\\n')\n",
    "\n",
    "MFWDs = [Model_forwarder_rk4default]\n",
    "dts = {Model_forwarder_predictorCorrector : args['dt']/10,\n",
    "       Model_forwarder_rk4default : args['dt']}\n",
    "\n",
    "RMSEs_states = np.zeros((len(MFWDs), len(exp_ids), len(T_start)))\n",
    "\n",
    "class Torch_solver(torch.nn.Module):\n",
    "    # numerical solver (from numpy/numba/Julia)\n",
    "    def __init__(self, fun):\n",
    "        self.fun = fun\n",
    "    def forward(self, x):\n",
    "        x = sortL96fromChannels(x.detach().cpu().numpy()).flatten()\n",
    "        return as_tensor(sortL96intoChannels(np.atleast_2d(self.fun(0., x)), J=J))\n",
    "    \n",
    "for mf_i, MFWD in enumerate(MFWDs):\n",
    "\n",
    "    print('\\n')\n",
    "    print(f'solver {MFWD}, dt = {dts[MFWD]}')\n",
    "    print('\\n')\n",
    "\n",
    "    model_forwarder_np = MFWD(Torch_solver(fun), dt=dts[MFWD])\n",
    "    \n",
    "    for m_i, model in enumerate(all_models):\n",
    "        \n",
    "        model_forwarder = MFWD(model=model, dt=dts[MFWD])\n",
    "        \n",
    "        for i in range(len(T_start)):\n",
    "            inputs = out[i_trial[i], T_start[i]] if N_trials > 1 else out[T_start[i]]\n",
    "            inputs_torch = as_tensor(sortL96intoChannels(np.atleast_2d(inputs.copy()),J=J))\n",
    "\n",
    "            out_np = model_forwarder_np(inputs_torch)\n",
    "            out_model = model_forwarder(inputs_torch)\n",
    "\n",
    "            RMSEs_states[mf_i, m_i, i] = np.sqrt(((out_np - out_model)**2).mean().detach().cpu().numpy())\n",
    "\n",
    "    plt.figure(figsize=(16,12))\n",
    "    for i in range(len(exp_id_model_sorted)):\n",
    "        plt.subplot(len(exp_id_model_sorted),2,1+2*i)\n",
    "        plt.semilogy(np.sort(RMSEs_states[mf_i][exp_id_model_sorted[i]],axis=1).T)\n",
    "        plt.ylabel('RMSE')\n",
    "        plt.xlabel('test data point (sorted)')\n",
    "        plt.title(model_names[i])\n",
    "        plt.legend(all_lgnd[exp_id_model_sorted[i]])\n",
    "        plt.subplot(len(exp_id_model_sorted),2,2+2*i)\n",
    "        plt.ylabel('RMSE')\n",
    "        plt.boxplot(RMSEs_states[mf_i][exp_id_model_sorted[i]].T, labels=all_lgnd[exp_id_model_sorted[i]])\n",
    "        plt.title(model_names[i])\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Jacobian error Frobenius norms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "from L96_emulator.networks import Model_forwarder_predictorCorrector, Model_forwarder_rk4default\n",
    "from L96_emulator.util import calc_jakobian_onelevelL96_tendencies, calc_jakobian_rk4, get_jacobian_torch\n",
    "import torch \n",
    "from torch.autograd.functional import jacobian\n",
    "\n",
    "def model_np(inputs):\n",
    "    return fun(0., inputs).copy()\n",
    "\n",
    "print('\\n')\n",
    "print('MSEs are on system state !')\n",
    "print('\\n')\n",
    "\n",
    "MFWDs = [Model_forwarder_rk4default]\n",
    "L2_jakobians = np.zeros((len(MFWDs), len(exp_ids), len(T_start)))\n",
    "\n",
    "  \n",
    "for mf_i, MFWD in enumerate(MFWDs):\n",
    "\n",
    "    print('\\n')\n",
    "    print(f'solver {MFWD}, dt = {dts[MFWD]}')\n",
    "    print('\\n')\n",
    "    \n",
    "    for m_i, model in enumerate(all_models):\n",
    "        \n",
    "        for p in model.parameters():\n",
    "            p.requires_grad = False\n",
    "        \n",
    "        model_forwarder = MFWD(model=model, dt=dts[MFWD])\n",
    "        print('\\n')\n",
    "        print(f'model forwarder for {model}')\n",
    "        print('\\n')\n",
    "        \n",
    "        for i in range(len(T_start)):\n",
    "            inputs = out[i_trial[i], T_start[i]] if N_trials > 1 else out[T_start[i]]\n",
    "            inputs_torch = as_tensor(sortL96intoChannels(np.atleast_2d(inputs.copy()),J=J))\n",
    "            inputs_torch.requires_grad = True\n",
    "            \n",
    "            #J_np = calc_jakobian_onelevelL96_tendencies(inputs, n=K)\n",
    "            J_np = calc_jakobian_rk4(inputs, calc_f=model_np, \n",
    "                         calc_J_f=calc_jakobian_onelevelL96_tendencies, dt=dt, n=K)\n",
    "\n",
    "            #J_torch = get_jacobian_torch(model_forwarder, inputs=inputs_torch, n=K)\n",
    "            J_torch = jacobian(model_forwarder, inputs_torch).detach().cpu().numpy()\n",
    "            \n",
    "            L2_jakobians[mf_i, m_i, i] = np.sqrt(((J_np - J_torch)**2).sum())\n",
    "\n",
    "    plt.figure(figsize=(16,12))\n",
    "    for i in range(len(exp_id_model_sorted)):\n",
    "        plt.subplot(2,2,1+2*i)\n",
    "        plt.semilogy(np.sort(L2_jakobians[mf_i][exp_id_model_sorted[i]],axis=1).T)\n",
    "        plt.ylabel('RMSE')\n",
    "        plt.xlabel('test data point (sorted)')\n",
    "        plt.title(model_names[i])\n",
    "        plt.legend(all_lgnd[exp_id_model_sorted[i]])\n",
    "        plt.subplot(2,2,2+2*i)\n",
    "        plt.ylabel('RMSE')\n",
    "        plt.boxplot(L2_jakobians[mf_i][exp_id_model_sorted[i]].T, labels=all_lgnd[exp_id_model_sorted[i]])\n",
    "        plt.title(model_names[i])\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# rollout RMSEs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "T_data=1      # length of base simulation (in [au]!) to get initial state for rollouts\n",
    "N_trials=1000 # number of base simulations \n",
    "\n",
    "n_start=1000 # number of rollouts \n",
    "T_dur=100    # length of rollouts (in steps!)\n",
    "F = 8.  # Lorenz-96 forcing parameter\n",
    "\n",
    "##########################\n",
    "#       test data        #\n",
    "##########################\n",
    "\n",
    "from L96_emulator.util import rk4_default, get_data\n",
    "\n",
    "spin_up_time = 5.\n",
    "T = T_data*args['dt'] + spin_up_time\n",
    "\n",
    "try:\n",
    "    assert data.shape == (N_trials, int(T/args['dt'])+1, args['K']*(args['J']+1))\n",
    "    print('found data with matching specs (shape)')\n",
    "except:\n",
    "    print('generating test data')\n",
    "    data, _ = get_data(K=args['K'], J=args['J'], T=T, dt=args['dt'], N_trials=N_trials, \n",
    "                      F=args['F_net'], \n",
    "                      resimulate=True, solver=rk4_default,\n",
    "                      save_sim=False, data_dir=data_dir)\n",
    "\n",
    "##########################\n",
    "#       rollouts         #\n",
    "##########################\n",
    "\n",
    "from L96sim.L96_base import f1\n",
    "\n",
    "T_start = np.linspace(int(spin_up_time/args['dt']), int(T/args['dt']), n_start).astype(np.int)\n",
    "i_trial = np.random.choice(N_trials, size=T_start.shape, replace=False)\n",
    "print('T_start, i_tria', (T_start, i_trial))\n",
    "\n",
    "class Torch_solver(torch.nn.Module):\n",
    "    # numerical solver (from numpy/numba/Julia)\n",
    "    def __init__(self, fun):\n",
    "        self.fun = fun\n",
    "    def forward(self, x):\n",
    "        J = x.shape[-2] - 1\n",
    "        x = sortL96fromChannels(x.detach().cpu().numpy()).flatten()\n",
    "        return as_tensor(sortL96intoChannels(np.atleast_2d(self.fun(0., x)), J=J))\n",
    "\n",
    "\n",
    "sols = np.nan * np.ones((len(all_model_forwarders)+1, n_start, T_dur+1, args['K']))\n",
    "\n",
    "for i_model in range(len(all_model_forwarders)): \n",
    "\n",
    "    model_forwarder_i = all_model_forwarders[i_model]\n",
    "\n",
    "    def model_simulate(y0, dy0, n_steps):\n",
    "        x = np.empty((n_steps+1, *y0.shape[1:]))\n",
    "        x[0] = y0.copy()\n",
    "        xx = as_tensor(x[0])\n",
    "        for i in range(1,n_steps+1):\n",
    "            xx = model_forwarder_i(xx.reshape(-1,J+1,args['K']))\n",
    "            x[i] = xx.detach().cpu().numpy().copy()\n",
    "        return x\n",
    "\n",
    "    print('forwarding model ' + str(model_forwarder_i))\n",
    "    X_init = []\n",
    "    for i in range(n_start):\n",
    "        X_init.append(data[i_trial[i], T_start[i]] if N_trials > 1 else data[T_start[i]])\n",
    "        X_init[-1] = sortL96intoChannels(X_init[-1][:args['K']*(J+1)],J=J)\n",
    "    X_init = np.stack(X_init)\n",
    "    X_init = X_init.reshape(1, *X_init.shape)\n",
    "    with torch.no_grad():\n",
    "        sol = model_simulate(y0=X_init, dy0=None, n_steps=T_dur)\n",
    "    sols[i_model,:,:,:] = sol[:,:,0,:].transpose(1,0,2)\n",
    "\n",
    "# not parallelising Numba model over initial values for rollouts\n",
    "model_forwarder_np = Model_forwarder_rk4default(Torch_solver(fun), dt=args['dt'])\n",
    "def model_simulate(y0, dy0, n_steps):\n",
    "    x = np.empty((n_steps+1, *y0.shape[1:]))\n",
    "    x[0] = y0.copy()\n",
    "    xx = as_tensor(x[0]).reshape(1,1,-1)\n",
    "    for i in range(1,n_steps+1):\n",
    "        xx = model_forwarder_np(xx.reshape(-1,args['J']+1,args['K']))\n",
    "        x[i] = xx.detach().cpu().numpy().copy()\n",
    "    return x\n",
    "\n",
    "print('forwarding np model')\n",
    "X_init = []\n",
    "for i in range(n_start):\n",
    "    X_init = data[i_trial[i], T_start[i]] if N_trials > 1 else data[T_start[i]]\n",
    "    X_init = sortL96intoChannels(X_init,J=args['J'])\n",
    "    X_init = X_init.reshape(1, *X_init.shape)\n",
    "\n",
    "    with torch.no_grad():\n",
    "        sol = model_simulate(y0=X_init, dy0=None, n_steps=T_dur)\n",
    "    sols[-1,i,:,:] = sol[:,0,:]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### actual figure"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "RMSEs_states.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fontsize=18\n",
    "\n",
    "xx = np.array([1200, 4000, 12000])\n",
    "idx_dn = np.arange(len(xx))\n",
    "\n",
    "plt.figure(figsize=(16,4))\n",
    "RMSES_all = [RMSEs_states[0][:,:len(xx)], L2_jakobians[0][:,:len(xx)]]\n",
    "titles = [r'state updates $x_{t+\\Delta}$', r'Jacobians $J_\\mathcal{M}$']\n",
    "yaxes = ['', 'Frobenius norm of error']\n",
    "for i in range(len(RMSES_all)):\n",
    "\n",
    "    RMSEs = RMSES_all[i]\n",
    "\n",
    "    ax = plt.subplot(1,3,i+2)    \n",
    "    plt.semilogx(xx[:len(idx_dn)], RMSEs[idx_dn,:].mean(axis=1), '*', color='black', label='deepNet')\n",
    "    plt.xticks([], fontsize=fontsize)\n",
    "    plt.ylabel(yaxes[i], fontsize=fontsize)\n",
    "    #if np.mod(i,2) == 1:\n",
    "    #    plt.legend(fontsize=fontsize)\n",
    "    plt.title(titles[i], fontsize=fontsize)\n",
    "    ax.spines['top'].set_visible(False)\n",
    "    ax.spines['right'].set_visible(False)\n",
    "    if i == 0:\n",
    "        plt.axis([xx[0]*0.95, xx[-1]*1.5, 0, 1.05*RMSEs[idx_dn,:].mean(axis=1).max()])\n",
    "        plt.yticks([0, 0.010], fontsize=fontsize)\n",
    "        box = ax.get_position()        \n",
    "        box.x1 -= 0.015\n",
    "        ax.set_position(box)    \n",
    "    else:\n",
    "        plt.axis([xx[0]*0.95, xx[-1]*1.5, 0.05, 1.05*RMSEs[idx_dn,:].mean(axis=1).max()])\n",
    "        plt.yticks([0.1, 0.2], fontsize=fontsize)\n",
    "        box = ax.get_position()        \n",
    "        box.x0 += 0.015\n",
    "        ax.set_position(box)    \n",
    "    plt.xticks([1e3, 1e4], #, 1e5],  \n",
    "               [r'$10^3$', r'$10^4$'], #, r'$10^5$'], \n",
    "               fontsize=fontsize)\n",
    "    plt.xlabel('training set size N', fontsize=fontsize)\n",
    "\n",
    "\n",
    "ax = plt.subplot(1,3,1)\n",
    "clrs = ['purple','blue', 'orange', 'green', 'k']\n",
    "\n",
    "rmses = np.zeros((len(xx), n_start, T_dur+1))\n",
    "for i_model in range(len(xx))[::-1]:\n",
    "    for i in range(n_start):\n",
    "        rmses[i_model,i,:] = np.sqrt(np.mean((sols[i_model,i] - sols[-1,i])**2, axis=1))\n",
    "    plt.plot(rmses[i_model].mean(axis=0), clrs[i_model], label='N='+str(int(xx[i_model])), linewidth=2.5)\n",
    "\n",
    "plt.legend(fontsize=fontsize, frameon=False, handlelength=0.9, loc=(0.02, 0.6))\n",
    "plt.ylabel('RMSE', fontsize=fontsize)\n",
    "plt.xlabel('time [au]', fontsize=fontsize)\n",
    "plt.xticks([0, T_dur/2, T_dur], [0, args['dt']*T_dur/2, args['dt']*T_dur], fontsize=fontsize)\n",
    "plt.yticks([0, 2, 4], fontsize=fontsize)\n",
    "ax.spines['top'].set_visible(False)\n",
    "ax.spines['right'].set_visible(False)\n",
    "plt.title('rollouts', fontsize=fontsize)\n",
    "box = ax.get_position()        \n",
    "box.x1 -= 0.015\n",
    "ax.set_position(box)    \n",
    "\n",
    "\n",
    "plt.gcf().text(0.1,  0.9, 'a)', fontsize=fontsize, weight='bold')\n",
    "plt.gcf().text(0.36, 0.9, 'b)', fontsize=fontsize, weight='bold')\n",
    "plt.gcf().text(0.62, 0.9, 'c)', fontsize=fontsize, weight='bold')\n",
    "\n",
    "#plt.savefig(res_dir + 'figs/emulator_eval.pdf', bbox_inches='tight', pad_inches=0, frameon=False)\n",
    "\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# local emulator training "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "import torch \n",
    "from L96_emulator.run import setup, sel_dataset_class\n",
    "from L96_emulator.eval import sortL96fromChannels, sortL96intoChannels, load_model_from_exp_conf\n",
    "from L96_emulator.networks import named_network, Model_forwarder_predictorCorrector, Model_forwarder_rk4default\n",
    "from L96_emulator.util import predictor_corrector, rk4_default, get_data, as_tensor\n",
    "from L96sim.L96_base import f1, f2, pf2\n",
    "\n",
    "# experiments to use: \n",
    "\n",
    "# 4x4 one-level L96, K=40, F=8 with various degrees of locality and batch-size\n",
    "exp_ids = np.concatenate((np.arange(49, 61), np.arange(61, 65)))\n",
    "exp_id_model_sorted=[np.arange(4), np.arange(4,8), np.arange(8,12), np.arange(12,16), ]\n",
    "model_names=['batch-size 1', 'batch-size 4', 'batch-size 16', 'batch-size 64']\n",
    "\n",
    "\n",
    "\n",
    "all_lgnd = []\n",
    "all_models, all_model_forwarders, all_training_outputs = [], [], []\n",
    "\n",
    "fig = plt.figure(figsize=(8,8))\n",
    "for exp_id in exp_ids:\n",
    "\n",
    "    exp_names = os.listdir('experiments/')   \n",
    "    conf_exp = exp_names[np.where(np.array([name[:2] for name in exp_names])==str(exp_id))[0][0]][:-4]\n",
    "\n",
    "    args = setup(conf_exp=f'experiments/{conf_exp}.yml')\n",
    "    args.pop('conf_exp')\n",
    "    args['model_forwarder'] = 'rk4_default'\n",
    "\n",
    "    K,J = args['K'], args['J']\n",
    "    assert args['dt_net'] == args['dt']\n",
    "\n",
    "    if J > 0:\n",
    "        F, h, b, c = 10., 1., 10., 10.\n",
    "    else:\n",
    "        h, b, c = 1., 10., 10.\n",
    "        F = 10. if K==36 else 8.\n",
    "\n",
    "    exp_str = 'N=' + str(args['N_trials'] * int(args['T']/args['dt'] - args['spin_up_time']/args['dt']))\n",
    "        \n",
    "    all_lgnd.append(exp_str)\n",
    "\n",
    "    if args['padding_mode'] == 'valid':\n",
    "        print('switching from local training to global evaluation')\n",
    "        args['padding_mode'] = 'circular'\n",
    "    model, model_forwarder, training_outputs = load_model_from_exp_conf(res_dir, args)\n",
    "    all_models.append(model)\n",
    "    all_model_forwarders.append(model_forwarder)\n",
    "    all_training_outputs.append(training_outputs)\n",
    "\n",
    "    if not training_outputs is None:\n",
    "        seq_length = args['seq_length']\n",
    "        plt.semilogy(training_outputs['validation_loss'], label=all_lgnd[-1])\n",
    "        print('final loss', np.min(training_outputs['validation_loss']))\n",
    "\n",
    "all_lgnd = np.array(all_lgnd)\n",
    "plt.title('training')\n",
    "plt.ylabel('validation error')\n",
    "plt.legend()\n",
    "fig.patch.set_facecolor('xkcd:white')\n",
    "plt.show()\n",
    "\n",
    "dX_dt = np.empty(K*(J+1), dtype=dtype_np)\n",
    "\n",
    "train_frac = args['train_frac']\n",
    "normalize_data = bool(args['normalize_data'])\n",
    "dt = args['dt']\n",
    "\n",
    "N_trials = 1000\n",
    "spin_up_time = 50\n",
    "T = (1000)*dt + spin_up_time\n",
    "\n",
    "out, _ = get_data(K=K, J=J, T=T, dt=dt, N_trials=N_trials, F=F, h=h, b=b, c=c, \n",
    "                  resimulate=True, solver=rk4_default,\n",
    "                  save_sim=False, data_dir=data_dir)\n",
    "\n",
    "if J > 0:\n",
    "    def fun(t, x):\n",
    "        return f2(x, F, h, b, c, dX_dt, K, J)\n",
    "else:\n",
    "    def fun(t, x):\n",
    "        return f1(x, F, dX_dt, K)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "T_start = np.arange(int(spin_up_time/dt), int(T/dt)) # grab initial states for rollout from long-running simulations\n",
    "i_trial = np.random.choice(N_trials, size=T_start.shape)\n",
    "idx_show = np.arange(0,len(T_start)-1, len(T_start)//3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from L96_emulator.networks import Model_forwarder_predictorCorrector, Model_forwarder_rk4default\n",
    "from L96_emulator.networks import Model_forwarder_forwardEuler\n",
    "import torch \n",
    "\n",
    "print('\\n')\n",
    "print('MSEs are on system state !')\n",
    "print('\\n')\n",
    "\n",
    "MFWDs = [Model_forwarder_rk4default]\n",
    "dts = {Model_forwarder_predictorCorrector : args['dt']/10,\n",
    "       Model_forwarder_rk4default : args['dt']}\n",
    "\n",
    "RMSEs_states = np.zeros((len(MFWDs), len(exp_ids), len(T_start)))\n",
    "\n",
    "class Torch_solver(torch.nn.Module):\n",
    "    # numerical solver (from numpy/numba/Julia)\n",
    "    def __init__(self, fun):\n",
    "        self.fun = fun\n",
    "    def forward(self, x):\n",
    "        x = sortL96fromChannels(x.detach().cpu().numpy()).flatten()\n",
    "        return as_tensor(sortL96intoChannels(np.atleast_2d(self.fun(0., x)), J=J))\n",
    "    \n",
    "for mf_i, MFWD in enumerate(MFWDs):\n",
    "\n",
    "    print('\\n')\n",
    "    print(f'solver {MFWD}, dt = {dts[MFWD]}')\n",
    "    print('\\n')\n",
    "\n",
    "    model_forwarder_np = MFWD(Torch_solver(fun), dt=dts[MFWD])\n",
    "    \n",
    "    for m_i, model in enumerate(all_models):\n",
    "        \n",
    "        model_forwarder = MFWD(model=model, dt=dts[MFWD])\n",
    "        \n",
    "        for i in range(len(T_start)):\n",
    "            inputs = out[i_trial[i], T_start[i]] if N_trials > 1 else out[T_start[i]]\n",
    "            inputs_torch = as_tensor(sortL96intoChannels(np.atleast_2d(inputs.copy()),J=J))\n",
    "\n",
    "            out_np = model_forwarder_np(inputs_torch)\n",
    "            out_model = model_forwarder(inputs_torch)\n",
    "\n",
    "            RMSEs_states[mf_i, m_i, i] = np.sqrt(((out_np - out_model)**2).mean().detach().cpu().numpy())\n",
    "\n",
    "    plt.figure(figsize=(16,12))\n",
    "    for i in range(len(exp_id_model_sorted)):\n",
    "        plt.subplot(len(exp_id_model_sorted),2,1+2*i)\n",
    "        plt.semilogy(np.sort(RMSEs_states[mf_i][exp_id_model_sorted[i]],axis=1).T)\n",
    "        plt.ylabel('RMSE')\n",
    "        plt.xlabel('test data point (sorted)')\n",
    "        plt.title(model_names[i])\n",
    "        plt.legend(all_lgnd[exp_id_model_sorted[i]])\n",
    "        plt.subplot(len(exp_id_model_sorted),2,2+2*i)\n",
    "        plt.ylabel('RMSE')\n",
    "        plt.boxplot(RMSEs_states[mf_i][exp_id_model_sorted[i]].T, labels=all_lgnd[exp_id_model_sorted[i]])\n",
    "        plt.title(model_names[i])\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fontsize=18\n",
    "plt.figure(figsize=(16,4))\n",
    "\n",
    "\n",
    "ax = plt.subplot(1,3,1, projection='polar')\n",
    "xx = np.arange(K)/K*2*np.pi\n",
    "state = out[0,-1]\n",
    "\n",
    "nsegs = 4\n",
    "r_outer_base, r_inner_base = 13, -10\n",
    "for i in range(nsegs):\n",
    "    alpha = 1 if i == 1 else 0.4\n",
    "    \n",
    "    xrange = xx[i*K//(4*nsegs):(i+1)*K//(4*nsegs)]\n",
    "    ax.plot(xrange, state[i*K//(4*nsegs):(i+1)*K//(4*nsegs)], 'k', alpha=alpha, linewidth=2.5)\n",
    "    r_outer = r_outer_base - 3 if i == 1 else r_outer_base\n",
    "    r_inner = r_inner_base + 3 if i == 1 else r_inner_base\n",
    "    ax.plot(xrange, r_outer*np.ones_like(xrange), 'b', alpha=alpha, linewidth=2.5)\n",
    "    ax.plot([xrange[0], xrange[0]], [r_outer-2, r_outer], 'b', alpha=alpha, linewidth=2.5)\n",
    "    ax.plot([xrange[-1], xrange[-1]], [r_outer-2, r_outer], 'b', alpha=alpha, linewidth=2.5)\n",
    "\n",
    "    xrange = np.linspace(xrange[0]-8/K*2*np.pi, xrange[-1]+4/K*2*np.pi, 40)\n",
    "    ax.plot(xrange, r_inner*np.ones_like(xrange), 'orange', alpha=alpha, linewidth=2.5)\n",
    "    ax.plot([xrange[0], xrange[0]], [r_inner, r_inner+2], 'orange', alpha=alpha, linewidth=2.5)\n",
    "    ax.plot([xrange[-1], xrange[-1]], [r_inner, r_inner+2], 'orange', alpha=alpha, linewidth=2.5)\n",
    "\n",
    "plt.plot(-1, -1, 'k', linewidth=2.5, label=r'state $x$')\n",
    "plt.plot(-1, -1, 'orange', linewidth=2.5, label='input region')\n",
    "plt.plot(-1, -1, 'b', linewidth=2.5, label='prediction region')\n",
    "plt.legend(fontsize=fontsize, frameon=False, loc=(-0.05,-0.15))\n",
    "\n",
    "ax.set_thetagrids([])\n",
    "ax.set_rorigin(-100)\n",
    "plt.axis([0, np.pi/2, -20, 20])\n",
    "plt.yticks([-5, 0, 5, 10], fontsize=fontsize)\n",
    "if i == 1:\n",
    "    ax.spines['polar'].set_visible(False)\n",
    "    ax.spines['inner'].set_visible(False)\n",
    "    ax.spines['end'].set_visible(False)\n",
    "ax.set_facecolor('xkcd:white')\n",
    "ax.axis('off')\n",
    "box = ax.get_position()\n",
    "box.x1 = box.x1 + 0.1 * (box.x1-box.x0)\n",
    "box.y1 = box.y1 + 0.1 * (box.y1-box.y0)\n",
    "ax.set_position(box)\n",
    "\n",
    "\n",
    "plt.subplot(1,3,2)\n",
    "plt.imshow(RMSEs_states[-1].mean(axis=1).reshape(4,4), vmax=3e-6, cmap='cividis')\n",
    "plt.xlabel('local region size', fontsize=fontsize)\n",
    "plt.xticks(range(4), ['640', '160', '40', '10'], fontsize=fontsize)\n",
    "plt.ylabel('batch-size', fontsize=fontsize)\n",
    "plt.yticks(range(4), ['1', '4', '16', '64'], fontsize=fontsize)\n",
    "plt.colorbar()\n",
    "plt.title(r'state update RMSE', fontsize=fontsize)\n",
    "for i in range(4):\n",
    "    plt.plot(np.array([-.5, .5])+i, np.array([-0.5, -0.5])+i, 'k--', linewidth=2)\n",
    "    plt.plot(np.array([.5, .5])+i, np.array([-.5, .5])+i, 'k--', linewidth=2)\n",
    "plt.plot([0.5, 3.5], [-0.5, -0.5], 'k--', linewidth=2)\n",
    "plt.plot([3.5, 3.5], [-0.5, 2.5], 'k--', linewidth=2)\n",
    "\n",
    "ax = plt.subplot(1,3,3)\n",
    "batch_sizes = [1, 1, 1, 1, 4, 4, 4, 4, 16, 16, 16, 16, 64, 64, 64, 64]\n",
    "region_sizes = [640, 160, 40, 10, 640, 160, 40, 10, 640, 160, 40, 10, 640, 160, 40, 10]\n",
    "plt.semilogx([640, 640], [0.1e-6, 2e-6], 'k--', linewidth=2.)\n",
    "for i in range(len(RMSEs_states[-1])):\n",
    "    plt.semilogx(batch_sizes[i]*region_sizes[i], RMSEs_states[-1][i].mean(), 'b.', markersize=8.0)\n",
    "plt.xlabel('locations per minibatch', fontsize=fontsize)\n",
    "plt.xticks(fontsize=fontsize)\n",
    "plt.yticks(np.array([0.5, 1., 1.5, 2.])*1e-6, fontsize=fontsize)\n",
    "ax.spines['top'].set_visible(False)\n",
    "ax.spines['right'].set_visible(False)\n",
    "plt.title(r'state update RMSE', fontsize=fontsize)\n",
    "plt.ylabel('RMSE', fontsize=fontsize)\n",
    "box = ax.get_position()\n",
    "box.x0 += 0.015\n",
    "ax.set_position(box)\n",
    "\n",
    "plt.gcf().text(0.125, 0.9, 'a)', fontsize=fontsize, weight='bold')\n",
    "plt.gcf().text(0.35,  0.9, 'b)', fontsize=fontsize, weight='bold')\n",
    "plt.gcf().text(0.625,  0.9, 'c)', fontsize=fontsize, weight='bold')\n",
    "\n",
    "plt.savefig(res_dir + 'figs/emulator_local.pdf', bbox_inches='tight', pad_inches=0, frameon=False)\n",
    "\n",
    "plt.show()\n"
   ]
  },
  {
   "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.7.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}