Implementation of benchmark 3

README.md

Implementation of Benchmark 3

Linting and testing

$ pylint *.py
$ flake8 *.py

TODO

• generate upload files from view.py
  • time.csv time data
  • contour plot data
• generate test and lint suite
• Dockerfile
• Fix README
• make main.py a command line tool
  • write to new directory
• expand out for other benchmarks

contour.csv

logo_dendrite.ipynb

main.py

"""Run Benchmark Problem 3 using FiPy

    $ python main.py --folder=my_folder

All the data files will be written to my_folder/
"""
# pylint: disable=no-value-for-parameter

import os
import time

import numpy as np
# pylint: disable=redefined-builtin
from toolz.curried import pipe, do, iterate, curry, map, compose
# pylint: disable=import-error
from fipy import CellVariable, DiffusionTerm, Grid2D, TransientTerm, Variable
# pylint: disable=import-error
from fipy import numerix
from fipy.solvers.pysparse import LinearGMRESSolver as Solver
# from fipy.solvers.trilinos import LinearGMRESSolver as Solver
import click


def get_params():
    """Get the parameters for the simulation
    """
    return dict(
        mm=4.0,
        epsilon_m=0.05,
        theta_0=np.pi / 4.,
        tau_0=1.0,
        diff=10.0,
        w_0=1.,
        delta=-0.3,
        dx=1.0,
        nx=960,
        radius=8.0,
        dt0=0.1,
        center=(0.0, 0.0),
        steps=15001,
        output_frequency=10
    )


def get_mesh(params):
    """Build the mesh
    """
    return Grid2D(nx=params['nx'],
                  ny=params['nx'],
                  dx=params['dx'],
                  dy=params['dx'])


def mask(x_cell, y_cell, params):
    """Mask for the initial solid portion
    """
    return (((x_cell - params['center'][0])**2 +
             (y_cell - params['center'][1])**2) <
            params['radius']**2).value


def get_phase(mesh, params, value=1.):
    """Generate the phase cell variable
    """
    return pipe(
        mesh,
        lambda x: CellVariable(mesh=x, hasOld=True, value=-1.0),
        do(
            lambda x:
            x.setValue(
                value,
                where=mask(mesh.cellCenters[0], mesh.cellCenters[1], params)
            )
        )
    )


def get_tau(phase, params):
    """Tau function for phase transient term
    """
    return pipe(
        phase,
        lambda x: numerix.arctan2(x.grad[1], x.grad[0]),
        lambda x: w_func(x, params) / params['w_0'],
        lambda x: params['tau_0'] * x**2
    )


def get_coupling(params):
    """Calculate the coupling constant, "lambda"
    """
    return params['diff'] * params['tau_0'] / 0.6267 / params['w_0']**2


def get_source(phase, heat, params):
    """Phase field source term
    """
    return (phase - get_coupling(params) * heat * (1 - phase**2)) * \
        (1 - phase**2)


def w_func(theta, params):
    """W function dependent on the orientation
    """
    return params['w_0'] * \
        (1 +
         params['epsilon_m'] *
         numerix.cos(params['mm'] * (theta - params['theta_0'])))


def w_theta(theta, params):
    """Derivative w.r.t theta of the W function
    """
    return -params['w_0'] * params['mm'] * params['epsilon_m'] * \
        numerix.sin(params['mm'] * (theta - params['theta_0']))


def get_phase_diff(phase, params):
    """Diffusion coefficient for the phase equation
    """
    return pipe(
        phase,
        lambda x: numerix.arctan2(x.faceGrad[1], x.faceGrad[0]),
        lambda x: (w_func(x, params), w_theta(x, params)),
        lambda x: x[0]**2 * Variable(value=((1, 0), (0, 1))) -
        x[0] * x[1] * Variable(value=((0, -1), (1, 0)))
    )


def get_heat_eqn(phase, delta_t, params):
    """Make the heat equation
    """
    return TransientTerm() == DiffusionTerm(params['diff']) + \
        (phase - phase.old) / delta_t / 2.


