source: branches/uq/puq/puq_analyze.py @ 5182

#!/usr/bin/env python
"""
Once submit has finished with the jobs, this function is called to have PUQ
process the results.
"""
import sys
import os
import numpy as np
import h5py
import re
from puq.jpickle import unpickle
import xml.etree.ElementTree as xml
from itertools import product, combinations

# Redirect stdout and stderr to files for debugging.
# Append to the files created in get_params.py
sys.stdout = open("uq_debug.out", 'a')
sys.stderr = open("uq_debug.err", 'a')


# Restore the state of a PUQ session from a HDF5 file.
def load_from_hdf5(name):
    h5 = h5py.File(name, 'r+')
    sw = unpickle(h5['private/sweep'].value)
    sw.fname = os.path.splitext(name)[0]
    h5.close()

    sw.psweep._sweep = sw

    if hasattr(sw.psweep, 'reinit'):
        sw.psweep.reinit()
    return sw


# variable names to labels
def subs_names(varl, h5):
    varlist = []
    for v in varl:
        try:
            lab = h5['/input/params/%s' % v[0]].attrs['label']
        except:
            lab = str(v[0])
        varlist.append(lab)
    return varlist


def plot_resp1(io, resp, name, h5):
    print 'plot_resp1', name
    varlist = subs_names(resp.vars, h5)
    try:
        tname = h5['/output/data/%s' % name].attrs['description']
    except:
        tname = name
    ydata = h5['/output/data/%s' % name].value
    xdata = h5['/input/param_array'].value.T[0]

    title = '%s (Response)' % (tname)
    curve = 'output.curve(response-%s)' % name
    io.put(curve + '.about.label', title)
    io.put(curve + '.about.group', title)
    io.put(curve + '.xaxis.label', varlist[0])
    io.put(curve + '.yaxis.label', tname)
    x = np.linspace(*resp.vars[0][1], num=50)
    y = np.array(resp.eval(x))
    for a, b in zip(x, y):
        io.put(curve + '.component.xy', "%s %s\n" % (a, b), append=1)

    # scatter plot sampled data on response surface
    curve = 'output.curve(response-%s-scatter)' % name
    io.put(curve + '.about.label', 'Data Points')
    io.put(curve + '.about.group', title)
    io.put(curve + '.about.type', 'scatter')
    io.put(curve + '.xaxis.label', varlist[0])
    io.put(curve + '.yaxis.label', tname)
    for a, b in zip(xdata, ydata):
        io.put(curve + '.component.xy', "%s %s\n" % (a, b), append=1)


def plot_resp2(io, resp, name, varl, h5):
    numpoints = 50
    m2d = 'output.mesh('+varl[0][0]+varl[1][0]+')'
    io.put(m2d+'.about.label', '2D Mesh')
    io.put(m2d+'.dim', '2')
    io.put(m2d+'.hide', 'yes')
    io.put(m2d+'.grid.xaxis.numpoints', '%s' % numpoints)
    io.put(m2d+'.grid.yaxis.numpoints', '%s' % numpoints)
    io.put(m2d+'.grid.xaxis.min', '%g' % (varl[0][1][0]))
    io.put(m2d+'.grid.xaxis.max', '%g' % (varl[0][1][1]))
    io.put(m2d+'.grid.yaxis.min', '%g' % (varl[1][1][0]))
    io.put(m2d+'.grid.yaxis.max', '%g' % (varl[1][1][1]))

    varlist = subs_names(varl, h5)
    fname = 'output.field(response-'+name+varlist[0]+varlist[1]+')'
    io.put(fname + '.about.xaxis.label', varlist[0])
    io.put(fname + '.about.yaxis.label', varlist[1])
    try:
        tname = h5['/output/data/%s' % name].attrs['description']
    except:
        tname = name
    title = tname + ' (Response %s)' % ' vs '.join(varlist)
    io.put(fname + '.about.label', title)
    io.put(fname + '.component.mesh', m2d)
    io.put(fname + '.about.view', 'heightmap')
    x = np.linspace(*varl[0][1], num=numpoints)
    y = np.linspace(*varl[1][1], num=numpoints)
    pts = np.array([(b, a) for a, b in product(y, x)])
    allpts = np.empty((numpoints**2, len(resp.vars)))
    for i, v in enumerate(resp.vars):
        if v[0] == varl[0][0]:
            allpts[:, i] = pts[:, 0]
        elif v[0] == varl[1][0]:
            allpts[:, i] = pts[:, 1]
        else:
            allpts[:, i] = np.mean(v[1])
    pts = np.array(resp.evala(allpts))
    for z in pts:
        io.put(fname + '.component.values', '%g\n' % z, append=1)


