Process ASPECT Exhumation Model Outputs

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.

# Imports
import os
import shutil

import cmcrameri.cm as cmc
import cv2
import matplotlib.cm as cm
import matplotlib.colors as mcolors
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import pyvista as pv
from IPython.display import Video
from tqdm import tqdm

import gdtchron as gdt

The cell below creates a list of 200 sorted files from the .pvtu files output by the ASPECT model

# Get files from model output
part_path = 'output_exhumation_box/particles'
files = [
    os.path.join(part_path, file) for file in os.listdir(part_path)
    if file.endswith('.pvtu')
    ]
files.sort()

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.

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.

# Run AHe
gdt.run_vtk(files=files, system='AHe', time_interval=0.1,
            file_prefix='meshes_tchron_exhumation', processes=os.cpu_count() / 2)

All AHe timesteps complete

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

# Set path to reuse files from AHe
new_path = 'meshes_tchron_exhumation'
files = [
    os.path.join(new_path, file) for file in os.listdir(new_path)
    if file.endswith('.vtu')
    ]
files.sort()

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.

# Add AFT
gdt.run_vtk(files=files, system='AFT', time_interval=0.1,
            file_prefix='meshes_tchron_exhumation', processes=os.cpu_count() / 2)

All AFT timesteps complete

# Add ZHe
gdt.run_vtk(files=files, system='ZHe', time_interval=0.1,
            file_prefix='meshes_tchron_exhumation', processes=os.cpu_count() / 2)

All ZHe timesteps complete

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.

# Set up video
directory = 'meshes_tchron_exhumation'

time_total = 20  # Myr
model_step = 0.1  # Myr
bounds = [35, 65, 0, 10]  # km, X position and depth

nsteps = int(time_total / model_step)

timesteps = np.arange(0, nsteps + 1, 1)

files = [
    os.path.join(directory, file) for file in os.listdir(directory)
    if file.endswith('.vtu')
    ]
files.sort()

# Load solution files for tempterature contours
dir_temp = 'output_exhumation_box/solution'
files_t = [
    os.path.join(dir_temp, file) for file in os.listdir(dir_temp)
    if file.endswith('.pvtu')
    ]
files_t.sort()

image_dir = 'images/'

ahe_cmap = cmc.lapaz_r
aft_cmap = cmc.lajolla_r
zhe_cmap = cmc.bamako_r
cat_cmap = cmc.batlowS
clim = [0, time_total]
bar = False

The cell below generates the images that will be stitched together to make the video

# Run image creation
try:
    shutil.rmtree(image_dir)
except Exception:
    print("Creating new image directory...")
else:
    print("Cleared existing image directory...")

os.makedirs(image_dir, exist_ok=False)

