Post processing script for tracer path plots. Script is for processing the output of Cenki et al. models.

PathPostProcessingScipt.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Underworld tracer particle path post processing script."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This post-processing script transforms the .h5 output files for passive tracers into Pressure Temperature Strain Rate Material... paths for these particles.\n",
    "Thanks to J. Giordani, P. Vernant, P. Rey and L. Mondy for their helpful advices. \n",
    "\n",
    "For Cenki et al., 2022:\n",
    "https://github.com/underworld-community/cenki-et-al-UHT-granulitic-terranes\n",
    "\n",
    "For Cenki-Tok et al., 2020\n",
    "https://www.underworldcode.org/articles/publication-news/ and\n",
    "https://github.com/underworld-community/cenki-tok-et-al-crustal-roots-in-stable-continents\n",
    "\n",
    "See the User Parameters cell below to set variables that effect the results of this scripts, e.g. 'h5_path', 'tracer_name', etc."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import glob, os, re\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import h5py "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Install seaborn (viz package) if required and setup default visuals\n",
    "try:\n",
    "    import seaborn as sns\n",
    "except ModuleNotFoundError:\n",
    "    import sys, subprocess\n",
    "    subprocess.check_call([sys.executable, \"-m\", \"pip\", \"install\", 'seaborn'])\n",
    "    import seaborn as sns\n",
    "    \n",
    "# Set a dark grey background for the plots and white axes\n",
    "sns.set(rc={'axes.facecolor':'lightgrey', 'figure.facecolor':'white'})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check the path to your csv file directory\n",
    "os.getcwd()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "### User Parameters\n",
    "# The following must be set by the user\n",
    "h5_path     = \"./outputs_script-test_v2-12_Fst2RHP_3\"\n",
    "# h5_path     = \"./outputs_script-G47301-Cenki-Tok-etal_ModelFig3B_1/\"\n",
    "\n",
    "tracer_name = \"Lag_grid\"\n",
    "\n",
    "##ACTIVATE FOR MODELS IN CENKI ET AL 2022\n",
    "## material labels and their corresponding id ranges. \n",
    "## note the range are float because of the old implementation of UWGeo. They can be changed later to integers.\n",
    "binsM   = [1, 1.91, 2.87, 9, 10, 11]\n",
    "labelsM = ['air', 'sediments', 'crust', 'granulite', 'upper mantle']\n",
    "\n",
    "##ACTIVATE FOR MODELS IN CENKI-TOK ET AL 2020\n",
    "#binsM = [1, 1.88, 2.74, 3.60, 4.46, 5.35]\n",
    "#labelsM = ['air', 'continental crust', 'retrogressed crust', 'granulite', 'upper mantle']\n",
    "\n",
    "\n",
    "##ACTIVATE FOR MODELS IN CENKI ET AL 2022\n",
    "color_dictM = dict({'air':'lightgrey',\n",
    "                   'crust':'orange',\n",
    "                   'sediments':'peachpuff',\n",
    "                   'granulite':'saddlebrown',\n",
    "                   'upper mantle':'dimgrey'\n",
    "                   })\n",
    "                \n",
    "##ACTIVATE FOR MODELS IN CENKI-TOK ET AL 2020\n",
    "#color_dictM = dict({'air':'lightgrey',\n",
    "#                   'continental crust':'orange',\n",
    "#                  'retrogressed crust':'aqua',\n",
    "#                   'granulite':'saddlebrown',\n",
    "#                  'upper mantle':'dimgrey',\n",
    "#                   })\n",
    "\n",
    "\n",
    "## To enable per particle plots, toggle the following to True/False\n",
    "\n",
    "pp_temperature_depth      = False\n",
    "pp_viscosity_strainrate   = False\n",
    "pp_temperature_strainrate = False\n",
    "pp_temperature_viscosity  = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "try: \n",
    "    os.chdir(h5_path)\n",
    "except FileNotFoundError:\n",
    "    print(\"Can't find the file directory. Change the variable 'h5_path' above.\")\n",
    "    raise\n",
    "except NameError:\n",
    "    print(\"No variable 'h5_path' defined, you must define it before this cell.\")\n",
    "    raise"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# glob search to get all the time steps available for swarm named, 'tracer_name'\n",
    "glob_string   = tracer_name+'-*.xdmf'\n",
    "\n",
    "# order the timestep files assuming file name appears as: 'tracer_name-xxx.xdmf'\n",
    "# where xxx is a digit for the time step\n",
    "key_gen       = lambda x : re.search('(\\d+)\\.xdmf$', x).group(1).zfill(5)\n",
    "ordered_files = sorted(glob.glob(glob_string), key=key_gen)\n",
    "\n",
    "# check if glob returns 0, if so error\n",
    "if len(ordered_files) == 0:\n",
    "    ss = f\"Error: Can't find tracer files with name '{glob_string}' in directory '{h5_path}'\\n\" + \\\n",
    "         f\"Update the 'tracer_name' or 'h5_path' variables in the python script\"\n",
    "    raise RuntimeError(ss)\n",
    "    \n",
    "# get the min max timesteps found from the glob\n",
    "steps_az = ( re.search( '(\\d+)\\.xdmf$', ordered_files[0]).group(1),\n",
    "             re.search( '(\\d+)\\.xdmf$', ordered_files[-1]).group(1) )\n",
    "steps_az = [int(x) if x.isdigit() else None for x in steps_az]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if None in steps_az:\n",
    "    raise RuntimeError(\"Something went wrong processing the timesteps\")\n",
    "\n",
    "tSteps = steps_az[1] - steps_az[0]\n",
    "    \n",
    "# get all the variables files at time steps_az[0] and grab their file names\n",
    "# assume these names are the variable names.\n",
    "files = sorted(glob.glob(tracer_name+'*-'+str(steps_az[0])+'.h5'))\n",
    "var_names = []\n",
    "\n",
    "for name in files:\n",
    "    name = re.findall(tracer_name+'_(.+)-\\d',name)\n",
    "    if name == []: # in this case assume file has the x,y positions of the swarm\n",
    "        var_names.append('Points:0')\n",
    "        var_names.append('Points:1')\n",
    "    else:\n",
    "        var_names.append(name[0])\n",
    "\n",
    "## Previously the var_names was hardcoded, but this has been disabled for the above loop.  \n",
    "## col_values = ['Points:0', 'Points:1', 'global_index','laggridmaterialfield', 'laggridmeltfield', 'laggridpress', 'laggridstrainrate','laggridtemp','laggridtimefield','laggridviscosity']\n",
    "## var_name = col_values\n",
    "        \n",
    "# make report\n",
    "report = \"*******************************************\\n\" + \\\n",
    "         \"Report from the output directory inspection\\n\" + \\\n",
    "         \"*******************************************\\n\" + \\\n",
    "         f\"In the directory:\\n {h5_path}\\n\" + \\\n",
    "         f\"The Swarm '{tracer_name}' has {tSteps} time steps available, starting from time step {steps_az[0]}\\n\" + \\\n",
    "         f\"The available variables names are:\\n\" + \\\n",
    "         f\"{var_names}\\n\"\n",
    "print(report)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# read the h5 files and fill in the matrix\n",
    "for i in np.arange(steps_az[0],steps_az[1]) :\n",
    "    files = sorted(glob.glob(tracer_name+'*-'+str(i)+'.h5'))\n",
    "    # print('\\nprocessing time step: ',i)\n",
    "    j = 0\n",
    "    for file in files :\n",
    "        # print('processing:',file)\n",
    "        f = h5py.File(file, 'r')\n",
    "        dset = np.array(f[list(f.keys())[0]])\n",
    "        f.close()\n",
    "        if j == 0 : # create a matrix to store all the data from the time step \n",
    "            nelem = len(dset) # number of elements\n",
    "            all_data_tstp = np.zeros((nelem,10))  # create the matrix\n",
    "        all_data_tstp[0:nelem,j:j+dset.shape[1]] = dset\n",
    "        j = j+dset.shape[1]\n",
    "        \n",
    "    if i == 0 :\n",
    "        all_data = all_data_tstp\n",
    "    else:\n",
    "        all_data = np.vstack((all_data,all_data_tstp))\n",
    "    \n",
    "print(f'\\nProcessed {steps_az[1]-steps_az[0]} time steps\\n')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export the matrix to a dataframe\n",
    "frame = pd.DataFrame(data = all_data, columns = var_names)\n",
    "print (frame)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "frame.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#ADD A TEMPERATURE COLUMN IN DEG C TO THE DATAFRAME\n",
    "frame['laggridtemp'] - 273\n",
    "frame.insert(2, \"temperature_degC\", frame['laggridtemp'] - 273, True)\n",
    "print(frame)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "frame['binnedM'] = pd.cut(frame['laggridmaterialfield'], bins=binsM, labels=labelsM)\n",
    "\n",
    "binsSR = [1e-16, 9e-15, 1e-14, 1e-13] \n",
    "labelsSR = ['lowSR', 'threshold', 'highSR']\n",
    "frame['binnedSR'] = pd.cut(frame['laggridstrainrate'], bins=binsSR, labels=labelsSR)\n",
    "\n",
    "binsMELT = [0, 0.0001, 1] \n",
    "labelsMELT = ['no melt', 'melt' ]\n",
    "frame['binnedMELT'] = pd.cut(frame['laggridmeltfield'], bins=binsMELT, labels=labelsMELT)\n",
    "\n",
    "print(frame)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Create an array with the colors you want to use\n",
    "\n",
    "markers_dictSR = dict({'lowSR':'.',\n",
    "                   'threshold':'o',\n",
    "                   'highSR':'*'})\n",
    "\n",
    "markers_dictMELT = dict({'no melt':'o',\n",
    "                   'melt':'.'})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# find out all particles ID in the global index column\n",
    "particules_ID = frame.global_index.unique()\n",
    "len(particules_ID)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_particles = len(particules_ID)\n",
    "num_particles"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "particules_dict = {nameid: frame.loc[frame['global_index'] == nameid] for nameid in particules_ID}\n",
    "\n",
    "# Uncomment to see the global indcies of the \n",
    "# list(particules_dict)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Temperature - Depth plot for all particles"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sns.set_style(\"ticks\")\n",
    "fig = sns.relplot(x=\"laggridtemp\", y=\"Points:1\", color=\"lightgrey\", edgecolor=\"none\", data=frame);\n",
    "fig.set(xlim=(273, 1273))\n",
    "fig.set(ylim=(0, -65))\n",
    "fig.savefig(\"Temperature_Depth_plot_all_particles.eps\",dpi=150)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Temperature - Depth plot for all particles colored according to Material "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sns.set_style(\"ticks\")\n",
    "fig = sns.relplot(x=\"laggridtemp\", y=\"Points:1\", hue=\"binnedM\", palette=color_dictM, edgecolor=\"none\", data=frame);\n",
    "fig.set(xlim=(273, 1273))\n",
    "fig.set(ylim=(0, -65))\n",
    "fig.savefig(\"Temperature_Depth_plot_all_particles_coloredMATERIAL.eps\",dpi=150)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Temperature - Depth plot for each particle colored according to Material"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if pp_temperature_depth:\n",
    "    for m in particules_dict:\n",
    "    \n",
    "        M_single_particle_through_time = []\n",
    "        M_single_particle_through_time = frame.loc[frame['global_index'] == m]\n",
    "        sns.set_style(\"ticks\")\n",
    "        fig = sns.relplot(x=\"laggridtemp\", y=\"Points:1\", hue=\"binnedM\", palette=color_dictM, edgecolor=\"none\", data=M_single_particle_through_time);\n",
    "        fig.set(xlim=(273, 1273))\n",
    "        fig.set(ylim=(0, -65))\n",
    "        fig.savefig(\"Temperature_Depth_plot_each_particle_coloredMATERIAL.%d.eps\"%m,dpi=150)\n",
    "        plt.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Viscosity - StrainRate plot for all particles colored according to Material"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = sns.relplot(x=\"laggridviscosity\", y=\"laggridstrainrate\", hue=\"binnedM\", palette=color_dictM, edgecolor=\"none\", data=frame);\n",
    "fig.set(xscale=\"log\")\n",
    "fig.set(yscale=\"log\")\n",
    "fig.savefig(\"Viscosity_StrainRate_plot_all_particles_coloredMATERIAL.eps\",dpi=150)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Viscosity - StrainRate plot for each particle colored according to Material"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if pp_viscosity_strainrate:\n",
    "    for m in particules_dict:\n",
    "\n",
    "        M_single_particle_through_time = []\n",
    "        M_single_particle_through_time = frame.loc[frame['global_index'] == m]\n",
    "        sns.set_style(\"ticks\")\n",
    "        fig = sns.relplot(x=\"laggridviscosity\", y=\"laggridstrainrate\", hue=\"binnedM\", palette=color_dictM, edgecolor=\"none\", data=M_single_particle_through_time);\n",
    "        fig.set(xscale=\"log\")\n",
    "        fig.set(yscale=\"log\")\n",
    "        fig.savefig(\"Viscosity_StrainRate_plot_each_particle_coloredMATERIAL.%d.eps\"%m,dpi=150)\n",
    "        plt.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Temperature - StrainRate plot for all particles colored according to Material"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sns.set_style(\"ticks\")\n",
    "fig = sns.relplot(x=\"laggridtemp\", y=\"laggridstrainrate\", hue=\"binnedM\", palette=color_dictM, edgecolor=\"none\", data=frame);\n",
    "fig.set(xlim=(273, 1273))\n",
    "fig.set(yscale=\"log\")\n",
    "fig.savefig(\"Temperature_StrainRate_plot_all_particles_coloredMATERIAL.eps\",dpi=150)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Temperature - StrainRate plot for each particle colored according to Material"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if pp_temperature_strainrate:\n",
    "    for m in particules_dict:\n",
    "\n",
    "        M_single_particle_through_time = []\n",
    "        M_single_particle_through_time = frame.loc[frame['global_index'] == m]\n",
    "        sns.set_style(\"ticks\")\n",
    "        fig = sns.relplot(x=\"laggridtemp\", y=\"laggridstrainrate\", hue=\"binnedM\", palette=color_dictM, edgecolor=\"none\", data=M_single_particle_through_time);\n",
    "        fig.set(xlim=(273, 1273))\n",
    "        fig.set(yscale=\"log\")\n",
    "        fig.savefig(\"Temperature_StrainRate_plot_each_particle_coloredMATERIAL.%d.eps\"%m,dpi=150)\n",
    "        plt.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Temperature - Viscosity plot for all particles colored according to Material "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sns.set_style(\"ticks\")\n",
    "fig = sns.relplot(x=\"laggridtemp\", y=\"laggridviscosity\", hue=\"binnedM\", palette=color_dictM, edgecolor=\"none\", data=frame);\n",
    "fig.set(xlim=(273, 1273))\n",
    "fig.set(yscale=\"log\")\n",
    "fig.savefig(\"Temperature_Viscosity_plot_all_particles_coloredMATERIAL.eps\",dpi=150)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Temperature - Viscosity plot for each particle colored according to Material"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if pp_temperature_viscosity:\n",
    "    for m in particules_dict:\n",
    "\n",
    "        M_single_particle_through_time = []\n",
    "        M_single_particle_through_time = frame.loc[frame['global_index'] == m]\n",
    "        sns.set_style(\"ticks\")\n",
    "        fig = sns.relplot(x=\"laggridtemp\", y=\"laggridviscosity\", hue=\"binnedM\", palette=color_dictM, edgecolor=\"none\", data=M_single_particle_through_time);\n",
    "        fig.set(xlim=(273, 1273))\n",
    "        fig.set(yscale=\"log\")\n",
    "        fig.savefig(\"Temperature_Viscosity_plot_each_particle_coloredMATERIAL.%d.eps\"%m,dpi=150)\n",
    "        plt.close()"
   ]
  }
 ],
 "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.10.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}