def plot_resp3(io, resp, name, varl, h5):
    numpoints = 40
    m3d = 'output.mesh(m3d)'
    io.put(m3d+'.about.label', '3D Mesh')
    io.put(m3d+'.dim', '3')
    io.put(m3d+'.hide', 'yes')
    io.put(m3d+'.grid.xaxis.numpoints', '%s' % numpoints)
    io.put(m3d+'.grid.yaxis.numpoints', '%s' % numpoints)
    io.put(m3d+'.grid.zaxis.numpoints', '%s' % numpoints)

    # for now, scale to [0,1]
    io.put(m3d+'.grid.xaxis.min', '0')
    io.put(m3d+'.grid.xaxis.max', '1')
    io.put(m3d+'.grid.yaxis.min', '0')
    io.put(m3d+'.grid.yaxis.max', '1')
    io.put(m3d+'.grid.zaxis.min', '0')
    io.put(m3d+'.grid.zaxis.max', '1')

    varlist = subs_names(varl, h5)
    try:
        tname = h5['/output/data/%s' % name].attrs['description']
    except:
        tname = name
    fname = 'output.field(response-'+name+varlist[0]+varlist[1]+varlist[2]+')'
    io.put(fname + '.about.xaxis.label',
                   '{%s [%.3g - %.3g]}' %
                   (varlist[0], varl[0][1][0], varl[0][1][1]))
    io.put(fname + '.about.yaxis.label',
                   '{%s [%.3g - %.3g]}' %
                   (varlist[1], varl[1][1][0], varl[1][1][1]))
    io.put(fname + '.about.zaxis.label',
                   '{%s [%.3g - %.3g]}' %
                   (varlist[2], varl[2][1][0], varl[2][1][1]))
    title = tname + ' (Response %s)' % ' vs '.join(varlist)
    io.put(fname + '.about.label', title)
    io.put(fname + '.component.mesh', m3d)
    io.put(fname + '.about.description', '3D Field Description')
    io.put(fname + '.about.view', 'vtkvolume')
    x = np.linspace(*varl[0][1], num=numpoints)
    y = np.linspace(*varl[1][1], num=numpoints)
    z = np.linspace(*varl[2][1], num=numpoints)

    # generate points in the right order, with x changing fastest
    pts = np.array([(c, b, a) for a, b, c in product(z, y, x)])
    allpts = np.empty((numpoints**3, len(resp.vars)))
    for i, v in enumerate(resp.vars):
        if v[0] == varl[0][0]:
            allpts[:, i] = pts[:, 0]
        elif v[0] == varl[1][0]:
            allpts[:, i] = pts[:, 1]
        elif v[0] == varl[2][0]:
            allpts[:, i] = pts[:, 2]
        else:
            allpts[:, i] = np.mean(v[1])
    pts = np.array(resp.evala(allpts))
    for z in pts:
        io.put(fname + '.component.values', '%g\n' % z, append=1)


def plot_resp(io, resp, name, h5):
    if resp is not None:
        if len(resp.vars) == 1:
            plot_resp1(io, resp, name, h5)
            return
        if len(resp.vars) >= 3:
            for v1, v2, v3 in combinations(resp.vars, 3):
                plot_resp3(io, resp, name, [v1, v2, v3], h5)

        # plot all combinations of 2 variables
        for v1, v2 in combinations(resp.vars, 2):
            plot_resp2(io, resp, name, [v1, v2], h5)


