This repository has been archived on 2022-02-01. You can view files and clone it, but cannot push or open issues or pull requests.
openfoam_project/post_process/animate/alpha_water.py

138 lines
3.5 KiB
Python

from pathlib import Path
import re
from multiprocessing.pool import ThreadPool
from pprint import pp
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(
bloc.index,
bloc,
color='k',
zorder=12
) for bloc 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')