streamplot.py

# TODO: Proper visualization of the streamlines

import argparse
import os

import numpy as np
import matplotlib.pyplot as plt
from line_integral_convolutions import lic

import utils

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

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

	# Load the grid file
	grid_vels = np.load(filename + '.npy')

	# Determine grid parameters
	grid_shape = grid_vels.shape[:-1]
	dim = len(grid_shape)
	dx = 1. / grid_shape[-1]
	lengths = np.array(grid_shape[::-1]) * dx
	print(f'{COMMT}[*] Grid shape: {FIGUR}{grid_shape}{COMMT}, dx: {FIGUR}{dx}{RESET}')
	print(f'{COMMT}[*] Lengths: {FIGUR}{lengths}{RESET}')
	
	return grid_vels, lengths

def plot(velocities, lengths, 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', 'streamlines', data_name))
	os.makedirs(dirname, exist_ok=True)
	filename = os.path.join(dirname, field_name)
	print(f'{COMMT}[*] Saving result to {filename}.png{RESET}')

	# Set parameters
	l_bounds = np.zeros(2)
	u_bounds = lengths
	# vels = np.linalg.norm(velocities, axis=-1)
	#maxval = vels.max() if args.maxval is None else args.maxval
	# print(f'{COMMT}[*] Maximum value for color mapping: {FIGUR}{maxval}{RESET}')

	# Generate image
	vfield = np.transpose(velocities, (2, 0, 1))
	sfield = lic.compute_lic_with_postprocessing(
		vfield=vfield,
		bool_periodic_BCs=False,
		streamlength=None if args.streamlen is None else args.streamlen,
		filter_sigma=3.0,
		num_iterations=3,
		num_repetitions=1)
	fig, ax = plt.subplots()
	ax.set_aspect('equal')
	ax.imshow(sfield, extent=(l_bounds[0], u_bounds[0], l_bounds[1], u_bounds[1]),
			  origin='lower', cmap='pink')
	ax.set_xlim(l_bounds[0], u_bounds[0])
	ax.set_ylim(l_bounds[1], u_bounds[1])
	ax.xaxis.set_visible(False)
	ax.yaxis.set_visible(False)

	# Save the images
	fig.savefig(filename + '.png', dpi=300, bbox_inches='tight', pad_inches=0.0)
	# fig.savefig(filename + '.eps', dpi=150, bbox_inches='tight', pad_inches=0.0)

def main():
	# Add auguments to the parser
	parser = argparse.ArgumentParser(description='Grid-based stream plotter for fluids')
	parser.add_argument('-n', '--filename' , type=str,   required=True, help='file name for the flow field')
	parser.add_argument('-l', '--streamlen', type=float,                help='stream length')
	
	# Parse the arguments
	args = parser.parse_args()
	
	# Load the files
	velocities, lengths = load_file(args.filename)

	# Plot
	plot(velocities, lengths, args)

if __name__ == '__main__':
	main()