postProcess/VideoAxi.py
#!/usr/bin/env python3
# -*- coding: utf-8 -*-Viscoelastic Visualization Tool
This script processes and visualizes fluid dynamics simulation data, particularly focused on droplet impact and deformable soft matter like liquid drops, sheets, and bubbles. It extracts interface positions and scalar field data from simulation files and creates visualizations showing physical quantities like strain rates and stresses.
The script is designed to process multiple simulation snapshots in parallel, extracting data using external executables and generating visualizations with proper colormaps, scales, and mathematical labels.
Features: - Extracts fluid interfaces and scalar fields from simulation files - Generates visualizations with proper colormaps and mathematical labels - Processes multiple timesteps in parallel using multiprocessing - Configurable via command-line arguments for different simulation cases - Creates publication-quality figures with LaTeX-rendered mathematical expressions
Usage: python fluid_vis.py [options]
Command-line Arguments: –CPUs Number of CPUs to use for parallel processing (default: all available) –nGFS Number of restart files to process (default: 550) –ZMAX Maximum Z coordinate for visualization (default: 4.0) –RMAX Maximum R coordinate for visualization (default: 2.0) –ZMIN Minimum Z coordinate for visualization (default: -4.0) –caseToProcess Path to simulation case directory (default: ‘../simulationCases/dropImpact’) –folderToSave Directory to save visualization images (default: ‘dropImpact’)
Dependencies: External executables: getFacet2D, getData-elastic-scalar2D Python libraries: numpy, matplotlib, subprocess, multiprocessing
Author: Vatsal Sanjay Email: [email protected] Affiliation: Physics of Fluids Last updated: Jul 24, 2024
import numpy as np
import os
import subprocess as sp
import matplotlib
import matplotlib.pyplot as plt
from matplotlib.collections import LineCollection
from matplotlib.ticker import StrMethodFormatter
import multiprocessing as mp
from functools import partial
import argparse
import matplotlib.colors as mcolors
custom_colors = ["white", "#DA8A67", "#A0522D", "#400000"]
custom_cmap = mcolors.LinearSegmentedColormap.from_list("custom_hot", custom_colors)
# Configure matplotlib for publication-quality figures with LaTeX rendering
matplotlib.rcParams['font.family'] = 'serif'
matplotlib.rcParams['text.usetex'] = True
matplotlib.rcParams['text.latex.preamble'] = r'\usepackage{amsmath}'
def gettingFacets(filename, includeCoat='true'):Extract interface positions (facets) from simulation files.
Args: filename (str): Path to simulation snapshot file includeCoat (str, optional): Whether to include coating layer. Defaults to ‘true’.
Returns: list: List of line segments defining fluid interfaces
exe = ["./getFacet2D", filename, includeCoat]
p = sp.Popen(exe, stdout=sp.PIPE, stderr=sp.PIPE)
stdout, stderr = p.communicate()
temp1 = stderr.decode("utf-8")
temp2 = temp1.split("\n")
segs = []
skip = False
if (len(temp2) > 1e2):
for n1 in range(len(temp2)):
temp3 = temp2[n1].split(" ")
if temp3 == ['']:
skip = False
pass
else:
if not skip:
temp4 = temp2[n1+1].split(" ")
r1, z1 = np.array([float(temp3[1]), float(temp3[0])])
r2, z2 = np.array([float(temp4[1]), float(temp4[0])])
segs.append(((r1, z1),(r2, z2)))
segs.append(((-r1, z1),(-r2, z2)))
skip = True
return segs
def gettingfield(filename, zmin, zmax, rmax, nr):Extract scalar field data from simulation files.
Args: filename (str): Path to simulation snapshot file zmin (float): Minimum Z coordinate zmax (float): Maximum Z coordinate rmax (float): Maximum R coordinate nr (int): Number of grid points in R direction
Returns: tuple: (R, Z, D2, vel, taup, nz) arrays of coordinates and field values
exe = ["./getData-elastic-scalar2D", filename, str(zmin), str(0), str(zmax), str(rmax), str(nr)]
p = sp.Popen(exe, stdout=sp.PIPE, stderr=sp.PIPE)
stdout, stderr = p.communicate()
temp1 = stderr.decode("utf-8")
temp2 = temp1.split("\n")
# print(temp2) #debugging
Rtemp, Ztemp, D2temp, veltemp, taupTemp = [],[],[],[],[]
for n1 in range(len(temp2)):
temp3 = temp2[n1].split(" ")
if temp3 == ['']:
pass
else:
Ztemp.append(float(temp3[0]))
Rtemp.append(float(temp3[1]))
D2temp.append(float(temp3[2]))
veltemp.append(float(temp3[3]))
taupTemp.append(float(temp3[4]))
R = np.asarray(Rtemp)
Z = np.asarray(Ztemp)
D2 = np.asarray(D2temp)
vel = np.asarray(veltemp)
taup = np.asarray(taupTemp)
nz = int(len(Z)/nr)
# print("nr is %d %d" % (nr, len(R))) # debugging
print("nz is %d" % nz)
R.resize((nz, nr))
Z.resize((nz, nr))
D2.resize((nz, nr))
vel.resize((nz, nr))
taup.resize((nz, nr))
return R, Z, D2, vel, taup, nz
# ----------------------------------------------------------------------------------------------------------------------
def process_timestep(ti, folder, nGFS, GridsPerR, rmin, rmax, zmin, zmax, lw, caseToProcess):Process a single timestep from simulation data and generate visualization.
Args: ti (int): Timestep index folder (str): Directory to save output images nGFS (int): Total number of timesteps GridsPerR (int): Grid points per unit length in R direction rmin (float): Minimum R coordinate rmax (float): Maximum R coordinate zmin (float): Minimum Z coordinate zmax (float): Maximum Z coordinate lw (float): Line width for plot elements caseToProcess (str): Path to simulation case directory
t = 0.01 * ti
place = f"{caseToProcess}/intermediate/snapshot-{t:.4f}"
name = f"{folder}/{int(t*1000):08d}.png"
if not os.path.exists(place):
print(f"{place} File not found!")
return
if os.path.exists(name):
print(f"{name} Image present!")
return
segs1 = gettingFacets(place)
segs2 = gettingFacets(place, 'false')
if not segs1 and not segs2:
print(f"Problem in the available file {place}")
return
nr = int(GridsPerR * rmax)
R, Z, taus, vel, taup, nz = gettingfield(place, zmin, zmax, rmax, nr)
zminp, zmaxp, rminp, rmaxp = Z.min(), Z.max(), R.min(), R.max()
# Plotting
AxesLabel, TickLabel = 50, 20
fig, ax = plt.subplots()
fig.set_size_inches(19.20, 10.80)
# Draw domain boundaries
ax.plot([0, 0], [zmin, zmax], '-.', color='grey', linewidth=lw)
ax.plot([rmin, rmin], [zmin, zmax], '-', color='black', linewidth=lw)
ax.plot([rmin, rmax], [zmin, zmin], '-', color='black', linewidth=lw)
ax.plot([rmin, rmax], [zmax, zmax], '-', color='black', linewidth=lw)
ax.plot([rmax, rmax], [zmin, zmax], '-', color='black', linewidth=lw)
# Add fluid interfaces
line_segments = LineCollection(segs2, linewidths=4, colors='green', linestyle='solid')
ax.add_collection(line_segments)
line_segments = LineCollection(segs1, linewidths=4, colors='blue', linestyle='solid')
ax.add_collection(line_segments)
# Plot scalar fields with colormaps
cntrl1 = ax.imshow(taus, cmap="hot_r", interpolation='Bilinear', origin='lower',
extent=[-rminp, -rmaxp, zminp, zmaxp], vmax=2.0, vmin=-3.0)
# TODO: fixme the colorbar bounds for taup must be set manually based on the simulated case.
cntrl2 = ax.imshow(taup, interpolation='Bilinear', cmap=custom_cmap, origin='lower',
extent=[rminp, rmaxp, zminp, zmaxp], vmax=2.0, vmin=-3.0)
# Set plot properties
ax.set_aspect('equal')
ax.set_xlim(rmin, rmax)
ax.set_ylim(zmin, zmax)
ax.set_title(f'$t/\\tau_\\gamma$ = {t:4.3f}', fontsize=TickLabel)
# Add colorbars
l, b, w, h = ax.get_position().bounds
# Left colorbar
cb1 = fig.add_axes([l-0.04, b, 0.03, h])
c1 = plt.colorbar(cntrl1, cax=cb1, orientation='vertical')
c1.set_label(r'$\log_{10}\left(\|\mathcal{D}\|\right)$', fontsize=TickLabel, labelpad=5)
c1.ax.tick_params(labelsize=TickLabel)
c1.ax.yaxis.set_ticks_position('left')
c1.ax.yaxis.set_label_position('left')
c1.ax.yaxis.set_major_formatter(StrMethodFormatter('{x:,.1f}'))
# Right colorbar
cb2 = fig.add_axes([l+w+0.01, b, 0.03, h])
c2 = plt.colorbar(cntrl2, cax=cb2, orientation='vertical')
c2.ax.tick_params(labelsize=TickLabel)
c2.set_label(r'$\log_{10}\left(\text{tr}\left(\mathcal{A}\right)-1\right)$', fontsize=TickLabel)
c2.ax.yaxis.set_major_formatter(StrMethodFormatter('{x:,.2f}'))
ax.axis('off')
plt.savefig(name, bbox_inches="tight")
plt.close()
def main():Main function that parses command-line arguments and parallelizes processing of timesteps.
# Set up command-line argument parsing
parser = argparse.ArgumentParser(description='Process fluid dynamics simulation data and create visualizations')
parser.add_argument('--CPUs', type=int, default=mp.cpu_count(),
help='Number of CPUs to use (default: all available)')
parser.add_argument('--nGFS', type=int, default=550,
help='Number of restart files to process (default: 550)')
parser.add_argument('--ZMAX', type=float, default=4.0,
help='Maximum Z value (default: 4.0)')
parser.add_argument('--RMAX', type=float, default=2.0,
help='Maximum R value (default: 2.0)')
parser.add_argument('--ZMIN', type=float, default=-4.0,
help='Minimum Z value (default: -4.0)')
parser.add_argument('--caseToProcess', type=str,
default='../simulationCases/dropImpact',
help='Case to process (default: ../simulationCases/dropImpact)')
parser.add_argument('--folderToSave', type=str, default='dropImpact',
help='Folder to save output images (default: dropImpact)')
args = parser.parse_args()
# Extract arguments
CPUStoUse = args.CPUs
nGFS = args.nGFS
ZMAX = args.ZMAX
RMAX = args.RMAX
ZMIN = args.ZMIN
num_processes = CPUStoUse
rmin, rmax, zmin, zmax = [-RMAX, RMAX, ZMIN, ZMAX]
GridsPerR = 128 # Grid resolution parameter
lw = 2 # Line width for plot elements
folder = args.folderToSave
caseToProcess = args.caseToProcess
# Create output directory if it doesn't exist
if not os.path.isdir(folder):
os.makedirs(folder)
# Create a pool of worker processes for parallel processing
with mp.Pool(processes=num_processes) as pool:
# Create partial function with fixed arguments
process_func = partial(process_timestep,
folder=folder, nGFS=nGFS,
GridsPerR=GridsPerR, rmin=rmin, rmax=rmax,
zmin=zmin, zmax=zmax, lw=lw, caseToProcess=caseToProcess)
# Map the process_func to all timesteps
pool.map(process_func, range(nGFS))
if __name__ == "__main__":
main()