# Make images for the video
for k, step in enumerate(tqdm(timesteps)):
    time_ma = time_total - (step * model_step)
    time_str = str(round(step / 2, 1)).zfill(5).replace('.', '-')

    fig, axs = plt.subplots(2, 2, dpi=300, figsize=(6.5, 6))
    axs = axs.flat

    title = 'Uplift' if 5 <= time_ma <= 10 else 'Quiescence'

    axs[3].set_title(str(round(time_ma, 1)) + ' Ma - ' + title, loc='left')

    mesh = pv.read(files[k])
    mesh.points = mesh.points / 1e3  # Convert to km
    mesh.points[:, 1] = -(mesh.points[:, 1] - 20)  # Convert Y position to depth

    gdt.plot_vtk_2d(mesh, 'AHe', bounds=bounds, ax=axs[0],
              cmap=ahe_cmap, colorbar=bar, clim=clim)

    gdt.plot_vtk_2d(mesh, 'AFT', bounds=bounds, ax=axs[1],
            cmap=aft_cmap, colorbar=bar, clim=clim)

    gdt.plot_vtk_2d(mesh, 'ZHe', bounds=bounds, ax=axs[2],
        cmap=zhe_cmap, colorbar=bar, clim=clim)

    # Load solution file for temperature contours
    mesh_t = pv.read(files_t[k])
    mesh_t.points /= 1e3  # Convert to km
    mesh_t['T'] -= 273.15  # Convert temp to °C
    mesh_t.points[:, 1] = -(mesh_t.points[:, 1] - 20)  # Convert Y position to depth
    mesh_cntrs = mesh_t.contour(isosurfaces=np.arange(100, 301, 200), scalars='T')

    for ax in axs[0:3]:
        gdt.plot_vtk_2d(mesh_cntrs, None, bounds=bounds, ax=ax,
          color='dark grey', line_width=8, colorbar=True)

    axs[0].set_title('AHe')
    axs[1].set_title('AFT')
    axs[2].set_title('ZHe')

    y = np.round(mesh.points[:, 1], 0)
    AHe = mesh['AHe']
    AFT = mesh['AFT']
    ZHe = mesh['ZHe']

    df = pd.DataFrame({'y': y, 'AHe': AHe, 'AFT': AFT, 'ZHe': ZHe})
    df_max = df.groupby('y').agg({'y': 'first', 'AHe': 'max', 'AFT': 'max',
                                   'ZHe': 'max'})

    axs[3].plot(df_max['AHe'], df_max['y'], c=cat_cmap.colors[6], label='AHe')
    axs[3].plot(df_max['AFT'], df_max['y'], c=cat_cmap.colors[4], label='AFT')
    axs[3].plot(df_max['ZHe'], df_max['y'], c=cat_cmap.colors[2], label='ZHe')

    axs[3].set_xlim(-0.25, time_total)
    axs[3].set_ylim(bounds[2], bounds[3])
    axs[3].set_xlabel('Maximum Age (Ma)', fontsize=6)
    axs[3].tick_params(labelsize=6)
    axs[3].text(0.5, 0.2,
                'Surface Ages:\n'
                  'AHe: ' + str(round(df_max['AHe'][0], 1)) + ' Ma\n'
                  'AFT: ' + str(round(df_max['AFT'][0], 1)) + ' Ma\n'
                  'ZHe: ' + str(round(df_max['ZHe'][0], 1)) + ' Ma\n',
                transform=axs[3].transAxes)
    axs[3].legend(fontsize=6, loc='lower right')

    for ax in [axs[0], axs[1], axs[2]]:
        ax.set_xlabel('X Position (km)', fontsize=6)

    for ax in axs:
        ax.set_ylabel('Depth (km)', fontsize=6, labelpad=2)
        ax.tick_params(labelsize=6)
        ax.invert_yaxis()

    cax_ahe = fig.add_axes([0.24, 0.68, 0.1, 0.015])
    cax_aft = fig.add_axes([0.7, 0.68, 0.1, 0.015])
    cax_zhe = fig.add_axes([0.24, 0.26, 0.1, 0.015])

    norm = mcolors.Normalize(vmin=clim[0], vmax=clim[1])

    sm_ahe = cm.ScalarMappable(cmap=ahe_cmap, norm=norm)
    sm_aft = cm.ScalarMappable(cmap=aft_cmap, norm=norm)
    sm_zhe = cm.ScalarMappable(cmap=zhe_cmap, norm=norm)
    sms = [sm_ahe, sm_aft, sm_zhe]

    for k, cax in enumerate([cax_ahe, cax_aft, cax_zhe]):
        cax.tick_params(labelsize=6, pad=2)
        plt.colorbar(sms[k], cax=cax, orientation='horizontal')
        cax.set_title('Age (Ma)', fontsize=6, pad=2)

    fig.savefig(image_dir + time_str + '.jpg')
    plt.close()

Cleared existing image directory...

100%|██████████| 201/201 [03:03<00:00,  1.10it/s]

The cell below stitches the images together and makes a video file: video_tchron.mp4.

# Make movie
img_paths = [
    image_dir + file for file in sorted(os.listdir(image_dir)) if file.endswith('.jpg')
    ]

frame = cv2.imread(img_paths[0])
height, width, layers = frame.shape

fourcc = cv2.VideoWriter_fourcc(*'avc1')

frate = 1 / model_step

video = cv2.VideoWriter('video_tchron.mp4', fourcc, frate, (width, height))

for img in img_paths:
    video.write(cv2.imread(img))

cv2.destroyAllWindows()
video.release()

This final cell displays the video in this Jupyter Notebook.

video_path = 'video_tchron.mp4'
Video(video_path, embed=True, width=500)