plot-output.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import gc\n",
    "import xarray as xr\n",
    "import numpy as np\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.colors import LogNorm, LinearSegmentedColormap, colorConverter\n",
    "from matplotlib import rc\n",
    "from mpl_toolkits.axes_grid1.inset_locator import inset_axes\n",
    "from matplotlib.transforms import Bbox\n",
    "# from matplotlib.patches import FancyBboxPatch\n",
    "from matplotlib.font_manager import FontProperties\n",
    "\n",
    "import matplotlib as mpl\n",
    "from matplotlib.colors import ListedColormap\n",
    "from matplotlib.colors import BoundaryNorm\n",
    "\n",
    "from functions.mandyocIO import read_mandyoc_output, read_datasets, read_particle_path"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Initial setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "path = os.getcwd().split('/') # Get local file\n",
    "machine_path = '/'+path[1]+'/'+path[2] # Select home according to OS.\n",
    "\n",
    "# Path to the model directory\n",
    "model_path = \"/home/kugelblitz/runs/strak_71\"\n",
    "# model_path = \"/home/kugelblitz/runs/_strain-softening-0-a\"\n",
    "\n",
    "read_ascii = True # if False, it will skip the ascii reading and saving processes and it will read the datasets from previous runs\n",
    "\n",
    "plot_first = False # Generate and plot the first image\n",
    "plot_last = False # Generate and plot the last image\n",
    "\n",
    "plot_step = False\n",
    "step_num = 180\n",
    "\n",
    "plot_sample = True # Generate and plot key images across the simulation\n",
    "sample_size = 5 # Number of key images\n",
    "\n",
    "save_images = False # Generate and save all the possible images\n",
    "follow_particle = False # Stores x and z position of a particle through the simulation\n",
    "make_video = False # Make video witll all the possible imagesTrue\n",
    "make_gif = False # Make an animated gif with the generated video\n",
    "plot_temperature_anomaly = False # Plot the 'horizontal temperature anomaly' instead of 'temperature'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Datasets handeling\n",
    "## Model path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create the output directory to save the dataset\n",
    "output_path = os.path.join(model_path, \"_output\")\n",
    "if not os.path.isdir(output_path):\n",
    "    os.makedirs(output_path)\n",
    "\n",
    "model_name = os.path.split(model_path)[1]\n",
    "\n",
    "# datasets = (\"velocity\")\n",
    "datasets = (\"temperature\", \"density\", \"strain_rate\", \"velocity\")\n",
    "# datasets = (\"temperature\", \"viscosity\", \"strain_rate\", \"surface\", \"velocity\", \"density\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Read ascii outputs and save them as xarray.Datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Read data and convert them to xarray.Dataset\n",
    "if (read_ascii):\n",
    "    ds_data = read_mandyoc_output(\n",
    "        model_path,\n",
    "        datasets=datasets,\n",
    "        parameters_file=\"param.txt\",\n",
    "        skip=2,\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if (follow_particle):    \n",
    "    particle_x, particle_z, particle_ID = read_particle_path(\n",
    "        model_path,\n",
    "        (2_400.0e3, -100.0e3),\n",
    "        unit_number=3,\n",
    "        ncores=np.nan\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if (follow_particle):\n",
    "    x_disp = abs(particle_x[-1] - particle_x[0])\n",
    "    z_disp = abs(particle_z[-1] - particle_z[0])\n",
    "    print(\"X displacement:\", x_disp/1.0e3, \"[km]\")\n",
    "    print(\"Z displacement:\", z_disp/1.0e3, \"[km]\")\n",
    "    plt.figure(figsize=(x_disp*2/z_disp, 4), facecolor=\"white\")\n",
    "    plt.title(f'Particle {int(particle_ID)} trajectory')\n",
    "    plt.plot(particle_x/1.0e3, particle_z/1.0e3, color='grey')\n",
    "    plt.scatter(particle_x/1.0e3, particle_z/1.0e3, color='grey')\n",
    "    plt.scatter(particle_x[0]/1.0e3, particle_z[0]/1.0e3, color='blue', label='Begin', zorder=10) # zorder must be big to be on top of the other plots\n",
    "    plt.scatter(particle_x[-1]/1.0e3, particle_z[-1]/1.0e3, color='red', label='End', zorder=10)\n",
    "    plt.legend()\n",
    "    plt.xlabel(\"x [km]\")\n",
    "    plt.ylabel(\"z [km]\")\n",
    "    plt.savefig(f\"{output_path}/{model_name}_output_trajectory_{int(particle_ID)}.png\", dpi=300)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Read and merge saved datasets into a single dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "to_mm_yr = 365 * 24 * 60 * 60 * 1000.0\n",
    "max_velocity = 100_000_000.0 # [mm/yr]\n",
    "\n",
    "if (save_images) or (plot_first) or (plot_last) or (plot_step) or (plot_sample):\n",
    "    dataset = read_datasets(model_path, datasets)\n",
    "\n",
    "    # Normalize velocity values\n",
    "    if (\"velocity_x\" and \"velocity_z\") in dataset.data_vars:\n",
    "        v_max = np.max((dataset.velocity_x**2 + dataset.velocity_z**2)**(0.5))\n",
    "        print(f\"ini v_max: {float(v_max):.2e} [m/s]\")\n",
    "        # if (v_max > max_velocity / to_mm_yr):\n",
    "        #     v_max = max_velocity / to_mm_yr\n",
    "        #     print(f\"new v_max: {v_max:.2e} [m/s]\")\n",
    "        #     # dataset.velocity_x.where(dataset.velocity_x < v_max, 1.0e100)\n",
    "        #     dataset.velocity_x.values[dataset.velocity_x.values > v_max] = v_max\n",
    "        #     dataset.velocity_x.values[dataset.velocity_x.values < -v_max] = -v_max\n",
    "        #     dataset.velocity_z.values[dataset.velocity_z.values > v_max] = v_max\n",
    "        #     dataset.velocity_z.values[dataset.velocity_z.values < -v_max] = -v_max\n",
    "        #     v_max = v_max**2\n",
    "        dataset.velocity_x[:] = dataset.velocity_x[:] / v_max\n",
    "        dataset.velocity_z[:] = dataset.velocity_z[:] / v_max\n",
    "\n",
    "    print(dataset)\n",
    "\n",
    "if (plot_temperature_anomaly):\n",
    "    dataset['temperature'] = (dataset['temperature'] - dataset['temperature'].mean(dim='x')) / float(dataset.x[1])\n",
    "# time = 0\n",
    "# plt.figure(figsize=(20,4))\n",
    "# plt.imshow(dataset.velocity_x.isel(time=time).T[::-1])\n",
    "# plt.colorbar()\n",
    "# plt.figure(figsize=(20,4))\n",
    "# plt.imshow(dataset.velocity_z.isel(time=time).T[::-1])\n",
    "# plt.colorbar()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Output images cofig"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "colorbar_text_pos_left = 1.05\n",
    "colorbar_width = \"5%\"\n",
    "colorbar_text_labelpad = -90\n",
    "air_thickness = 44.0 # in km"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Steps configuration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "start = 0\n",
    "end = dataset.time.size\n",
    "step = 1\n",
    "\n",
    "if (plot_first) and (not plot_last) and (not plot_sample):\n",
    "    print(\"Plotting ONLY the first step\")\n",
    "    start = 0\n",
    "    end = 1\n",
    "    step = 1\n",
    "    \n",
    "if (not plot_first) and (plot_last) and (not plot_sample):\n",
    "    print(\"Plotting ONLY the last step\")\n",
    "    start = dataset.time.size - 1\n",
    "    end = dataset.time.size\n",
    "    step = 1 \n",
    "\n",
    "if (plot_first) and (plot_last) and (not plot_sample):\n",
    "    print(\"Plotting the first and last steps\")\n",
    "    start = 0\n",
    "    end = dataset.time.size \n",
    "    step = dataset.time.size - 1\n",
    "\n",
    "if (plot_step) and (not plot_sample):\n",
    "    print(f\"Plotting ONLY the step {step_num}\")\n",
    "    start = step_num\n",
    "    end = step_num + 1\n",
    "    step = 1\n",
    "\n",
    "times = np.arange(start, end, step)\n",
    "\n",
    "if (plot_sample):\n",
    "    print(\"Plotting sample steps\")\n",
    "    times = np.linspace(start, end-1, sample_size)\n",
    "    times = times.astype(np.int64)\n",
    "\n",
    "print(\"step stamps: \", times)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Generate and save output images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if (save_images) or (plot_first) or (plot_last) or (plot_step) or (plot_sample):\n",
    "    aux = 0\n",
    "    if (\"temperature\" or (\"velocity_x\" and \"velocity_z\")) in dataset.data_vars:\n",
    "        aux += 1\n",
    "        if \"temperature\" in dataset.data_vars:\n",
    "            if (plot_temperature_anomaly):\n",
    "                t_min = float(dataset.temperature.min())\n",
    "            else:\n",
    "                t_min = 0.0\n",
    "            t_max = np.max([float(dataset.temperature.max()), np.abs(t_min)])\n",
    "            if (plot_temperature_anomaly):\n",
    "                t_min = - np.max([float(dataset.temperature.max()), np.abs(t_min)])\n",
    "        if (\"velocity_x\" and \"velocity_x\") in dataset.data_vars:\n",
    "            num_vectors = 20\n",
    "            desired_mm_per_year_value = 50.0\n",
    "            print(f\"v_max: {(float(v_max)/to_mm_yr):.2f} [mm/yr]\")\n",
    "            v_key = desired_mm_per_year_value / to_mm_yr / v_max\n",
    "            v_scale = 0.75 # [tip:0.5] If set to <n>, max velocity arrows will have the size of dataset.x.max()/<n>, ie the bigger, the smaller will be the arrow\n",
    "    if \"strain\" in dataset.data_vars:\n",
    "        s_max = dataset.strain.max()\n",
    "        s_min = np.max([dataset.strain.min(), s_max/1.0E3])\n",
    "        print(\"strain: \", float(s_min), float(s_max))\n",
    "        aux += 1\n",
    "    if \"strain_rate\" in dataset.data_vars:\n",
    "        e_max = dataset.strain_rate.max()\n",
    "        e_min = np.max([dataset.strain_rate.min(), e_max/1.0E3])\n",
    "        print(\"strain_rate: \", float(e_min), float(e_max))\n",
    "        aux += 1\n",
    "    if \"viscosity\" in dataset.data_vars:\n",
    "        eta_min, eta_max = float(dataset.attrs[\"viscosity_min\"]), float(dataset.attrs[\"viscosity_max\"])\n",
    "        aux += 1\n",
    "    if \"surface\" in dataset.data_vars:\n",
    "        w_min, w_max = -10, 10 # dataset.surface.min()/1.0e3, dataset.surface.max()/1.0e3\n",
    "        aux += 1\n",
    "    if \"density\" in dataset.data_vars:\n",
    "        aux += 1\n",
    "\n",
    "    aux_size_x = np.round(float(abs(dataset.x[-1]/1.0e6)), 1)\n",
    "    aux_size_z = np.round(float(abs(dataset.z[0]/1.0e6)), 1)\n",
    "    \n",
    "    aux_size_x = np.max([aux_size_x, 1])\n",
    "    aux_size_z = np.max([aux_size_z, 1])\n",
    "    print(aux_size_x, aux_size_z)\n",
    "    \n",
    "    # Create the color map to plot the density\n",
    "    cr = 255.0\n",
    "    color_air = (255.0 / cr, 255.0 / cr, 255.0 / cr)\n",
    "    color_continental_crust = (240.0 / cr, 209.0 / cr, 188.0 / cr)\n",
    "    color_continental_mantle = (228.0 / cr, 156.0 / cr, 124.0 / cr)\n",
    "    color_oceanic_crust = (207.0 / cr, 226.0 / cr, 205.0 / cr)\n",
    "    color_oceanic_mantle = (155.0 / cr, 194.0 / cr, 155.0 / cr)\n",
    "    color_mantle = (155.0 / cr, 0.0 / cr, 155.0 / cr)\n",
    "    colors = [\n",
    "        color_continental_crust, \n",
    "        color_continental_mantle, \n",
    "        color_oceanic_crust,\n",
    "    ]\n",
    "    levels = [\n",
    "        2750.0, # density will show up with <color_continental_crust>\n",
    "        3354.0, # density will show up with <color_continental_mantle>\n",
    "        3378.0, # density will show up with <color_oceanic_crust>\n",
    "        3400.0  # density will show up with <color_oceanic_mantle>\n",
    "    ]\n",
    "    cmap_de = ListedColormap(colors)\n",
    "    cmap_de.set_under(color_air)\n",
    "    cmap_de.set_over(color_oceanic_mantle)\n",
    "    norm = BoundaryNorm(levels, ncolors=cmap_de.N)\n",
    "\n",
    "    # Create the color map to plot the strain rate\n",
    "    color_0 = colorConverter.to_rgba('white', alpha = 0)\n",
    "    color_1 = colorConverter.to_rgba('purple', alpha = 1)\n",
    "    color_2 = colorConverter.to_rgba('black', alpha = 1)\n",
    "    cmap_sr = LinearSegmentedColormap.from_list('cmap_sr', [color_0, color_1, color_2], 512)\n",
    "\n",
    "    SMALL_SIZE, MEDIUM_SIZE, BIGGER_SIZE = 16, 16, 16\n",
    "\n",
    "    plt.rc('font', size=SMALL_SIZE)          # controls default text sizes\n",
    "    plt.rc('axes', titlesize=SMALL_SIZE)     # fontsize of the axes title\n",
    "    plt.rc('axes', labelsize=MEDIUM_SIZE)    # fontsize of the x and y labels\n",
    "    plt.rc('xtick', labelsize=SMALL_SIZE)    # fontsize of the tick labels\n",
    "    plt.rc('ytick', labelsize=SMALL_SIZE)    # fontsize of the tick labels\n",
    "    plt.rc('legend', fontsize=SMALL_SIZE)    # legend fontsize\n",
    "    plt.rc('figure', titlesize=BIGGER_SIZE)  # fontsize of the figure title\n",
    "\n",
    "    hspace = 0.2\n",
    "\n",
    "    # Corrects an unintended fontsize on the first run for the quiverkey\n",
    "    quiverkey_font = FontProperties()\n",
    "    quiverkey_font.set_size(MEDIUM_SIZE)\n",
    "\n",
    "    # for i in range(start, end, step):\n",
    "    for i in times:\n",
    "        subplot_num = 0\n",
    "        per = np.round(100*(i+1-start)/(end-start), 2)\n",
    "        print(f'output {i}/{end-1}, {per:.2f}%', end='\\r')\n",
    "\n",
    "        data = dataset.isel(time=i)\n",
    "        \n",
    "        if (i==start):\n",
    "            z_text_pos = float(data.z.mean()) / 1.0E3 + (1.0 * air_thickness)\n",
    "            x_text_pos = (1.0 + 0.5* (colorbar_text_pos_left - 1.0)) * float(data.x.max()) / 1.0E3\n",
    "            extent = [data.x.min()/1.0e3, \n",
    "                      data.x.max()/1.0e3, \n",
    "                      data.z.min()/1.0e3 + air_thickness, \n",
    "                      data.z.max()/1.0e3 + air_thickness] \n",
    "\n",
    "        fig, axs = plt.subplots(aux, 1, figsize=(3*aux_size_x, aux*4*aux_size_z), sharex=True, facecolor=\"white\")\n",
    "        fig.subplots_adjust(hspace=hspace)\n",
    "\n",
    "        if \"surface\" in dataset.data_vars:\n",
    "            if (i == 0) or (start != 0): \n",
    "                data_0 = dataset.isel(time=0)\n",
    "                h_air = np.round(np.mean(data_0.surface/1.0e3)) # Estimates air layer thickness based on the first step\n",
    "            im = axs[subplot_num].hlines(0, data.surface.x[0]/1.0e3, data.surface.x[-1]/1.0e3, linestyle=\"solid\", color=\"grey\")\n",
    "            im = axs[subplot_num].plot(data.surface.x/1.0e3, data.surface/1.0e3 - h_air)\n",
    "            axs[subplot_num].set_ylim(-10, 10)\n",
    "            axs[subplot_num].set_ylim(w_min, w_max)\n",
    "            axs[subplot_num].set_ylabel(\"Elevation [km]\")\n",
    "            surface_subplot_num = subplot_num\n",
    "            imshow_box = np.array(axs[subplot_num].get_position()) # Backup Bbox to be used if no other data_vars is in the xarray.Dataset\n",
    "            subplot_num += 1\n",
    "\n",
    "        if \"viscosity\" in dataset.data_vars:\n",
    "            im = axs[subplot_num].imshow(data.viscosity.T[::-1], \n",
    "                                         extent=extent, \n",
    "                                         norm=LogNorm(vmin=eta_min, vmax=eta_max),\n",
    "                                         cmap=\"viridis\",\n",
    "                                         interpolation='none')\n",
    "            cax = inset_axes(axs[subplot_num], \n",
    "                             width=colorbar_width, \n",
    "                             height=\"100%\", \n",
    "                             loc='center left', \n",
    "                             bbox_to_anchor=(colorbar_text_pos_left, 0., 0.3, 1), \n",
    "                             bbox_transform=axs[subplot_num].transAxes, \n",
    "                             borderpad=0)\n",
    "            axs[subplot_num].set_aspect(\"equal\")\n",
    "            axs[subplot_num].set_ylabel(\"Depth [km]\")\n",
    "            cbar = fig.colorbar(im, cax=cax)\n",
    "            axs[subplot_num].text(x_text_pos, \n",
    "                                  z_text_pos, \n",
    "                                  r'Viscosity, $\\eta$ [Pa.s]', \n",
    "                                  rotation=90, \n",
    "                                  horizontalalignment='center', \n",
    "                                  verticalalignment='center', \n",
    "                                  zorder=9999,\n",
    "                                  fontsize=SMALL_SIZE)\n",
    "            imshow_box = np.array(axs[subplot_num].get_position())\n",
    "            if (follow_particle):\n",
    "                im = axs[subplot_num].plot(particle_x/1.0e3, particle_z/1.0e3, color=\"red\")\n",
    "                im = axs[subplot_num].scatter(particle_x[i]/1.0e3, particle_z[i]/1.0e3, color=\"blue\", zorder=10)\n",
    "            subplot_num += 1\n",
    "            \n",
    "        if \"density\" in dataset.data_vars:\n",
    "            im = axs[subplot_num].imshow(data.density.T[::-1], \n",
    "                                         extent=extent, \n",
    "                                         cmap=cmap_de,\n",
    "                                         norm=norm,\n",
    "                                         interpolation='none')\n",
    "            cax = inset_axes(axs[subplot_num], \n",
    "                             width=colorbar_width, \n",
    "                             height=\"100%\", \n",
    "                             loc='center left', \n",
    "                             bbox_to_anchor=(colorbar_text_pos_left, 0., 0.3, 1), \n",
    "                             bbox_transform=axs[subplot_num].transAxes, \n",
    "                             borderpad=0)\n",
    "            axs[subplot_num].set_aspect(\"equal\")\n",
    "            axs[subplot_num].set_ylabel(\"Depth [km]\")\n",
    "            cbar = fig.colorbar(im, cax=cax, extend='both')\n",
    "            axs[subplot_num].text(x_text_pos, \n",
    "                                  z_text_pos, \n",
    "                                  r'Density, $\\rho$ [kg/m$^3$]', \n",
    "                                  rotation=90, \n",
    "                                  horizontalalignment='center', \n",
    "                                  verticalalignment='center', \n",
    "                                  zorder=9999,\n",
    "                                  fontsize=SMALL_SIZE)\n",
    "            imshow_box = np.array(axs[subplot_num].get_position())\n",
    "            subplot_num += 1\n",
    "\n",
    "        if (\"temperature\" or (\"velocity_x\" and \"velocity_x\")) in dataset.data_vars:\n",
    "            if \"temperature\" in dataset.data_vars:\n",
    "                im = axs[subplot_num].imshow(data.temperature.T[::-1], \n",
    "                                             extent=extent, \n",
    "                                             vmin=t_min, vmax=t_max, \n",
    "                                             cmap=\"RdYlBu_r\",\n",
    "                                             interpolation='none')\n",
    "                cax = inset_axes(axs[subplot_num], \n",
    "                                 width=colorbar_width, \n",
    "                                 height=\"100%\", \n",
    "                                 loc='center left', \n",
    "                                 bbox_to_anchor=(colorbar_text_pos_left, 0., 0.3, 1), \n",
    "                                 bbox_transform=axs[subplot_num].transAxes, \n",
    "                                 borderpad=0)\n",
    "                cbar = fig.colorbar(im, cax=cax)\n",
    "                if (plot_temperature_anomaly):\n",
    "                    txt_label = r'$\\delta_{x}$T [$^{\\circ}$C/km]'\n",
    "                else:\n",
    "                    txt_label = r'Temperature, T [$^{\\circ}$C]'\n",
    "                axs[subplot_num].text(x_text_pos, \n",
    "                                  z_text_pos, \n",
    "                                  txt_label, \n",
    "                                  rotation=90, \n",
    "                                  horizontalalignment='center', \n",
    "                                  verticalalignment='center', \n",
    "                                  zorder=9999,\n",
    "                                  fontsize=SMALL_SIZE)\n",
    "            if (\"velocity_x\" and \"velocity_x\") in dataset.data_vars:\n",
    "                vel_aux = data[dict(x=slice(None, None, num_vectors), z=slice(None, None, num_vectors))]\n",
    "                v_label = '$\\overrightarrow{v} =$' + f'{np.round(float(v_key*v_max*to_mm_yr))} [mm/yr]'\n",
    "                im = axs[subplot_num].quiver(vel_aux.x/1.0e3, \n",
    "                                             vel_aux.z/1.0e3 + air_thickness, \n",
    "                                             vel_aux.velocity_x.values.T, \n",
    "                                             vel_aux.velocity_z.values.T, \n",
    "                                             scale=v_scale)\n",
    "                # arrow_size = 0.085 * dataset.x.max()/1.0e3\n",
    "                # dd = 100\n",
    "                # dx, dz = 1150 + arrow_size, 2 * dd\n",
    "                # x0, z0 = dataset.x.max()/1.0e3 - dx - dd,  dataset.z.min()/1.0e3 + dd\n",
    "                im = axs[subplot_num].quiverkey(im, \n",
    "                                                X=0.5,#(x0+arrow_size)/(dataset.x.max()/1.0e3),\n",
    "                                                Y=-0.1,#1-(z0+dd)/(dataset.z.min()/1.0e3),\n",
    "                                                U=v_key,\n",
    "                                                label=v_label, \n",
    "                                                labelpos='E', \n",
    "                                                fontproperties=quiverkey_font) # fontproperties corrects an unintended behaviour where the incorrect fontsize was used during 1st run\n",
    "                # im = axs[subplot_num].add_patch(FancyBboxPatch((x0, z0), dx, dz, boxstyle='round, rounding_size=25', facecolor = 'white'))\n",
    "            # axs[subplot_num].set_aspect(\"equal\")\n",
    "            axs[subplot_num].set_aspect(\"equal\")\n",
    "            axs[subplot_num].set_ylabel(\"Depth [km]\")\n",
    "            imshow_box = np.array(axs[subplot_num].get_position())\n",
    "            subplot_num += 1\n",
    "            \n",
    "        if \"strain\" in dataset.data_vars:\n",
    "            palette = plt.cm.viridis.with_extremes(over=plt.cm.viridis.colors[-1], under=plt.cm.viridis.colors[0], bad=plt.cm.viridis.colors[0])\n",
    "            im = axs[subplot_num].imshow(data.strain.T[::-1],\n",
    "                                         extent=extent, \n",
    "                                         norm=LogNorm(vmin=s_min, vmax=s_max), \n",
    "                                         alpha=1.0,\n",
    "                                         cmap=palette,\n",
    "                                         interpolation='none')\n",
    "            cax = inset_axes(axs[subplot_num], \n",
    "                             width=colorbar_width, \n",
    "                             height=\"100%\", \n",
    "                             loc='center left', \n",
    "                             bbox_to_anchor=(colorbar_text_pos_left, 0., 0.3, 1), \n",
    "                             bbox_transform=axs[subplot_num].transAxes, \n",
    "                             borderpad=0)\n",
    "            axs[subplot_num].set_aspect(\"equal\")\n",
    "            axs[subplot_num].set_ylabel(\"Depth [km]\")\n",
    "            cbar = fig.colorbar(im, cax=cax, extend='min')\n",
    "            axs[subplot_num].text(x_text_pos, \n",
    "                                  z_text_pos, \n",
    "                                  r'Strain, $\\epsilon$', \n",
    "                                  rotation=90, \n",
    "                                  horizontalalignment='center', \n",
    "                                  verticalalignment='center', \n",
    "                                  zorder=9999,\n",
    "                                  fontsize=SMALL_SIZE)\n",
    "            imshow_box = np.array(axs[subplot_num].get_position())\n",
    "            subplot_num += 1\n",
    "            \n",
    "        if \"strain_rate\" in dataset.data_vars:\n",
    "            if \"density\" in dataset.data_vars:\n",
    "                im = axs[subplot_num].imshow(data.density.T[::-1], \n",
    "                                            extent=extent, \n",
    "                                            cmap=cmap_de,\n",
    "                                            norm=norm,\n",
    "                                            interpolation='none',\n",
    "                                            alpha=0.5)\n",
    "            im = axs[subplot_num].imshow(data.strain_rate.T[::-1],  \n",
    "                                         extent=extent, \n",
    "                                         norm=LogNorm(vmin=e_min, vmax=e_max), \n",
    "                                         alpha=1.0,\n",
    "                                         cmap=cmap_sr,\n",
    "                                         interpolation='none')\n",
    "            cax = inset_axes(axs[subplot_num], \n",
    "                             width=colorbar_width, \n",
    "                             height=\"100%\", \n",
    "                             loc='center left', \n",
    "                             bbox_to_anchor=(colorbar_text_pos_left, 0., 0.3, 1), \n",
    "                             bbox_transform=axs[subplot_num].transAxes, \n",
    "                             borderpad=0)\n",
    "            axs[subplot_num].set_aspect(\"equal\")\n",
    "            axs[subplot_num].set_ylabel(\"Depth [km]\")\n",
    "            cbar = fig.colorbar(im, cax=cax, extend='min')\n",
    "            axs[subplot_num].text(x_text_pos, \n",
    "                                  z_text_pos, \n",
    "                                  r'Strain rate, $\\dot{\\epsilon}$ [s$^{-1}$]', \n",
    "                                  rotation=90, \n",
    "                                  horizontalalignment='center', \n",
    "                                  verticalalignment='center', \n",
    "                                  zorder=9999,\n",
    "                                  fontsize=SMALL_SIZE)\n",
    "            imshow_box = np.array(axs[subplot_num].get_position())\n",
    "            subplot_num += 1\n",
    "\n",
    "        # Fix surface plot aspect ratio and position\n",
    "        if \"surface\" in dataset.data_vars:\n",
    "            axs[surface_subplot_num].set_aspect(dataset.surface.x.max()/1.0e3/120)\n",
    "            plot_box = np.array(axs[surface_subplot_num].get_position())    \n",
    "            new_plot_box = Bbox([[imshow_box[0, 0], plot_box[0, 1]], [imshow_box[1, 0], plot_box[1, 1]]])\n",
    "            axs[surface_subplot_num].set_position(new_plot_box)   \n",
    "\n",
    "        axs[subplot_num-1].set_xlabel(\"Length [km]\")\n",
    "        plt.suptitle(\"time = {:.2f} My, step = {}\".format(np.round(data.time.item(), 2), data.step.item()), ha='center', y=0.9, x=0.5)\n",
    "        fig.align_ylabels(axs[:])\n",
    "        plt.savefig(f\"{output_path}/{model_name}_output_{str(i).zfill(6)}.png\", dpi=300)\n",
    "        if (plot_first==False):\n",
    "            plt.close('all')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Make video with output images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "line = f\"ffmpeg -r 7 -i {model_path}/_output/{model_name}_output_%06d.png -c:v libx264 -vf fps=25 -pix_fmt yuv420p {model_path}/_output/{model_name}.mp4\"\n",
    "print(line)\n",
    "make_video = True\n",
    "if (make_video):\n",
    "    !rm {model_path}/_output/{model_name}.mp4\n",
    "    !ffmpeg -r 7 -i {model_path}/_output/{model_name}_output_%06d.png -c:v libx264 -vf fps=25 -pix_fmt yuv420p {model_path}/_output/{model_name}.mp4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "line = 'ffmpeg -ss 0 -t 10 -i {model_path}/_output/{model_name}.mp4 -vf \"fps=10,scale=1080:-1:flags=lanczos,split[s0][s1];[s0]palettegen[p];[s1][p]paletteuse\" -loop 0 {model_path}/_output/{model_name}.gif'\n",
    "print(line)\n",
    "if (make_gif):\n",
    "    !rm {model_path}/_output/{model_name}.gif\n",
    "    !ffmpeg -ss 0 -t 10 -i {model_path}/_output/{model_name}.mp4 -vf \"fps=100,scale=1080:-1:flags=lanczos,split[s0][s1];[s0]palettegen[p];[s1][p]paletteuse\" -loop 0 {model_path}/_output/{model_name}.gif"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.9.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}