138 lines
3.5 KiB
Python
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')
|