%matplotlib inline
%load_ext autoreload
%autoreload 2

import numpy as np
import matplotlib.pyplot as plt

import json

params = {'N' : 10, 'suite' : 'trilinos', 'iterations' : 100}

with open('params.json', 'w') as fp:
    json.dump(params, fp)

%%writefile fipy_timing.py

"""
Usage: fipy_timing.py [<jsonfile>]

"""

from docopt import docopt
import json
import timeit
import numpy as np
import fipy as fp
import os

arguments = docopt(__doc__, version='Run FiPy timing')
jsonfile = arguments['<jsonfile>']

if jsonfile:
    with open(jsonfile, 'rb') as ff:
        params = json.load(ff)
else:
    params = dict()
    
N = params.get('N', 10)
iterations = params.get('iterations', 100)
suite = params.get('suite', 'trilinos')
sumatra_label = params.get('sumatra_label', '')

attempts = 3

setup_str = '''
import fipy as fp
import numpy as np
np.random.seed(1)
L = 1.
N = {N:d}
m = fp.GmshGrid3D(nx=N, ny=N, nz=N, dx=L / N, dy=L / N, dz=L / N)
v0 = np.random.random(m.numberOfCells)
v = fp.CellVariable(mesh=m)
v0 = np.resize(v0, len(v)) ## Gmsh doesn't always give us the correct sized grid!
eqn = fp.TransientTerm(1e-3) == fp.DiffusionTerm()
v[:] = v0.copy()

import fipy.solvers.{suite} as solvers
solver = solvers.linearPCGSolver.LinearPCGSolver(precon=None, iterations={iterations}, tolerance=1e-100)

eqn.solve(v, dt=1., solver=solver)
v[:] = v0.copy()
'''

timeit_str = '''
eqn.solve(v, dt=1., solver=solver)
fp.parallelComm.Barrier()
'''

timer = timeit.Timer(timeit_str, setup=setup_str.format(N=N, suite=suite, iterations=iterations))
times = timer.repeat(attempts, 1)

if fp.parallelComm.procID == 0:
    filepath = os.path.join('Data', sumatra_label)
    filename = 'data.txt'
    np.savetxt(os.path.join(filepath, filename), times)

!python fipy_timing.py params.json

!more Data/data.txt

!git log -1

%%bash

\rm -rf .smt
smt init smt-demo
smt configure --executable=python --main=fipy_timing.py
smt configure --launch_mode=distributed
smt configure -g uuid
smt configure -c store-diff
smt configure --addlabel=parameters

!mkdir Data

import itertools

nprocs = (1, 2, 4, 8)#
iterations_ = (100,)
Ns = (10, 40)
suites = ('trilinos',)
tag='demo4'

for nproc, iterations, N, suite in itertools.product(nprocs, iterations_, Ns, suites):
    !smt run --tag=$tag -n $nproc params.json N=$N iterations=$iterations suite=$suite

import json
import pandas

!smt export
with open('.smt/records_export.json') as ff:
    data = json.load(ff)

df = pandas.DataFrame(data)

print df

print df[['label', 'duration']]

import json
df = df.copy()
df['nproc'] = df.launch_mode.map(lambda x: x['parameters']['n'])
for p in 'N', 'iterations', 'suite':
    df[p] = df.parameters.map(lambda x: json.loads(x['content'])[p])

columns = ['label', 'nproc', 'N', 'iterations', 'suite', 'tags']
print df[columns].sort('nproc')

import os

datafiles = df['output_data'].map(lambda x: x[0]['path'])
datapaths = df['datastore'].map(lambda x: x['parameters']['root'])
data = [np.loadtxt(os.path.join(x, y)) for x, y in zip(datapaths, datafiles)]
df['run_time'] = [min(d) for d in data]

columns.append('run_time')
print df[columns].sort('nproc')

tag_mask = df.tags.map(lambda x: 'demo4' in x)
df_tmp = df[tag_mask]
m10 = df_tmp.N.map(lambda x: x == 10)
m40 = df_tmp.N.map(lambda x: x == 40)
df_N10 = df_tmp[m10]
df_N40 = df_tmp[m40]
print df_N10[columns].sort('nproc')
print df_N40[columns].sort('nproc')

ax = df_N10.plot('nproc', 'run_time', label='N={0}'.format(df_N10.N.iat[0]))
df_N40.plot('nproc', 'run_time', ylim=0, ax=ax, label='N={0}'.format(df_N40.N.iat[0]))
plt.ylabel('Run Time (s)')
plt.xlabel('Number of Processes')
plt.legend()

df.to_hdf('store.h5', 'df')

store = pandas.HDFStore('store.h5')
print store.df.dependencies