openfoam_project/post_process/animate/alpha_water.py

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']).expanduser()
    
    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')