visualize.py

import argparse
import os
import sys

import imageio.v3 as iio
import matplotlib as mpl
import numpy as np

from tqdm import tqdm

COMMT = '\033[36m'
ERROR = '\033[31m'
FIGUR = '\033[32m'
RESET = '\033[0m'

def detect_data(frame_data, slice):
	# Mapping of axis to numpy array index
	axis_map = {'x':-2, 'y':-3, 'z':-4}

	# Detect dimension and slice variables
	dim = len(frame_data.shape) - 1
	print(f'{COMMT}[*] Detected data type: {dim}D{RESET}')
	if dim == 3:
		if slice is None or slice[0] not in axis_map:
			print('{ERROR}[*] Slice axis not provided or invalid{RESET}')
			sys.exit(1)
		slice_axis = axis_map[slice[0]]
		slice_idx = int(slice[2:]) if len(slice) > 1 else frame_data.shape[slice_axis] // 2
		print(f'{COMMT}[*] Slicing along {slice[0]}-axis at index {FIGUR}{slice_idx}{RESET}')
		return slice_axis, slice_idx
	else:
		return None, None

def slice_data(frame_data, slice_axis, slice_idx):
	if slice_axis != None:
		frame_data = np.take(frame_data, slice_idx, slice_axis)
		if slice_axis == -2:
			frame_data = np.swapaxes(frame_data, 0, 1)
			frame_data = np.flip(frame_data, axis=0)
		elif slice_axis == -4:
			frame_data = np.flip(frame_data, axis=0)
		# Project 3D data to 2D
		if frame_data.shape[-1] == 3:
			frame_data = np.delete(frame_data, -slice_axis-2, axis=-1)
			if slice_axis == -2:
				frame_data[..., [0, 1]] = frame_data[..., [1, 0]]
			elif slice_axis == -3:
				frame_data[..., 1] = -frame_data[..., 1]
	else:
		frame_data = np.flip(frame_data, axis=0)
	
	return frame_data

def load_file(filename, slice):
	print(f'{COMMT}[*] Loading file: {filename}.npy{RESET}')

	frame_data = np.load(filename + '.npy')
	slice_axis, slice_idx = detect_data(frame_data, slice)
	frame_data = slice_data(frame_data, slice_axis, slice_idx)

	print(f'{COMMT}[*] Loaded data with shape: {FIGUR}{frame_data.shape}{RESET}')
	return frame_data

def load_files(fn_prefix, begin, end, slice):
	print(f'{COMMT}[*] Loading files: {fn_prefix}_{begin:04d}.npy to {fn_prefix}_{end-1:04d}.npy{RESET}')

	frame_data = np.load(f'{fn_prefix}_{begin:04d}.npy')
	slice_axis, slice_idx = detect_data(frame_data, slice)

	# Iterate over the range of frames
	for i in tqdm(range(begin, end), desc='Loading files'):
		if i != begin:
			filename = f'{fn_prefix}_{i:04d}.npy'
			frame_data = np.load(filename)
		
		frame_data = slice_data(frame_data, slice_axis, slice_idx)

		frame_data = np.expand_dims(frame_data, axis=0)
		data = frame_data if i == begin else np.concatenate((data, frame_data), axis=0)

	print(f'{COMMT}[*] Loaded data with shape: {FIGUR}{data.shape}{RESET}')
	return data

def visualize_1d(data, maxval):
	# pixels = mpl.colormaps['viridis']((data / maxval).clip(0, 1).squeeze(axis=-1))
	pixels = mpl.colormaps['coolwarm']((((data / maxval) + 1) / 2).squeeze(axis=-1))
	pixels = np.delete(pixels, 3, axis=-1)
	pixels = (pixels * 255).astype(np.uint8)
	return pixels

def visualize_2d(data, maxval):
	# Use angle as H and norm as S
	angles = np.arctan2(data[..., 1], data[..., 0]) / (2 * np.pi) + 0.5
	norms = (np.linalg.norm(data, axis=-1) / maxval).clip(0, 1)
	# Convert HSV to RGB
	pixels = mpl.colors.hsv_to_rgb(np.stack((angles, norms, np.ones_like(norms)), axis=-1))
	pixels = (pixels * 255).astype(np.uint8)
	return pixels

