from pathlib import Path import re from multiprocessing.pool import ThreadPool from time import time import logging import configparser from contextlib import redirect_stdout from io import StringIO import matplotlib as mpl import matplotlib.pyplot as plt from matplotlib import cm from matplotlib import patches import matplotlib.animation as animation from scipy.interpolate import griddata import numpy as np import pandas as pd mpl.use('agg') import fluidfoam #import stl def animate(config_path): tt0 = time() config = configparser.ConfigParser() config.read(config_path) log = logging.getLogger('alpha_water') logging.basicConfig(encoding='utf-8', level=config['main']['logging']) log.info('Starting') root = Path(config['main']['root']) t = [] log.info(f'Finding timesteps in {root}') for item in root.iterdir(): if item.is_dir() and re.match(r'^[0-9]+(\.[0-9]+)?$', item.name): t.append(item.name) timesteps = pd.Series(t, dtype='string', name='t')\ .sort_values(key=np.float64) log.info('Reading mesh') fluidfoam_log = StringIO() with redirect_stdout(fluidfoam_log): mesh = fluidfoam.readmesh(str(root)) log.info('Reading alpha.water for each time step') with ThreadPool(4) as pool: with redirect_stdout(fluidfoam_log): aw = zip(timesteps, pool.map( lambda t:fluidfoam.readscalar(root, t, 'alpha.water'), timesteps)) log.info('Building alpha_water') alpha_water = pd.DataFrame(dict(aw)).round(2) del aw alpha_water['x'] = mesh[0] alpha_water['y'] = mesh[2] with pd.HDFStore( config['data']['path'], mode='r', ) as hdf: bathy = hdf.get('bathy') rubble = hdf.get('rubble') blocs = {} for bloc in config['data']['blocs'].split(','): blocs[bloc] = hdf.get(bloc) log.info('Starting plot') fig, ax = plt.subplots(figsize=(19.2,10.8), dpi=100) base_artists = [ ax.fill_between( bathy.index, bathy, np.floor(bathy.min()/10)*10, color='k', zorder=10 ), ax.fill_between( rubble.index, rubble, rubble.min(), color='k', alpha=.1, lw=2, zorder=11 ), *[ ax.fill( a.index.values, a.values, color='k', zorder=12 ) for a in blocs.values() ] ] log.info('Generating artists') with ThreadPool(4) as pool: ims = pool.map( lambda t:ax.tricontourf( alpha_water.x, alpha_water.y, alpha_water[t], cmap=cm.Blues, levels=[0.1,0.9], vmin=0, vmax=1, zorder=1, extend='both' ), timesteps ) ax.set( aspect='equal', xlim=(alpha_water.x.min(), alpha_water.x.max()), ylim=alpha_water.y.min() ) fig.colorbar(ims[-1], orientation='horizontal') fig.tight_layout() log.info('Generating animation') ani = animation.ArtistAnimation( fig, [im.collections for im in ims], interval=1000/10, blit=True, repeat=True ) ani.save(config['main']['out']) log.info('Program ended successfully')