2021-CVAE_v3.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Conditional VAEs for Uncertainty Quantification in Fatigue Simulations\n",
    "If running on google colab remember to change your runtime to \"GPU\" :\n",
    "\n",
    "\n",
    "**(Runtime’ > ‘Change runtime type’ > ‘Hardware accelerator’ > select GPU)**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#@title Clone repository, and prepare directory for the rest of computations:\n",
    "!git clone https://github.com/mylonasc/fatigue_cvae.git\n",
    "!cp -r fatigue_cvae/data .\n",
    "!cp fatigue_cvae/utils.py .\n",
    "!pip install tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Imports\n",
    "import tensorflow as tf\n",
    "import tensorflow_probability as tfp\n",
    "tfd = tfp.distributions\n",
    "import numpy as np\n",
    "from utils import BladeCSMeshPlotter, default_mesh_file,DELDatasetPreProc, PlotFatvals\n",
    "import matplotlib.pyplot as pplot\n",
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tf.__version__, tfp.__version__\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# A conditional VAE (IWAE) model using the functional Keras API\n",
    "class CVAEModel:\n",
    "    def __init__(self, n_conditioning, input_dims, n_latent,layer_width = 64, n_layers_enc = 2,n_layers_dec = 2,\n",
    "                 activation = \"relu\", n_iwae = 40, enc_use_cond = True):\n",
    "        \"\"\"\n",
    "        A simple FFNN conditional VAE model [1] for tf 2.4\n",
    "        Implements the importance weighted autoencoder [2] by sampling the latent space multiple times.\n",
    "        A similar implementation (initially in tensorflow 1.13) was used for [3].\n",
    "        \n",
    "        parameters:\n",
    "          n_conditioning : number of conditioning variables\n",
    "          input_dims     : number of dimensions of the input\n",
    "          n_latent       : the size of the latent variable.\n",
    "          layer_width    : (default: 50) a width value for all layers\n",
    "          n_layers_enc   : (default: 2) the number of encoder layers (with activations). The part of the\n",
    "                           encoder that parametrizes the std of the latent is passed through a softplus.\n",
    "          n_layers_dec   : (default: 2) number of decoder layers\n",
    "          activation     : ['relu'] what activation to use \n",
    "          n_iwae         : [10] number of samples from the latent space passed through the decoder for each pass. \n",
    "          \n",
    "          \n",
    "        references:\n",
    "        \n",
    "          [1] Kingma, Diederik P., and Max Welling. \"Auto-encoding variational bayes.\" \n",
    "              arXiv preprint arXiv:1312.6114 (2013).\n",
    "              \n",
    "          [2] Burda, Yuri, Roger Grosse, and Ruslan Salakhutdinov. \"Importance weighted autoencoders.\" \n",
    "              arXiv preprint arXiv:1509.00519 (2015).\n",
    "              \n",
    "          [3] Mylonas, Charilaos, Imad Abdallah, and E. N. Chatzi. \"Deep unsupervised learning\n",
    "              for condition monitoring and prediction of high dimensional data with application \n",
    "              on windfarm scada data.\" Model Validation and Uncertainty Quantification, Volume 3. \n",
    "              Springer, Cham, 2020. 189-196.\n",
    "\n",
    "        \"\"\"\n",
    "        all_params = [n_conditioning, input_dims, layer_width, n_layers_dec, n_layers_enc, n_latent,activation, n_iwae]\n",
    "        pnames = [\"n_conditioning\", \"input_dims\", \"layer_width\", \"n_layers_dec\", \"n_layers_enc\", \"n_latent\",\"activation\",\"n_iwae\"]\n",
    "        self.cvae_params = {k: v for v,k in zip(all_params,pnames)}\n",
    "        \n",
    "        if enc_use_cond:\n",
    "            input_dims_xx = input_dims  + n_conditioning\n",
    "        else:\n",
    "            input_dims_xx = input_dims\n",
    "        \n",
    "        self.encoder = CVAEModel.make_encoder(input_dims = input_dims_xx,\n",
    "                                              layer_width = layer_width,\n",
    "                                              n_layers = n_layers_enc , \n",
    "                                              n_latent= n_latent)\n",
    "        \n",
    "        \n",
    "        prior = tfd.MultivariateNormalDiag(loc = tf.zeros(n_latent), scale_diag=tf.ones(n_latent))\n",
    "        posterior = tfp.layers.DistributionLambda(make_distribution_fn = lambda t: tfd.MultivariateNormalDiag(loc = t[:,0:n_latent], \n",
    "                                                                                                              scale_diag= 1e-5 + tf.nn.softplus(t[:,n_latent:])),\n",
    "                                                  convert_to_tensor_fn= lambda s : s.sample(n_iwae))\n",
    "        \n",
    "        \n",
    "        X = tf.keras.Input(shape = (input_dims,), name = \"X_in\")\n",
    "        W = tf.keras.Input(shape = (n_conditioning,) , name = \"W\")\n",
    "        \n",
    "        if enc_use_cond:\n",
    "            X_in = tf.keras.layers.concatenate([X,W], axis = -1)\n",
    "        else:\n",
    "            X_in = X\n",
    "        \n",
    "        vae_params = self.encoder(X_in)\n",
    "        \n",
    "        posterior_out = posterior(tf.keras.layers.concatenate(vae_params))\n",
    "        \n",
    "        kl_div = tfd.kl_divergence(tfd.MultivariateNormalDiag(loc = vae_params[0], scale_diag=vae_params[1] ), prior)\n",
    "        self.posterior_out = posterior_out\n",
    "        \n",
    "        dec_output = input_dims\n",
    "        self.decoder = CVAEModel.make_decoder(output_dims = dec_output,\n",
    "                                      layer_width = layer_width,\n",
    "                                      conditioning_dims = n_conditioning,\n",
    "                                      n_layers = n_layers_dec,\n",
    "                                      n_latent= n_latent,\n",
    "                                      n_iwae = n_iwae, \n",
    "                                      activation = activation,\n",
    "                                      input_tensor = posterior_out)\n",
    "        \n",
    "\n",
    "        decoder_out = self.decoder([posterior_out, W])\n",
    "        y_decoder_out = decoder_out\n",
    "        self.vae_model = tf.keras.Model(inputs = [X,W] , outputs = [y_decoder_out, posterior_out, kl_div])\n",
    "        \n",
    "        self.posterior_out = posterior_out\n",
    "        \n",
    "        self.prior = prior\n",
    "        self.posterior = posterior\n",
    "    \n",
    "    \n",
    "        \n",
    "    @staticmethod\n",
    "    def make_encoder(input_dims = None,layer_width = None, n_layers = None,n_latent = None,activation = None):\n",
    "        x_in = tf.keras.layers.Input(shape = (input_dims,))\n",
    "\n",
    "        ln = tf.keras.layers.Dense(layer_width)(x_in)\n",
    "        for l in range(n_layers - 1):\n",
    "            ln = tf.keras.layers.Dense(layer_width, activation = activation)(ln)\n",
    "            ln = tf.keras.layers.Dropout(rate = 0.2)(ln)\n",
    "\n",
    "        l_out_mean = tf.keras.layers.Dense(n_latent)(ln)\n",
    "        l_out_std = tf.keras.layers.Dense(n_latent, activation = tf.nn.softplus)(ln)\n",
    "\n",
    "        return tf.keras.Model(x_in, [l_out_mean, l_out_std])\n",
    "\n",
    "    @staticmethod\n",
    "    def make_decoder(output_dims = None,\n",
    "                     n_latent= None ,\n",
    "                     layer_width = None, \n",
    "                     conditioning_dims = None , \n",
    "                     n_layers = None ,\n",
    "                     activation = None,\n",
    "                     n_iwae = None,\n",
    "                    input_tensor = None):\n",
    "        \n",
    "        # This is to circumvent a keras bug (see https://github.com/tensorflow/probability/issues/1200)\n",
    "        dd = tf.identity(input_tensor, name = \"z_in\")\n",
    "        \n",
    "        z_in = tf.keras.layers.Input(tensor = dd, name = \"z_in\")\n",
    "        \n",
    "        w_in = tf.keras.layers.Input(shape = (conditioning_dims, ), name =  \"w_in\")\n",
    "        w_in_tiled = tf.tile(w_in[tf.newaxis, ... ] , [n_iwae, 1,1] )\n",
    "        \n",
    "        yyin = tf.keras.layers.concatenate([z_in, w_in_tiled])\n",
    "        \n",
    "        ln = tf.keras.layers.Dense(layer_width)(yyin)\n",
    "\n",
    "        for l in range(n_layers - 1):\n",
    "            ln = tf.keras.layers.Dense(output_dims, activation = activation)(ln)\n",
    "            ln = tf.keras.layers.Dropout(rate = 0.2)(ln)\n",
    "\n",
    "        y_out = tf.keras.layers.Dense(output_dims)(ln)\n",
    "        return tf.keras.Model(inputs=[z_in, w_in], outputs = y_out)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load the fatigue dataset\n",
    "The dataset consists of 1999 simulations of fatigue of the blade root. Please refer back to the paper for more details."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = DELDatasetPreProc()\n",
    "x_norm = dataset.get_normalized_data_DEL()\n",
    "w_norm = dataset.get_normalized_data_X()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#@Title some code that is needed for plotting.\n",
    "plot_blade = BladeCSMeshPlotter(default_mesh_file, dataset.hasval_inds)\n",
    "mesh = plot_blade.mesh\n",
    "nl_inds = np.array(mesh.nl_2d[:,0], dtype = np.int)\n",
    "el_loc = mesh.nl_2d[:,1:3]\n",
    "elnums = mesh.el_2d[:,0]\n",
    "node_1 = el_loc[mesh.el_2d[:,1]-1]\n",
    "node_2 = el_loc[mesh.el_2d[:,2]-1]\n",
    "node_3 = el_loc[mesh.el_2d[:,3]-1]\n",
    "node_4 = el_loc[mesh.el_2d[:,4]-1]\n",
    "node_pos_avg = (node_1 + node_2 + node_3 + node_4 )/4\n",
    "\n",
    "r0 =( node_pos_avg[:,0]**2 + node_pos_avg[:,1] ** 2 ) **0.5\n",
    "theta_el =np.arctan2(  node_pos_avg[:,1] , node_pos_avg[:,0]) \n",
    "\n",
    "# Filter the nodes that have a value for DEL:\n",
    "r0 = r0[plot_blade.mesh_finite_mask]\n",
    "theta_el = theta_el[plot_blade.mesh_finite_mask]\n",
    "\n",
    "pfv = PlotFatvals(theta_el, r0)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Training a CVAE on fatigue data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# initialize the optimizer:\n",
    "opt = tf.keras.optimizers.Adam(learning_rate= 0.001)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Train-test split:\n",
    "from sklearn.model_selection import train_test_split\n",
    "xtrain,xtest, wtrain , wtest = train_test_split(x_norm, w_norm, test_size=0.1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_dims = x_norm.shape[-1]\n",
    "cond_size = w_norm.shape[-1]\n",
    "nlatent = 30\n",
    "cvae = CVAEModel( cond_size,input_dims,  nlatent, layer_width=300)\n",
    "\n",
    "def eval_loss(X,W, beta = 0.1 ):\n",
    "    Xhat ,z_out, kl_loss = cvae.vae_model([X,W])\n",
    "    rec_loss = tf.reduce_mean(tf.pow(Xhat-X,2),-1)\n",
    "    return rec_loss + beta * kl_loss/nlatent , rec_loss, kl_loss, Xhat, z_out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_losses = []\n",
    "test_losses = []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#@title A function to anneal the beta parameter\n",
    "def beta_cyclic_anneal(epoch, beta_0 = 0.01,init_beta_epochs = 10, burnin_beta = 20,\n",
    "                       beta_max = 0.5, beta_min = 2.,\n",
    "                       period = 10):\n",
    "    \"\"\"\n",
    "    A funcion for changing the \"beta \" parameter during training according to a  cyclic annealing scheme, \n",
    "    a burn-in perior, and a cyclic annealing. Feel free to experiment with experimenting with it!\n",
    "    \n",
    "    \n",
    "    parameters:\n",
    "    \n",
    "      beta_0    : starting value for beta\n",
    "      init_beta_epochs : number of epochs that beta remains un-changed with beta-0\n",
    "      burnin_beta      : the epoch that the beta value is (linearly) ramped to, \n",
    "                         before the start of the annealing schedule\n",
    "      beta_max         : the max value for beta\n",
    "      beta_min         : the min value for beta\n",
    "      period           : how many epochs to go from min to max beta\n",
    "    \"\"\"\n",
    "    if epoch <init_beta_epochs:\n",
    "        return beta_0\n",
    "    \n",
    "    \n",
    "    if epoch>=init_beta_epochs and epoch < burnin_beta:\n",
    "        # ramp-up\n",
    "        return beta_0 + ((beta_max - beta_0)/(burnin_beta-init_beta_epochs)) * (epoch-init_beta_epochs)\n",
    "    \n",
    "    if epoch >= init_beta_epochs:\n",
    "        # sawtooth cycles - max/min beta:\n",
    "        db = (beta_max - beta_min)/period\n",
    "        return beta_max - (epoch - init_beta_epochs)%period * db\n",
    "        \n",
    "\n",
    "pplot.plot([beta_cyclic_anneal(ee) for ee in range(100)],'.-')\n",
    "pplot.title(\"Beta annealing schedule.\")\n",
    "pplot.grid()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Training loop:\n",
    "epochs = 400\n",
    "\n",
    "for e in tqdm(range(epochs)):\n",
    "    beta_curr = beta_cyclic_anneal(e)\n",
    "    with tf.GradientTape() as tape:\n",
    "        tot_loss,rec_loss, kl_loss,Xhat, z_out = eval_loss(xtrain, wtrain,beta = beta_curr)\n",
    "        grads = tape.gradient(tot_loss, cvae.vae_model.weights)\n",
    "        opt.apply_gradients(zip(grads,  cvae.vae_model.weights))\n",
    "        train_losses.append([np.mean(tot_loss.numpy().flatten()), np.mean(rec_loss), np.mean(kl_loss)])\n",
    "        \n",
    "\n",
    "    if e % 50 == 0:\n",
    "        test_tot_loss, test_rec_loss, test_kl_loss, test_xout, test_zout = eval_loss(xtest, wtest, beta = beta_curr)\n",
    "        test_losses.append([np.mean(test_tot_loss), np.mean(test_rec_loss), np.mean(test_kl_loss)])\n",
    "        pplot.subplot(2,1,1)\n",
    "        pplot.plot(test_xout[0][0::10].numpy(),  xtest[0::10],'.C0')\n",
    "        pplot.plot(xtest[0::10],xtest[0::10])\n",
    "        pplot.show()\n",
    "        \n",
    "        ii = 5\n",
    "        xout  = cvae.vae_model([xtest, wtest])\n",
    "        pplot.plot(tf.reduce_mean(xout[0],axis=0)[ii])\n",
    "        pplot.plot(xtest[ii])        \n",
    "        pplot.show()\n",
    "        pplot.pause(0.1)\n",
    "        "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Plot CVAE results\n",
    "The approximate posterior is replaced with the prior. If the KL part of the loss is small and the latent space can be captured with transformations of a spherical Gaussian, the prior should yield realistic samples from the distribution of the MC simulation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Evaluate the decoder using the spherical gaussian prior:\n",
    "wcond = wtest\n",
    "xvals = xtest\n",
    "\n",
    "niwae = cvae.cvae_params['n_iwae']\n",
    "z_prior = cvae.prior.sample(niwae)\n",
    "xout_from_prior = cvae.decoder([np.tile(z_prior[:,np.newaxis,:], [1,wcond.shape[0],1]), wcond])\n",
    "\n",
    "# Evaluate the VAE:\n",
    "xout = cvae.vae_model([xvals,wcond])\n",
    "def scale_to_plot(F_, mm2 = 0.0957):\n",
    "    ee = 0.2\n",
    "    F_ = dataset.unnormalize_DEL(F_)\n",
    "    F_ = (F_**ee )/mm2\n",
    "    return F_\n",
    "\n",
    "F_post = xout[0].numpy()#\n",
    "F_prior = xout_from_prior.numpy()\n",
    "Fsc_post = scale_to_plot(F_post)\n",
    "Fsc_prior = scale_to_plot(F_prior)\n",
    "Fact = scale_to_plot(xtest)\n",
    "Fact_train = scale_to_plot(xtrain)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#@title: a function to plot the CVAE samples together with the neighbors from the dataset:\n",
    "def plot_with_neighbs(wsp, ti, shear, xbounds = 3):\n",
    "    wunnorm = dataset.unnormalize_X(w_norm)\n",
    "\n",
    "    fig, ax = pplot.subplots(nrows=2,ncols=2, figsize = (20,10))\n",
    "    gs = ax[0,0].get_gridspec()\n",
    "    for ax_ in ax[:,0]:\n",
    "        ax_.remove()\n",
    "    axleft = fig.add_subplot(gs[:,0])\n",
    "\n",
    "    vv_unnorm = np.array([[ti,wsp,shear]])\n",
    "\n",
    "    vv = (vv_unnorm-dataset.data_input_mean) / dataset.data_input_std\n",
    "    w_unnorm = dataset.data_input\n",
    "\n",
    "    # Compute 10 samples from the CVAE using the prior and plot:\n",
    "    niwae = cvae.cvae_params['n_iwae']\n",
    "    z_prior = cvae.prior.sample(niwae)\n",
    "\n",
    "    wcond = np.array([])\n",
    "    xout_from_prior = cvae.decoder([np.tile(z_prior[:,np.newaxis,:], [1,vv.shape[0],1]), vv])\n",
    "    xout_from_prior\n",
    "    Fsc_prior_ = scale_to_plot(xout_from_prior)\n",
    "    for k in range(niwae):\n",
    "        pfv.make_plot(Fsc_prior_[k]*500000, 0, opacity=0.1, color=\"C0\",ax = axleft, tight = False)\n",
    "    \n",
    "    axleft.set_title(\"Fatigue estimates\\n ti/wsp/se: %2.1f/%2.2f/%2.1f\"%(wsp, ti, shear), fontsize = 20)\n",
    "    axleft.set_xlim([-xbounds, xbounds])\n",
    "    axleft.set_ylim([-xbounds, xbounds])\n",
    "    axleft.grid()\n",
    "    # pfv.make_plot(Fact*500000, ii, opacity=0.5, color=\"C1\", ax = ax_curr, tight = False)\n",
    "    # Find nearest neighbors from training set:\n",
    "    nneighbors = 2\n",
    "    inds_closest = np.argsort(np.sum(np.square(vv_unnorm-w_unnorm)*[100.,10.,1],1))[0:nneighbors]\n",
    "    # for kk in range\n",
    "    for n in inds_closest:\n",
    "        pfv.make_plot(scale_to_plot(x_norm[[n]])*500000, 0 , opacity=0.9, color=\"C1\",ax = axleft, tight = False)\n",
    "\n",
    "    # plot nearest neighbor:\n",
    "\n",
    "    ax[0,1].plot(w_unnorm[:,0], w_unnorm[:,1],'.', label = \"all data\")\n",
    "    ax[0,1].plot(w_unnorm[inds_closest,0], w_unnorm[inds_closest,1],'*', markersize = 15, label = \"neighbors\")\n",
    "    ax[0,1].plot(vv_unnorm[0,0], vv_unnorm[0,1],'*r', markersize = 15, label = \"CVAE conditioning\")\n",
    "    ax[0,1].set_xlabel(\"Ti[pct]\"); ax[0,1].set_ylabel(\"Wsp[m/s]\")\n",
    "    ax[0,1].legend(fontsize = 15)\n",
    "\n",
    "    ax[1,1].plot(w_unnorm[:,1], w_unnorm[:,2],'.')\n",
    "    ax[1,1].plot(w_unnorm[inds_closest,1], w_unnorm[inds_closest,2],'*',markersize = 15)\n",
    "    ax[1,1].plot(vv_unnorm[0,1], vv_unnorm[0,2],'*r', markersize = 15, label  = \"CVAE conditioning\")\n",
    "\n",
    "    ax[1,1].set_xlabel(\"Wsp[m/s]\"); ax[1,1].set_ylabel(\"shear exp\")\n",
    "    ax[0,1].grid(True)\n",
    "    ax[1,1].grid(True)\n",
    "    return fig"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from ipywidgets import FloatSlider, interact\n",
    "\n",
    "xbounds = 3\n",
    "\n",
    "figidx = 0\n",
    "@interact(wsp = FloatSlider(min = 0,max = 20, value = 5.90),\n",
    "          ti  = FloatSlider(min = 0, max = 0.4, value = 0.16, step = 0.4/20),\n",
    "          shear = FloatSlider(min =-1.64, max = 2.15, value = 0))\n",
    "def plot_with_inputs(wsp, ti, shear):\n",
    "    plot_with_neighbs(wsp, ti, shear)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#@title Run the following for creating an animation\n",
    "\n",
    "\n",
    "#plot_with_neighbs\n",
    "ngridpoints = 20\n",
    "ti0 = 0.05\n",
    "tigrid = np.linspace(ti0,0.15,ngridpoints)\n",
    "wsp0 = 5\n",
    "wspgrid = np.linspace(wsp0,20,ngridpoints)\n",
    "\n",
    "frame_idx = 0\n",
    "for kk in range(ngridpoints):\n",
    "    ti = tigrid[kk]\n",
    "    ff = plot_with_neighbs(wsp0, ti, 0.)\n",
    "    ff.savefig(\"anim_%03i.jpg\"%frame_idx)\n",
    "    frame_idx += 1\n",
    "\n",
    "for kk in reversed(range(ngridpoints)):\n",
    "    ti = tigrid[kk]\n",
    "    ff = plot_with_neighbs(wsp0, ti, 0.)\n",
    "    ff.savefig(\"anim_%03i.jpg\"%frame_idx)\n",
    "    frame_idx += 1\n",
    "\n",
    "    \n",
    "for kk in range(ngridpoints):\n",
    "    ff = plot_with_neighbs(wspgrid[kk], ti0, 0.)\n",
    "    ff.savefig(\"anim_%03i.jpg\"%frame_idx)\n",
    "    frame_idx += 1\n",
    "\n",
    "for kk in reversed(range(ngridpoints)):\n",
    "    ff = plot_with_neighbs(wspgrid[kk], ti0, 0.)\n",
    "    ff.savefig(\"anim_%03i.jpg\"%frame_idx)\n",
    "    frame_idx += 1\n"
   ]
  }
 ],
 "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.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}