def visualize_3d(data, maxval):
	vx, vy, vz = data[..., 0], data[..., 1], data[..., 2]
	phis = np.arctan2(vz, vx) / (2 * np.pi) + 0.5
	vx2vz2 = vx**2 + vz**2
	thetas = np.arctan2(np.abs(vy), np.sqrt(vx2vz2)) / np.pi * 2
	norms = (np.sqrt(vx2vz2 + vy**2) / maxval).clip(0, 1)
	# Convert HSV to RGB
	pixels = mpl.colors.hsv_to_rgb(np.stack((phis, 1. - thetas * 16 / 17, norms), axis=-1))
	pixels = (pixels * 255).astype(np.uint8)
	return pixels

def generate(data, maxval, args):
	# Create directory
	data_name, field_name = os.path.split(args.filename)
	data_name = os.path.basename(data_name)
	dirname = os.path.abspath(os.path.join(os.getcwd(), 'results', 'visualization', data_name))
	os.makedirs(dirname, exist_ok=True)
	if args.slice is None:
		filename = os.path.join(dirname, f'{field_name}.png')
	else:
		filename = os.path.join(dirname, f'{args.slice}-{field_name}.png')
	print(f'{COMMT}[*] Saving result to {filename}{RESET}')

	# Generate image
	pixels = globals()[f'visualize_{data.shape[-1]}d'](data, maxval)
	iio.imwrite(filename, pixels)

def generate_sequence(data, maxval, args):
	begin, end = args.range

	# Create directories
	data_name, field_name = os.path.split(args.filename)
	data_name = os.path.basename(data_name)
	if args.slice is None:
		dirname = os.path.abspath(os.path.join(os.getcwd(), 'results', 'visualization', data_name, f'{field_name}'))
	else:
		dirname = os.path.abspath(os.path.join(os.getcwd(), 'results', 'visualization', data_name, f'{args.slice}_{field_name}'))
	os.makedirs(dirname, exist_ok=True)
	print(f'{COMMT}[*] Saving results to {dirname}{RESET}')

	# Generate images and video
	video_fn = os.path.join(dirname, f'v{begin:04d}-{end-1:04d}.mp4')
	with iio.imopen(video_fn, 'w', plugin='pyav') as video:
		video.init_video_stream('libx264', fps=25)
		for i in tqdm(range(begin, end), desc='Generating images'):
			filename = os.path.join(dirname, f'i{i:04d}.png')
			pixels = globals()[f'visualize_{data.shape[-1]}d'](data[i-begin], maxval)
			iio.imwrite(filename, pixels)
			video.write_frame(pixels)

def convert_to_vorticity(data):
	vx, vy = data[..., 0], data[..., 1]
	vx_D = np.concatenate((vx[..., 1:, :], np.expand_dims(vx[..., -1, :], -2)), axis=-2)
	vx_U = np.concatenate((np.expand_dims(vx[..., 0, :], -2), vx[..., :-1, :]), axis=-2)
	vy_R = np.concatenate((vy[..., :, 1:], np.expand_dims(vy[..., :, -1], -1)), axis=-1)
	vy_L = np.concatenate((np.expand_dims(vy[..., :, 0], -1), vy[..., :, :-1]), axis=-1)
	vort = (vy_R - vy_L) - (vx_U - vx_D)
	return np.expand_dims(vort, axis=-1)

def main():
	# Add auguments to the parser
	parser = argparse.ArgumentParser(description='Grid-based field visualizer for fluids')
	parser.add_argument('-n', '--filename', type=str,          required=True,           help='file name for field files')
	parser.add_argument('-r', '--range',    type=int, nargs=2, metavar=('BEGIN','END'), help='range of frames to visualize')
	parser.add_argument('-s', '--slice',    type=str,          metavar='AXIS[=N]',      help='slice to visualize (3D only)')
	parser.add_argument('-m', '--maxval',   type=float,                                 help='maximum value for color mapping')
	parser.add_argument('-v', '--vort',     action='store_true',                        help='visualize vorticity field')
	
	# Parse the arguments
	args = parser.parse_args()
	
	# Load the files
	if args.range is None:
		data = load_file(args.filename, args.slice)
	else:
		begin, end = args.range
		data = load_files(args.filename, begin, end, args.slice)

	# Remove NaNs
	np.nan_to_num(data, copy=False)

	# Convert velocity to vorticity
	if args.vort and data.shape[-1] == 2:
		data = convert_to_vorticity(data)

	# Set maximum value for color mapping
	if args.maxval is None:
		maxval = np.linalg.norm(data, axis=-1).max()
	else:
		maxval = args.maxval
	print(f'{COMMT}[*] Maximum value for color mapping: {FIGUR}{maxval}{RESET}')

	# Generate
	if args.range is None:
		generate(data, maxval, args)
	else:
		generate_sequence(data, maxval, args)



if __name__ == '__main__':
	main()