def get_phase_eqn(phase, heat, params):
    """make the phase equation
    """
    return TransientTerm(get_tau(phase, params)) == \
        DiffusionTerm(get_phase_diff(phase, params)) + \
        get_source(phase, heat, params)


def get_equations(mesh, delta_t, params):
    """Make the heat and phase equation systems
    """
    return pipe(
        dict(phase=get_phase(mesh, params),
             heat=CellVariable(mesh=mesh,
                               hasOld=True,
                               value=params['delta'])),
        lambda x: (
            dict(name='heat',
                 eqn=get_heat_eqn(x['phase'], delta_t, params),
                 var=x['heat'],
                 solver=Solver(precon=None)),
            dict(name='phase',
                 eqn=get_phase_eqn(x['phase'], x['heat'], params),
                 var=x['phase'],
                 solver=Solver(precon=None))
        )
    )


@curry
def iterate_(func, times, value):
    """Use toolz iterate function to actually iterate
    rather than return an iterator object
    """
    iter_ = iterate(func, value)
    for _ in range(times):
        next(iter_)
    return next(iter_)


@curry
def map_(func, value):
    """Force evaluation of toolz map
    """
    return list(map(func, value))


@curry
def sweep_output(data, residuals):
    """Output from a sweep
    """
    print('sweep: {4}, {0} res: {1}, {2} res: {3}'.format(
        data['equations'][0]['name'],
        residuals[0],
        data['equations'][1]['name'],
        residuals[1],
        data['sweep_counter']
    ))
    data['sweep_counter'] += 1
    return data


@curry
def one_sweep(delta_t, equation):
    """Do one sweep
    """
    return equation['eqn'].sweep(equation['var'],
                                 dt=float(delta_t.value),
                                 solver=equation['solver'])


def sweep(data):
    """Do one sweep and output results
    """
    return pipe(
        data,
        lambda x: map_(one_sweep(data['dt']), x['equations']),
        sweep_output(data),
    )


@curry
def step_output(folder, data):
    """Output before time step
    """
    print ""
    print "write data to: {0}".format(folder)
    print "timestep: {0}, elapsed time: {1}".format(
        data['step_counter'],
        data['elapsed']
    )
    return data


def update(data):
    """Update after time step
    """
    data['step_counter'] += 1
    data['elapsed'] += float(data['dt'].value)
    data['dt'].setValue(float(data['dt'].value) * 1.0)
    return data


@curry
def timestep(folder, data):
    """Do a time step
    """
    return pipe(
        data,
        step_output(folder),
        dump_data(folder),
        do(lambda x: map_(lambda x: x['var'].updateOld(), x['equations'])),
        iterate_(sweep, 3),
        update
    )


def get_data(params, delta_t):
    """Generate the initial data structure
    """
    return dict(
        dt=delta_t,
        equations=get_equations(mesh=get_mesh(params),
                                delta_t=delta_t,
                                params=params),
        step_counter=0,
        sweep_counter=0,
        elapsed=0.0,
        params=params
    )


@curry
def dump_data(folder, data):
    """Dump data to a file
    """
    if data['step_counter'] % data['params']['output_frequency'] == 0:
        np.savez_compressed(
            os.path.join(
                folder,
                'data{0}.npz'.format(
                    str(data['step_counter']).zfill(7)
                )
            ),
            heat=np.array(data['equations'][0]['var']),
            phase=np.array(data['equations'][1]['var']),
            elapsed_time=float(data['elapsed']),
            step_counter=int(data['step_counter']),
            params=data['params'],
            wall_time=time.time()
        )
    return data


def sequence(*args):
    """Reverse compose
    """
    return compose(*args[::-1])


add_options = sequence(
    lambda x: x.items(),
    map_(lambda x: click.option('--' + x[0], default=x[1])),
    lambda x: compose(*x)
)