# Plots probability curves
def plot_pdf_curve(xvals, vname, percent, desc):
    print 'plot_pdf_curve %s %s %s' % (vname, percent, desc)
    # compute upper and lower percentiles
    pm = (100 - percent)/200.0
    pp = 1 - pm

    # collect data into an array
    xarr = np.empty(len(xvals[vname]))
    yp = np.empty(len(xvals[vname]))
    ym = np.empty(len(xvals[vname]))
    for vindex in sorted(xvals[vname].keys()):
        xarr[vindex] = xvals[vname][vindex]
        yp[vindex] = pcurves[vname][vindex].ppf(pp)
        ym[vindex] = pcurves[vname][vindex].ppf(pm)

    curve = xml.SubElement(dout, 'curve', {'id': 'curve_pdf-%s-%s' % (vname, percent)})
    about = xml.SubElement(curve, 'about')
    if percent == 0:
        xml.SubElement(about, 'label').text = "mean"
    else:
        xml.SubElement(about, 'label').text = "middle %s%%" % percent
    xml.SubElement(about, 'group').text = '%s (Probability)' % desc
    comp = xml.SubElement(curve, 'component')
    xy = xml.SubElement(comp, 'xy')
    pts = ""
    for x, y in zip(xarr, yp):
        pts += "%s %s " % (x, y)
    if percent == 0:
        pts += '\n'
    else:
        for x, y in reversed(zip(xarr, ym)):
            pts += "%s %s " % (x, y)
        pts += "%s %s\n" % (xarr[0], yp[0])
    xy.text = pts


def dist(x1, y1, x2, y2, x3, y3): # x3,y3 is the point
    diffx = x2-x1
    diffy = y2-y1
    sqdiff = float(diffx**2 + diffy**2)

    u = ((x3 - x1) * diffx + (y3 - y1) * diffy) / sqdiff

    if u > 1:
        u = 1
    elif u < 0:
        u = 0

    x = x1 + u * diffx
    y = y1 + u * diffy

    dx = x - x3
    dy = y - y3

    return np.sqrt(dx*dx + dy*dy)


# Plots advanced probability curves
def plot_pdf_acurve(hf, acurves, vname, percent, desc):
    print 'plot_pdf_acurve %s %s %s' % (vname, percent, desc)
    # compute upper and lower percentiles
    pm = (100 - percent)/200.0
    pp = 1 - pm
    vlen = len(acurves[vname])

    print 'vlen=', vlen
    # collect data into an array
    xp = np.empty(vlen)
    xm = np.empty(vlen)
    yp = np.empty(vlen)
    ym = np.empty(vlen)

    for vindex in sorted(acurves[vname].keys()):
        print 'vindex=', vindex
        x1 = acurves[vname][vindex][0].ppf(pp)
        x2 = acurves[vname][vindex][0].ppf(pm)
        y1 = acurves[vname][vindex][1].ppf(pp)
        y2 = acurves[vname][vindex][1].ppf(pm)
        xd = hf['/output/data/%s.x[%d]' % (vname, int(vindex))].value
        yd = hf['/output/data/%s.y[%d]' % (vname, int(vindex))].value
        #p1, p2 = get_closest2()
        if int(vindex) == 2:
            print x1, x2, y1, y2
            print 'xd=', repr(xd)
            print 'yd=', repr(yd)

        #xp[vindex] = x1
        #xm[vindex] = x2
        #yp[vindex] = y1
        #ym[vindex] = y2
    return
    curve = xml.SubElement(dout, 'curve', {'id': 'curve_pdf-%s-%s' % (vname, percent)})
    about = xml.SubElement(curve, 'about')
    if percent == 0:
        xml.SubElement(about, 'label').text = "mean"
    else:
        xml.SubElement(about, 'label').text = "middle %s%%" % percent
    xml.SubElement(about, 'group').text = '%s (Probability)' % desc
    comp = xml.SubElement(curve, 'component')
    xy = xml.SubElement(comp, 'xy')
    pts = ""
    for x, y in zip(xarr, yp):
        pts += "%s %s " % (x, y)
    if percent == 0:
        pts += '\n'
    else:
        for x, y in reversed(zip(xarr, ym)):
            pts += "%s %s " % (x, y)
        pts += "%s %s\n" % (xarr[0], yp[0])
    xy.text = pts


def plot_pdf(v, pdf, desc):
    id = '%s (PDF)' % desc
    elem = xml.SubElement(dout, 'curve', {'id': id})
    about = xml.SubElement(elem, 'about')
    xml.SubElement(about, 'label').text = id
    yaxis = xml.SubElement(elem, 'yaxis')
    xml.SubElement(yaxis, 'label').text = 'Probability'

    component = xml.SubElement(elem, 'component')
    xy = xml.SubElement(component, 'xy')
    pts = "%s 0\n" % pdf.x[0]
    for x, y in zip(pdf.x, pdf.y):
        pts += "%s %s\n" % (x, y)
    pts += "%s 0\n" % pdf.x[-1]
    xy.text = pts


