{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Process ASPECT Exhumation Model Outputs\n",
"This notebook demonstrates how to use GDTchron to predict AHe, AFT, and ZHe ages from the results of the accompanying ASPECT model of simplified exhumation defined in `exhumation_box.prm` and run by `run_aspect_model.ipynb`."
]
},
{
"cell_type": "code",
"execution_count": 1,
"metadata": {},
"outputs": [],
"source": [
"# Imports\n",
"import os\n",
"import shutil\n",
"\n",
"import cmcrameri.cm as cmc\n",
"import cv2\n",
"import matplotlib.cm as cm\n",
"import matplotlib.colors as mcolors\n",
"import matplotlib.pyplot as plt\n",
"import numpy as np\n",
"import pandas as pd\n",
"import pyvista as pv\n",
"from IPython.display import Video\n",
"from tqdm import tqdm\n",
"\n",
"import gdtchron as gdt"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"The cell below creates a list of 200 sorted files from the `.pvtu` files output by the ASPECT model"
]
},
{
"cell_type": "code",
"execution_count": 2,
"metadata": {},
"outputs": [],
"source": [
"# Get files from model output\n",
"part_path = 'output_exhumation_box/particles'\n",
"files = [\n",
" os.path.join(part_path, file) for file in os.listdir(part_path) \n",
" if file.endswith('.pvtu')\n",
" ]\n",
"files.sort()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"This cell feeds those files into GDTchron to predict AHe ages. The list of files, the thermochronologic system, the time between model outputs (Myr), the prefix for the file names, and the number of processes/cores to use are specified in this example.\n",
"\n",
"Note that this and the other calls to `gdt.run_vtk` may take a couple hours each to run, depending on the system and number of processes."
]
},
{
"cell_type": "code",
"execution_count": 3,
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
" "
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"All AHe timesteps complete\n"
]
}
],
"source": [
"# Run AHe\n",
"gdt.run_vtk(files=files, system='AHe', time_interval=0.1, \n",
" file_prefix='meshes_tchron_exhumation', processes=os.cpu_count() / 2)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"The cell below sets the input files to now be the output `.vtu` files containing the AHe ages, so that all three systems will be saved to the same set of files that are separate from the original model output"
]
},
{
"cell_type": "code",
"execution_count": 4,
"metadata": {},
"outputs": [],
"source": [
"# Set path to reuse files from AHe\n",
"new_path = 'meshes_tchron_exhumation'\n",
"files = [\n",
" os.path.join(new_path, file) for file in os.listdir(new_path)\n",
" if file.endswith('.vtu')\n",
" ]\n",
"files.sort()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"The next two cells below simply take the files with AHe ages and add AFT and ZHe. Note again that these may each take a couple hours to run."
]
},
{
"cell_type": "code",
"execution_count": 5,
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
" "
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"All AFT timesteps complete\n"
]
}
],
"source": [
"# Add AFT\n",
"gdt.run_vtk(files=files, system='AFT', time_interval=0.1, \n",
" file_prefix='meshes_tchron_exhumation', processes=os.cpu_count() / 2)"
]
},
{
"cell_type": "code",
"execution_count": 6,
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
" "
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"All ZHe timesteps complete\n"
]
}
],
"source": [
"# Add ZHe\n",
"gdt.run_vtk(files=files, system='ZHe', time_interval=0.1, \n",
" file_prefix='meshes_tchron_exhumation', processes=os.cpu_count() / 2)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"The next set of cells uses the final set of `.vtu` files with the three systems to make a video showing the evolution of thermochronologic ages over the model runtime. The cell below sets some parameters for the image files that will be used in the video."
]
},
{
"cell_type": "code",
"execution_count": 5,
"metadata": {},
"outputs": [],
"source": [
"# Set up video\n",
"directory = 'meshes_tchron_exhumation'\n",
"\n",
"time_total = 20 # Myr\n",
"model_step = 0.1 # Myr\n",
"bounds = [35, 65, 0, 10] # km, X position and depth\n",
"\n",
"nsteps = int(time_total / model_step)\n",
"\n",
"timesteps = np.arange(0, nsteps + 1, 1)\n",
"\n",
"files = [\n",
" os.path.join(directory, file) for file in os.listdir(directory) \n",
" if file.endswith('.vtu')\n",
" ]\n",
"files.sort()\n",
"\n",
"# Load solution files for tempterature contours\n",
"dir_temp = 'output_exhumation_box/solution'\n",
"files_t = [\n",
" os.path.join(dir_temp, file) for file in os.listdir(dir_temp)\n",
" if file.endswith('.pvtu')\n",
" ]\n",
"files_t.sort()\n",
"\n",
"image_dir = 'images/'\n",
"\n",
"ahe_cmap = cmc.lapaz_r\n",
"aft_cmap = cmc.lajolla_r\n",
"zhe_cmap = cmc.bamako_r\n",
"cat_cmap = cmc.batlowS\n",
"clim = [0, time_total]\n",
"bar = False"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"The cell below generates the images that will be stitched together to make the video"
]
},
{
"cell_type": "code",
"execution_count": 6,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Cleared existing image directory...\n"
]
},
{
"name": "stderr",
"output_type": "stream",
"text": [
"100%|██████████| 201/201 [03:03<00:00, 1.10it/s]\n"
]
}
],
"source": [
"# Run image creation\n",
"try:\n",
" shutil.rmtree(image_dir)\n",
"except Exception:\n",
" print(\"Creating new image directory...\")\n",
"else:\n",
" print(\"Cleared existing image directory...\")\n",
"\n",
"os.makedirs(image_dir, exist_ok=False)\n",
"\n",
"# Make images for the video\n",
"for k, step in enumerate(tqdm(timesteps)):\n",
" time_ma = time_total - (step * model_step)\n",
" time_str = str(round(step / 2, 1)).zfill(5).replace('.', '-')\n",
" \n",
" fig, axs = plt.subplots(2, 2, dpi=300, figsize=(6.5, 6))\n",
" axs = axs.flat\n",
"\n",
" title = 'Uplift' if 5 <= time_ma <= 10 else 'Quiescence'\n",
"\n",
" axs[3].set_title(str(round(time_ma, 1)) + ' Ma - ' + title, loc='left')\n",
"\n",
" mesh = pv.read(files[k])\n",
" mesh.points = mesh.points / 1e3 # Convert to km\n",
" mesh.points[:, 1] = -(mesh.points[:, 1] - 20) # Convert Y position to depth\n",
" \n",
" gdt.plot_vtk_2d(mesh, 'AHe', bounds=bounds, ax=axs[0],\n",
" cmap=ahe_cmap, colorbar=bar, clim=clim)\n",
" \n",
" gdt.plot_vtk_2d(mesh, 'AFT', bounds=bounds, ax=axs[1],\n",
" cmap=aft_cmap, colorbar=bar, clim=clim)\n",
" \n",
" gdt.plot_vtk_2d(mesh, 'ZHe', bounds=bounds, ax=axs[2],\n",
" cmap=zhe_cmap, colorbar=bar, clim=clim)\n",
"\n",
" # Load solution file for temperature contours\n",
" mesh_t = pv.read(files_t[k])\n",
" mesh_t.points /= 1e3 # Convert to km\n",
" mesh_t['T'] -= 273.15 # Convert temp to °C\n",
" mesh_t.points[:, 1] = -(mesh_t.points[:, 1] - 20) # Convert Y position to depth\n",
" mesh_cntrs = mesh_t.contour(isosurfaces=np.arange(100, 301, 200), scalars='T')\n",
"\n",
" for ax in axs[0:3]:\n",
" gdt.plot_vtk_2d(mesh_cntrs, None, bounds=bounds, ax=ax,\n",
" color='dark grey', line_width=8, colorbar=True)\n",
"\n",
" axs[0].set_title('AHe')\n",
" axs[1].set_title('AFT')\n",
" axs[2].set_title('ZHe')\n",
"\n",
" y = np.round(mesh.points[:, 1], 0)\n",
" AHe = mesh['AHe']\n",
" AFT = mesh['AFT']\n",
" ZHe = mesh['ZHe']\n",
"\n",
" df = pd.DataFrame({'y': y, 'AHe': AHe, 'AFT': AFT, 'ZHe': ZHe})\n",
" df_max = df.groupby('y').agg({'y': 'first', 'AHe': 'max', 'AFT': 'max',\n",
" 'ZHe': 'max'})\n",
"\n",
" axs[3].plot(df_max['AHe'], df_max['y'], c=cat_cmap.colors[6], label='AHe')\n",
" axs[3].plot(df_max['AFT'], df_max['y'], c=cat_cmap.colors[4], label='AFT')\n",
" axs[3].plot(df_max['ZHe'], df_max['y'], c=cat_cmap.colors[2], label='ZHe')\n",
"\n",
" axs[3].set_xlim(-0.25, time_total)\n",
" axs[3].set_ylim(bounds[2], bounds[3])\n",
" axs[3].set_xlabel('Maximum Age (Ma)', fontsize=6)\n",
" axs[3].tick_params(labelsize=6)\n",
" axs[3].text(0.5, 0.2,\n",
" 'Surface Ages:\\n'\n",
" 'AHe: ' + str(round(df_max['AHe'][0], 1)) + ' Ma\\n'\n",
" 'AFT: ' + str(round(df_max['AFT'][0], 1)) + ' Ma\\n'\n",
" 'ZHe: ' + str(round(df_max['ZHe'][0], 1)) + ' Ma\\n',\n",
" transform=axs[3].transAxes)\n",
" axs[3].legend(fontsize=6, loc='lower right')\n",
"\n",
" for ax in [axs[0], axs[1], axs[2]]:\n",
" ax.set_xlabel('X Position (km)', fontsize=6)\n",
" \n",
" for ax in axs:\n",
" ax.set_ylabel('Depth (km)', fontsize=6, labelpad=2)\n",
" ax.tick_params(labelsize=6)\n",
" ax.invert_yaxis()\n",
"\n",
" cax_ahe = fig.add_axes([0.24, 0.68, 0.1, 0.015])\n",
" cax_aft = fig.add_axes([0.7, 0.68, 0.1, 0.015])\n",
" cax_zhe = fig.add_axes([0.24, 0.26, 0.1, 0.015])\n",
" \n",
" norm = mcolors.Normalize(vmin=clim[0], vmax=clim[1])\n",
" \n",
" sm_ahe = cm.ScalarMappable(cmap=ahe_cmap, norm=norm)\n",
" sm_aft = cm.ScalarMappable(cmap=aft_cmap, norm=norm)\n",
" sm_zhe = cm.ScalarMappable(cmap=zhe_cmap, norm=norm)\n",
" sms = [sm_ahe, sm_aft, sm_zhe]\n",
" \n",
" for k, cax in enumerate([cax_ahe, cax_aft, cax_zhe]):\n",
" cax.tick_params(labelsize=6, pad=2)\n",
" plt.colorbar(sms[k], cax=cax, orientation='horizontal')\n",
" cax.set_title('Age (Ma)', fontsize=6, pad=2)\n",
" \n",
" fig.savefig(image_dir + time_str + '.jpg')\n",
" plt.close()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"The cell below stitches the images together and makes a video file: `video_tchron.mp4`."
]
},
{
"cell_type": "code",
"execution_count": 7,
"metadata": {},
"outputs": [],
"source": [
"# Make movie\n",
"img_paths = [\n",
" image_dir + file for file in sorted(os.listdir(image_dir)) if file.endswith('.jpg')\n",
" ]\n",
"\n",
"frame = cv2.imread(img_paths[0])\n",
"height, width, layers = frame.shape\n",
"\n",
"fourcc = cv2.VideoWriter_fourcc(*'avc1')\n",
"\n",
"frate = 1 / model_step\n",
"\n",
"video = cv2.VideoWriter('video_tchron.mp4', fourcc, frate, (width, height))\n",
"\n",
"for img in img_paths:\n",
" video.write(cv2.imread(img))\n",
"\n",
"cv2.destroyAllWindows()\n",
"video.release()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"This final cell displays the video in this Jupyter Notebook."
]
},
{
"cell_type": "code",
"execution_count": 8,
"metadata": {},
"outputs": [
{
"data": {
"text/html": [
""
],
"text/plain": [
""
]
},
"execution_count": 8,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"video_path = 'video_tchron.mp4'\n",
"Video(video_path, embed=True, width=500)"
]
}
],
"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.12.12"
}
},
"nbformat": 4,
"nbformat_minor": 4
}