@click.command()
@click.option('--folder', default='data', help="name for data directory")
@add_options(get_params())
def run(folder, **params):
    if os.path.exists(folder):
        click.echo(
            "{0} directory already exists, remove to continue".format(folder)
        )
    else:
        os.makedirs(folder)
        iterate_(timestep(folder),
                 params['steps'],
                 get_data(params, Variable(params['dt0'])))


if __name__ == '__main__':
    run()

time.csv

view.py

"""View results for benchmark problem 3

To plot the solid fraction from directory "data_xxx"

    $ python view.py --folder data_xxx solid_fraction

To plot the free energy

    $ python view.py --folder data_xxx free_energy

To plot the tip position

    $ python view.py --folder data_xxx tip_position

To view the latest time step

    $ python view.py --folder data_xxx view_step --latest

To view a particular time step

    $ python view.py --folder data_xxx view_step --step 500

To save the data to CSV files for upload to the website

    $ python view.py --folder data_xxx save_data
"""
# pylint: disable=no-value-for-parameter

import glob
import os

from scipy.interpolate import interp1d
from scipy.optimize import brentq
import numpy as np
import click
from fipy import Viewer, CellVariable, numerix
from toolz.curried import curry, pipe, do, assoc, juxt, get
import matplotlib.pyplot as plt
import progressbar
import pandas

from main import get_mesh, map_, sequence, get_phase, w_func, get_coupling


@click.group()
@click.option('--folder',
              default='data',
              help="the name of the data directory to parse")
@click.pass_context
def cli(ctx, folder):
    """Use a group function to allow subcommands.
    """
    ctx.params['folder'] = folder


def get_filename(step, latest, folder):
    """Get the filename given the step
    """
    if latest:
        return sorted(glob.glob(os.path.join(folder, 'data*.npz')))[-1]
    return os.path.join(folder, 'data{0}.npz'.format(str(step).zfill(7)))


@curry
def get_var(data, name):
    """Make a variale
    """
    return CellVariable(
        get_mesh(data['params'].item()),
        value=data[name],
        name=name
    )


def get_vars(ctx, step, latest):
    return pipe(
        get_filename(step, latest, ctx.parent.params['folder']),
        np.load,
        lambda x: map_(get_var(x), ['heat', 'phase'])
    )


@curry
def save_vars(save, vars):
    if save:
        pipe(
            vars,
            get(1),
            lambda x: dict(phase=np.array(x),
                           x=np.array(x.mesh.cellCenters[0]),
                           y=np.array(x.mesh.cellCenters[1])),
            pandas.DataFrame,
            lambda x: x.to_csv('space.csv', index=False)
        )


# def debug(x):
#     import ipdb; ipdb.set_trace()
#     return x

@curry
def save_contour(save, vars):
    def reshape(arr, nx):
        return np.array(arr).reshape((nx, nx))

    if save:
        pipe(
            vars,
            get(1),
            lambda x: dict(phase=reshape(x, x.mesh.nx),
                           x=reshape(x.mesh.cellCenters[0], x.mesh.nx),
                           y=reshape(x.mesh.cellCenters[1], x.mesh.nx)),
            lambda x: plt.contour(x['x'], x['y'], x['phase'], (0.0,)),
            lambda x: x.collections[0].get_paths()[0].vertices,
            lambda x: dict(x=x[:, 0], y=x[:, 1]),
            pandas.DataFrame,
            lambda x: x.to_csv('contour.csv', index=False)
        )


@cli.command()
@click.option('--step', default=0, help='step to view')
@click.option('--latest/--no-latest',
              default=False,
              help='view the latest result available')
@click.option('--save/--no-save',
              default=False,
              help='save the spatial data to a CSV')
@click.pass_context
def view_step(ctx, step, latest, save):
    """Use Matplotlib to view the phase and heat variables
    """
    return pipe(
        get_vars(ctx, step, latest),
        do(save_vars(save)),
        do(save_contour(save)),
        map_(lambda x: Viewer(x, title=x.name).plot()),
        do(lambda _: click.pause())
    )


def plot2d(arr):
    """Plot a 2D graph
    """
    plt.plot(arr[:, 0], arr[:, 1])
    plt.show()