sw = load_from_hdf5(sys.argv[1])
sw.analyze()

h5 = h5py.File(sys.argv[1], 'r+')
dtree = xml.parse('run_uq.xml')
droot = dtree.getroot()
dout = droot.find('output')

# curves built from pdfs
pcurves = {}
xvals = {}
acurves = {}

reg1 = re.compile('([ \da-zA-Z]+)\[([ \d]+)\]')
reg2 = re.compile('([ \da-zA-Z]+)\.([xy])\[([ \d]+)\]')

uqtype = h5.attrs['UQtype']
for v in h5[uqtype]:
    rs = unpickle(h5['/%s/%s/response' % (uqtype, v)].value)
    pdf = rs.pdf(fit=False)
    odata = h5['/output/data/%s' % v]

    # For curves built from pdfs, just put them in a dict for now
    if 'x' in odata.attrs:
        matches = reg1.findall(v)
        vname, vindex = matches[0]
        vindex = int(vindex)
        if vname not in pcurves:
            pcurves[vname] = {}
            xvals[vname] = {}
        xvals[vname][vindex] = odata.attrs['x']
        pcurves[vname][vindex] = pdf
    elif reg2.findall(v) != []:
        matches = reg2.findall(v)
        vname, xy, vindex = matches[0]
        if vname not in acurves:
            acurves[vname] = {}
        if vindex not in acurves[vname]:
            acurves[vname][vindex] = {}
        if xy == 'x':
            acurves[vname][vindex][0] = pdf
        else:
            acurves[vname][vindex][1] = pdf
    else:
        desc = h5['/output/data/%s' % v].attrs['description']
        plot_pdf(v, pdf, desc)

# now do pdf curves
for vname in xvals:
    desc = h5['/output/data/%s[0]' % vname].attrs['description']
    plot_pdf_curve(xvals, vname, 0, desc)
    plot_pdf_curve(xvals, vname, 50, desc)
    plot_pdf_curve(xvals, vname, 95, desc)

# for vname in acurves:
#    desc = h5['/output/data/%s.x[0]' % vname].attrs['description']
#    plot_pdf_acurve(h5, acurves, vname, 0, desc)
#    plot_pdf_acurve(h5, acurves, vname, 50, desc)
#    plot_pdf_acurve(h5, acurves, vname, 95, desc)


# If more than one variable, display sensitivity.
# Curves have indexed variables, so skip them.
if len(h5['/input/params']) > 1 and ['[' in x for x in h5[uqtype]].count(False):
    for v in h5[uqtype]:
        if '[' in v:
            continue
        desc = h5['/output/data/%s' % v].attrs['description']
        sens = unpickle(h5['/%s/%s/sensitivity' % (uqtype, v)].value)
        elem = xml.SubElement(dout, 'histogram', {'id': 'sens-%s' % v})
        about = xml.SubElement(elem, 'about')
        xml.SubElement(about, 'label').text = '%s (Sensitivity)' % desc
        xml.SubElement(elem, 'type').text = 'scatter'
        xaxis = xml.SubElement(elem, 'xaxis')
        xml.SubElement(xaxis, 'label').text = 'Parameters'
        yaxis = xml.SubElement(elem, 'yaxis')
        xml.SubElement(yaxis, 'label').text = 'Sensitivity'
        comp = xml.SubElement(elem, 'component')
        xy = xml.SubElement(comp, 'xy')
        pts = ''
        for name in sens:
            n = name[0]
            try:
                n = h5['/input/params/%s' % n].attrs['label']
            except:
                pass
            pts += "\"%s\" %s\n" % (n, name[1]['ustar'])
        xy.text = pts

with open('run_uq.xml', 'w') as f:
    f.write("<?xml version=\"1.0\"?>\n")
    dtree.write(f)


import Rappture
io = Rappture.library('run_uq.xml')

for v in h5[uqtype]:
    if '[' in v:
        continue
    rs = unpickle(h5['/%s/%s/response' % (uqtype, v)].value)
    print 'Plotting response', v
    plot_resp(io, rs, v, h5)

Rappture.result(io)

h5.close()