VTK File Reader Code

This code handles the reading and processing of VTK files from PERSEUS (a magnetohydrodynamics (MHD) code). It provides functionality to combine tiled VTK files, convert between structured/unstructured grids, and extract vector field components.

The pyvista library scares me greatly. I will never modify this code again, out of fear that I will inevitably break it. For all intents and purposes, this code is a black box that takes a variable from your vtk files and outputs a neat numpy array. Just call "string_to_array(vtk variable)" and you'll be fine.

import numpy as np
import matplotlib.pyplot as plt
import pyvista as pv
import os

"""
Original Functions:
- combine_tiles()
- unstructured_to_structured() 
- get_mesh_subset()
Author: James Young (github: badwolfend)

Additional Functions & Documentation:
- string_to_array()
- Documentation and comments
Author: Cole Jerum (github: CJerum)
"""

# Independent Variables
run_path = r'./Resistivity_Tests/data'  # path to the directory containing the VTK files
time_analyze = 32  # timestep you want to analyze

def combine_tiles(time, run_dir):
    """Combines VTK tiles into a single grid."""
    time_str = f'{time:04d}'
    vtk_files = [os.path.join(run_dir, vtk_file) 
                 for vtk_file in os.listdir(run_dir) 
                 if vtk_file.endswith(time_str+'.vtr')]
    tiles = [pv.read(vtk_file) for vtk_file in vtk_files]
    combined_grid = tiles[0].copy()
    for tile in tiles[1:]:
        combined_grid = combined_grid.merge(tile)
    return combined_grid

def unstructured_to_structured(grid, variable_name):
    """Converts unstructured grid to structured grid."""
    x_coords = np.unique(grid.points[:, 0])
    y_coords = np.unique(grid.points[:, 1])
    z_coords = np.unique(grid.points[:, 2])
    X, Y, Z = np.meshgrid(x_coords, y_coords, z_coords, indexing='ij')
    
    structured_grid = pv.StructuredGrid(X, Y, Z)
    structured_grid_with_data = structured_grid.sample(grid)
    return structured_grid_with_data

def get_mesh_subset(mesh, nxr, nzr, var):
    """Extracts a subset of the mesh."""
    nxrange = np.arange(nxr[0], nxr[1])
    nzrange = np.arange(nzr[0], nzr[1])
    data = mesh.point_data[var]
    nx, nz, nu = mesh.dimensions
    
    Fy = np.zeros((nzr[1]-nzr[0], nxr[1]-nxr[0]))
    is_vector = len(data.shape) > 1 and data.shape[1] > 1
    
    for zi in nzrange:
        for xi in nxrange:
            if is_vector:
                value = data[(zi * nx) + xi, 0]
            else:
                value = data[(zi * nx) + xi]
            Fy[zi-nzr[0], xi-nxr[0]] = float(value)
    return Fy

def string_to_array(str_var):
    """Converts a variable string to numpy arrays of its components."""
    str_var = str(str_var)
    mesh = combine_tiles(time_analyze, run_path)
    current_mesh = unstructured_to_structured(mesh, str_var)
    nx, nz, nu = current_mesh.dimensions
    
    var = current_mesh.point_data[str_var]
    
    if len(var.shape) > 1 and var.shape[1] == 3:
        ndimensions = 3
        var_x, var_y, var_z = var[:, 0], var[:, 1], var[:, 2]
        
        str_var_x = f"{str_var}_x"
        str_var_y = f"{str_var}_y"
        str_var_z = f"{str_var}_z"
        
        current_mesh.point_data[str_var_x] = var_x
        current_mesh.point_data[str_var_y] = var_y
        current_mesh.point_data[str_var_z] = var_z
        
        var_array_x = get_mesh_subset(current_mesh, [0,nx], [0,nz], str_var_x)
        var_array_y = get_mesh_subset(current_mesh, [0,nx], [0,nz], str_var_y)
        var_array_z = get_mesh_subset(current_mesh, [0,nx], [0,nz], str_var_z)
    else:
        ndimensions = 1
        var_array_x = get_mesh_subset(current_mesh, [0,nx], [0,nz], str_var)
        var_array_y = np.zeros_like(var_array_x)
        var_array_z = np.zeros_like(var_array_x)
    
    return ndimensions, var_array_x, var_array_y, var_array_z

Documentation (HTML) formatted with assistance from Anthropic's Claude, an LLM

VTK File Reader for PERSEUS Simulation Data