postProcess/2-LidDrivenCavity-Newtonian-dyeInjection.py
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
from mpl_toolkits.axes_grid1 import make_axes_locatable
import argparse # Add at top with other imports
import multiprocessing as mp
from functools import partial
matplotlib.rcParams['font.family'] = 'serif'
matplotlib.rcParams['text.usetex'] = True
matplotlib.rcParams['text.latex.preamble'] = r'\usepackage{amsmath}'
def gettingfield(filename, zmin, zmax, rmin, rmax, nr):
Extract simulation field data from a simulation file.
This function executes the external program “getData-LidDriven” with the given simulation file and boundary parameters to extract numerical data for axial and radial coordinates, temperature, velocity, and stream function. The output is parsed into NumPy arrays, reshaped to a (nz, nr) grid based on the number of radial grid points (nr), rotated 90° counterclockwise, and flipped vertically to ensure the correct orientation for further processing. It returns the processed field arrays along with the total number of axial grid points (nz).
Args: filename: Path to the simulation data file. zmin: Minimum axial coordinate. zmax: Maximum axial coordinate. rmin: Minimum radial coordinate. rmax: Maximum radial coordinate. nr: Number of grid points in the radial direction.
Returns: tuple: A tuple containing: - R (numpy.ndarray): 2D array of radial coordinate values. - Z (numpy.ndarray): 2D array of axial coordinate values. - T (numpy.ndarray): 2D array of temperature values. - vel (numpy.ndarray): 2D array of velocity values. - psi (numpy.ndarray): 2D array of stream function values. - nz (int): Number of grid points in the axial direction.
exe = ["./getData-LidDriven", filename, str(zmin), str(rmin), 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, Ttemp, veltemp, Psitemp = [],[],[],[], []
for n1 in range(len(temp2)):
temp3 = temp2[n1].split(" ")
if temp3 == ['']:
pass
else:
Ztemp.append(float(temp3[0]))
Rtemp.append(float(temp3[1]))
Ttemp.append(float(temp3[2]))
veltemp.append(float(temp3[3]))
Psitemp.append(float(temp3[4]))
R = np.asarray(Rtemp)
Z = np.asarray(Ztemp)
T = np.asarray(Ttemp)
vel = np.asarray(veltemp)
psi = np.asarray(Psitemp)
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))
T.resize((nz, nr))
vel.resize((nz, nr))
psi.resize((nz, nr))
# rotate the arrays by 90 degrees
R = np.rot90(R, k=1)
Z = np.rot90(Z, k=1)
T = np.rot90(T, k=1)
vel = np.rot90(vel, k=1)
psi = np.rot90(psi, k=1)
# flip the array
R = np.flip(R, axis=0)
Z = np.flip(Z, axis=0)
T = np.flip(T, axis=0)
vel = np.flip(vel, axis=0)
psi = np.flip(psi, axis=0)
return R, Z, T, vel, psi, nz
# ----------------------------------------------------------------------------------------------------------------------
def process_timestep(ti, caseToProcess, folder, tsnap, GridsPerR, rmin, rmax, zmin, zmax, lw):
Generates and saves a two-panel plot for a simulation timestep.
Calculates the simulation time using the timestep index and a time increment, then retrieves field data from an intermediate snapshot file. Creates a figure with two subplots—one displaying the dye color (with a coolwarm heat map) and the other showing velocity magnitude (with a viridis heat map). Both panels include domain boundaries and streamlines for the stream function. If the snapshot file is missing or the output image already exists, the function prints a warning and exits.
t = tsnap * ti
place = f"{caseToProcess}/intermediate/snapshot-{t:.4f}"
name = f"{folder}/{int(t*1000):08d}.png"
# Check if the file exists
if not os.path.exists(place):
print(f"{place} File not found!")
return
# if name exits, return
if os.path.exists(name):
print(f"{name} already exists!")
return
# if folder does not exist, create it
if not os.path.exists(folder):
os.makedirs(folder)
# Calculate number of grid points in r-direction based on domain size
nr = int(GridsPerR * rmax)
# Extract field data from the simulation file
R, Z, T, vel, psi, nz = gettingfield(place, zmin, zmax, rmin, rmax, nr)
# Get actual domain bounds from the data
zminp, zmaxp, rminp, rmaxp = Z.min(), Z.max(), R.min(), R.max()
# Set up the figure with two subplots
AxesLabel, TickLabel = 50, 20
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(24, 11), constrained_layout=True)
# First subplot - Rate of Strain Tensor (D2)
# Draw domain boundaries
ax1.plot([rmin, rmin], [zmin, zmax], '-', color='black', linewidth=lw)
ax1.plot([rmin, rmax], [zmin, zmin], '-', color='black', linewidth=lw)
ax1.plot([rmin, rmax], [zmax, zmax], '-', color='black', linewidth=lw)
ax1.plot([rmax, rmax], [zmin, zmax], '-', color='black', linewidth=lw)
# Plot the dye with a heat map
cntrl1 = ax1.pcolormesh(Z, R, T, cmap="coolwarm", edgecolor='face', vmax=1, vmin=0)
# Add streamlines using the stream function
ax1.contour(Z, R, psi, 20, colors='black', linewidths=2)
# Configure the first subplot
ax1.set_aspect('equal')
ax1.set_xlim(rmin, rmax)
ax1.set_ylim(zmin, zmax)
ax1.set_title(r'Rate of Strain Tensor', fontsize=TickLabel)
# Add colorbar to the first subplot
divider1 = make_axes_locatable(ax1)
cax1 = divider1.append_axes("right", size="5%", pad=0.1)
c1 = plt.colorbar(cntrl1, cax=cax1)
c1.set_label(r'$log_{10}(\|\mathcal{D}_{ij}\|)$', fontsize=TickLabel, labelpad=5)
c1.ax.tick_params(labelsize=TickLabel)
# Second subplot - Velocity Magnitude
# Draw domain boundaries
ax2.plot([rmin, rmin], [zmin, zmax], '-', color='black', linewidth=lw)
ax2.plot([rmin, rmax], [zmin, zmin], '-', color='black', linewidth=lw)
ax2.plot([rmin, rmax], [zmax, zmax], '-', color='black', linewidth=lw)
ax2.plot([rmax, rmax], [zmin, zmax], '-', color='black', linewidth=lw)
# Plot velocity magnitude with viridis colormap
cntrl2 = ax2.pcolormesh(Z, R, vel, cmap="viridis", edgecolor='face', vmax = 1, vmin = 0)
# Add streamlines using the stream function
ax2.contour(Z, R, psi, 20, colors='black', linewidths=2)
# Configure the second subplot
ax2.set_aspect('equal')
ax2.set_xlim(rmin, rmax)
ax2.set_ylim(zmin, zmax)
ax2.set_title(r'Velocity Magnitude', fontsize=TickLabel)
# Add colorbar to the second subplot
divider2 = make_axes_locatable(ax2)
cax2 = divider2.append_axes("right", size="5%", pad=0.1)
c2 = plt.colorbar(cntrl2, cax=cax2)
c2.set_label(r'Velocity', fontsize=TickLabel, labelpad=5)
c2.ax.tick_params(labelsize=TickLabel)
# Turn off axes for cleaner visualization
ax1.axis('off')
ax2.axis('off')
plt.savefig(name, bbox_inches='tight')
# Display the figure
plt.close()
def main():
# Get number of CPUs from command line argument, or use all available
Parses command-line arguments and initiates parallel processing of simulation timesteps.
This function reads simulation configuration and processing parameters from the command line, including grid settings, simulation bounds, and performance options. It ensures that the output directory exists and then sets up a multiprocessing pool to concurrently process simulation data timesteps by mapping a worker function that generates visualizations.
parser = argparse.ArgumentParser()
parser.add_argument('--CPUs', type=int, default=mp.cpu_count(), help='Number of CPUs to use')
parser.add_argument('--nGFS', type=int, default=150, help='Number of restart files to process')
parser.add_argument('--GridsPerR', type=int, default=512, help='Number of grids per R')
parser.add_argument('--ZMAX', type=float, default=0.5, help='Maximum Z value')
parser.add_argument('--RMAX', type=float, default=0.5, help='Maximum R value')
parser.add_argument('--ZMIN', type=float, default=-0.5, help='Minimum Z value')
parser.add_argument('--RMIN', type=float, default=-0.5, help='Minimum R value')
parser.add_argument('--tsnap', type=float, default=0.01, help='Time snap')
parser.add_argument('--caseToProcess', type=str, default='../simulationCases/2-LidDrivenCavity-Newtonian-dyeInjection', help='Case to process')
parser.add_argument('--folderToSave', type=str, default='2-LidDrivenCavity-Newtonian-dyeInjection', help='Folder to save')
args = parser.parse_args()
num_processes = args.CPUs
nGFS = args.nGFS
tsnap = args.tsnap
ZMAX = args.ZMAX
RMAX = args.RMAX
ZMIN = args.ZMIN
RMIN = args.RMIN
rmin, rmax, zmin, zmax = [RMIN, RMAX, ZMIN, ZMAX]
GridsPerR = args.GridsPerR
lw = 2
folder = args.folderToSave
caseToProcess = args.caseToProcess
if not os.path.isdir(folder):
os.makedirs(folder)
# Create a pool of worker processes
with mp.Pool(processes=num_processes) as pool:
# Create partial function with fixed arguments
process_func = partial(process_timestep, caseToProcess=caseToProcess,
folder=folder, tsnap=tsnap,
GridsPerR=GridsPerR, rmin=rmin, rmax=rmax,
zmin=zmin, zmax=zmax, lw=lw)
# Map the process_func to all timesteps
pool.map(process_func, range(nGFS))
if __name__ == "__main__":
main()