@curry
def save2d(filename, column_names, data):
    pandas.DataFrame(
        dict(zip(column_names, data.transpose()))
    ).to_csv(filename, index=False)
    click.echo('CSV file written to {0}'.format(filename))


def pbar(items):
    pbar_ = progressbar.ProgressBar(
        widgets=[progressbar.Percentage(),
                 ' ',
                 progressbar.Bar(marker=progressbar.RotatingMarker()),
                 ' ',
                 progressbar.ETA()],
        maxval=len(items)).start()
    for counter, item in enumerate(items):
        yield item
        pbar_.update(counter + 1)
    pbar_.finish()


def read_and_calc(f_calc):
    """Read in the data and return a result
    """
    return sequence(
        lambda x: x.parent.params['folder'],
        lambda x: os.path.join(x, 'data*.npz'),
        glob.glob,
        sorted,
        pbar,
        map_(
            sequence(
                np.load,
                f_calc
            )
        ),
        np.array
    )


def read_and_plot(f_calc):
    """Read in a file and plot using f_calc
    """
    return sequence(
        read_and_calc(juxt(get('elapsed_time'), f_calc)),
        plot2d
    )


def read_and_save(filename, column_names, f_calc):
    """Read in a file and save data to CSV
    """
    return sequence(
        read_and_calc(f_calc),
        save2d(filename, column_names)
    )


def calc_solid_fraction(ctx):
    return np.sum((1 + ctx['phase']) / 2) / ctx['params'].item()['nx']**2


@cli.command()
@click.pass_context
def solid_fraction(ctx):
    """Command to plot the solid fraction
    """
    read_and_plot(calc_solid_fraction)(ctx)


calc_tip_position = sequence(
    lambda x: dict(dx=x['params'].item()['dx'],
                   phase=x['phase'],
                   nx=x['params'].item()['nx']),
    lambda x: dict(xx=x['dx'] * (0.5 + np.arange(x['nx'])),
                   yy=x['phase'][:x['nx']]),
    lambda x: brentq(interp1d(x['xx'], x['yy']), x['xx'][0], x['xx'][-1])
)


def f_chem(phase, heat, params):
    return -phase**2 / 2 \
        + phase**4 / 4 \
        + get_coupling(params) * heat * phase * \
        (1 - 2 * phase**2 / 3 + phase**4 / 5)


def w_func_(phase, params):
    return w_func(
        numerix.arctan2(phase.grad[1], phase.grad[0]),
        params
    )


def grad_mag(phase):
    return phase.grad[1]**2 + phase.grad[0]**2


def free_energy_func(phase, heat, params):
    return np.array(
        (
            0.5 * w_func_(phase, params) * grad_mag(phase)
            + f_chem(phase, heat, params)
        ) * params['dx']**2
    ).sum()


calc_free_energy = sequence(
    lambda x: assoc(x, 'params', x['params'].item()),
    lambda x: assoc(x, 'mesh', get_mesh(x['params'])),
    lambda x: assoc(x,
                    'phase',
                    get_phase(x['mesh'], x['params'], value=x['phase'])),
    lambda x: assoc(x,
                    'heat',
                    CellVariable(mesh=x['mesh'], value=x['heat'])),
    lambda x: free_energy_func(x['phase'], x['heat'], x['params'])
)


@cli.command()
@click.pass_context
def tip_position(ctx):
    """Command to plot the tip position versus time
    """
    read_and_plot(calc_tip_position)(ctx)


@cli.command()
@click.pass_context
def free_energy(ctx):
    """Command to calculate the free energy
    """
    read_and_plot(calc_free_energy)(ctx)


@cli.command()
@click.pass_context
def save_time_data(ctx):
    read_and_save(
        'time.csv',
        ['time', 'solid_fraction', 'tip_position', 'free_energy'],
        juxt(get('elapsed_time'),
             calc_solid_fraction,
             calc_tip_position,
             calc_free_energy)
    )(ctx)


if __name__ == '__main__':
    cli()