File diff 000000000000 → d6faa5ffcedf
python/analysis.py
Show inline comments
 
new file 100644
 
#! /cwi/bin/python
 

	
 
from csplot import *
 
import sys
 

	
 
# Update and use the corresponding SciPy functions
 
#from Scientific.Functions import Interpolation
 
import math
 

	
 
import numpy as num
 
import scipy as sci
 
from scipy.optimize.optimize import fmin_powell as fmin
 
from scipy.optimize.minpack import leastsq
 
from optparse import OptionParser, Option, OptionValueError
 

	
 
INFTY = 1e10
 

	
 
try:
 
    BASE_PATH = os.environ['CSTREAM_DEFAULT_OUT_PATH']
 
except KeyError:
 
    BASE_PATH = '/export/scratch1/luque/cstream/'
 

	
 
def max_in_axis(lz, m, sign, zmin=None, zmax=None):
 
    """ Calculates and returns the maximum of a var in the axis and
 
    its position.    """
 

	
 
    max_i = -1
 

	
 
    # No infinity symbol in some old versions of Python
 
    max_v = 0
 
    min_v = 1e100
 
    corr = 0
 

	
 
    dz = lz[1] - lz[0]
 
    
 
    if zmin is None:
 
        zmin = 0
 

	
 
    if zmax is None:
 
        zmax = INFTY
 
        
 
    for i in range(len(lz)):
 
        if lz[i] < zmin or lz[i] > zmax:
 
            continue
 
        
 
        v = m[i]
 
        if abs(v) > abs(max_v) and sign * v >= 0:
 
            max_i = i
 
            max_v = v
 

	
 
        if abs(v) < abs(min_v) and sign * v >= 0:
 
            min_i = i
 
            min_v = v
 

	
 
    if max_i >= 1:
 
        corr, max_v = quadratic_max(m[max_i - 1: max_i + 2])
 

	
 
    return lz[max_i] + corr * dz, max_v, min_v
 

	
 

	
 
def max_off_axis(lr, lz, m, sign):
 
    Z, R = num.meshgrid(lz, lr)
 
    indx = (sign * m).argmax(axis=None)
 
    return R.flat[indx], Z.flat[indx]
 

	
 
    
 
def find_front(var, grid, path=BASE_PATH, z0=None, ret_z0=False,
 
               fmin=0, fmax=INFTY, rmin=None, rmax=INFTY,
 
               pre_margin=400, post_margin=20, zmin=None, zmax=None,
 
               shift_origin=True, sign=None):
 

	
 

	
 
    if sign is None:
 
        dict_signs = {'charge': -1, 'ez': 0, 'eabs': 0}
 
        sign = dict_signs.get(var, 0)
 
        
 
    data = []
 
    lr, lz, m = load_var(var, grid, path=path)
 
    
 
    if rmin is None:
 
        rmin = 0.5 * (lr[1] + lr[2])
 
        
 
    if z0 is None:
 
        z0, max_var, _ = max_in_axis(lz, m[0, :], sign,
 
                                     zmin=zmin, zmax=zmax)
 

	
 
    fmin, fmax = abs(max_var) * fmin, abs(max_var) * fmax
 
    print fmin, fmax
 

	
 
    
 
    if shift_origin:
 
        z0_shift = z0
 
    else:
 
        z0_shift = 0
 

	
 
    for j in xrange(len(lz)):
 
        # This is a small WTF, really
 
        if lz[j] < z0 - pre_margin or lz[j] > z0 + post_margin:
 
            continue
 
        
 
        r, f = find_max_row(lr, sign * m[:, j])
 
        
 
        if abs(f) >= fmin and abs(f) < fmax and r > rmin and r < rmax:
 
            data.append((lz[j] - z0_shift, r, f))
 
        #else:
 
        #    print ("f out of (%f, %f) or r out of (%f, %f)" %
 
        #           (fmin, fmax, rmin, rmax))
 
            
 
    if ret_z0:
 
        return data, z0
 
    else:
 
        return data
 

	
 
class ThresholdNotFound(ValueError):
 
    pass
 
    
 
def find_envelope(grid, path=BASE_PATH, eps=0.0005, tail=350.0):
 
    var = 'charge'
 
    lr, lz, m = [num.array(a) for a in load_var(var, grid, path=path)]
 
    
 
    m = -m
 
    front = num.zeros((len(lz), 2), 'd')
 
    found = num.array([False for i in xrange(len(lz))])
 

	
 
    dz = lz[1] - lz[0]
 
    ntail = int(tail / dz)
 
    
 
    lr_inv = lr[-1::-1]
 
                      
 
    for i in xrange(len(lz)):
 
        try:
 
            a, rmax  = find_threshold(eps, m[-1::-1, i], lr_inv)
 
            front[i, 0:2] = lz[i], rmax
 
            found[i] = True
 
        except ThresholdNotFound:
 
            pass
 

	
 
    front = front.compress(found, axis=0)
 
    zlen, _ = front.shape
 
    if zlen > ntail:
 
        zlen = ntail
 
        
 
    return front[-zlen:, :]
 

	
 
def find_threshold(threshold, a, r):
 
    for ia, ir in zip(a, r):
 
        if ia >= threshold:
 
            return ia, ir
 

	
 
    raise ThresholdNotFound
 

	
 
def simple_find_front(var, grid, path=BASE_PATH, tail=350.0, eps=1e-4):
 
    var = 'charge'
 
    lr, lz, m = [num.array(a) for a in load_var(var, grid, path=path)]
 
    m = num.abs(m)
 

	
 
    maxes = m.argmax(0).squeeze()
 
    front = num.array([(lz[i], lr[maxes[i]])
 
                       for i in xrange(len(lz))
 
                       if maxes[i] > 0 and m[maxes[1], i] > eps])
 

	
 
    zlen, _ = front.shape
 
    if zlen > ntail:
 
        zlen = ntail
 
        
 
    return front[-zlen:, :]
 

	
 

	
 
def fit_radius(var, grid, sign, path=BASE_PATH, zmax=None):
 
    front, z0 = find_front(var, grid, path=path, pre_margin=400.0, sign=sign,
 
                           post_margin=400.0, ret_z0=True,
 
                           fmax=1.0, fmin=0.5, zmax=zmax)
 

	
 
    if not front:
 
        return num.NAN
 
    
 
    front = num.array(front)
 
    
 
    def errors(R):
 
        return num.sqrt(front[:, 1]**2 + (front[:, 0] - R)**2) - abs(R)
 

	
 
    try:
 
        x, _ = leastsq(errors, 20)
 
    except TypeError:
 
        raise
 
        #ipshell()
 
        
 
    return x
 

	
 
    
 
def find_max_row(lr, row):
 
    max_f = 0
 
    min_f = 1e100
 
    max_i = 0
 

	
 
    dr = lr[1] - lr[0]
 

	
 
    for i in xrange(len(lr)):
 
        f = row[i]
 
        if abs(f) > abs(max_f):
 
            max_i = i
 
            max_f = f
 

	
 
    if max_i >= 1 and max_i < len(lr) - 1:
 
        corr, max_f = quadratic_max(row[max_i - 1: max_i + 2])
 
    else:
 
        corr = 0
 
        
 
    return lr[max_i] + corr * dr, max_f
 

	
 
    
 

	
 
def inward_norm(nx, ny, y):
 
    """ gets a normal vector from the array we created making sure that
 
    it points always inward"""
 

	
 
    if num.sign(y) == num.sign(ny):
 
        nx, ny = -nx, -ny
 

	
 
    return nx, ny
 

	
 
def normal_to_curve(m):
 
    """ given an array as [(x0, y1), ...(xN, yN)] returns
 
    [(nx0, ny0), ... (nxN, nyN)] of components of a vector normal to the
 
    curve """
 
    print m.shape
 
    data = zeros(m.shape, typecode='fd')
 
    n = m.shape[0]
 
    
 
    for i in xrange(1, n - 1):                
 
        x, y = m[i, 0], m[i, 1]
 
        
 
        deltas = [(m[j, 0] - m[i, 0], m[j, 1] - m[i, 1])
 
                  for j in (i - 1, i + 1)]
 
        nx, ny = normalize(normal(*deltas))
 
        nx, ny = inward_norm(nx, ny, y)
 
        
 
        data[i, 0] = nx
 
        data[i, 1] = ny
 
        
 
    nx, ny = normalize((-(m[1, 0] - m[0, 0]),
 
                        m[1, 1] - m[0, 1]))
 

	
 
    nx, ny = inward_norm(nx, ny, m[0, 1])
 

	
 
    data[0, 0] = nx
 
    data[0, 1] = ny
 

	
 
    nx, ny = normalize((m[n - 2, 1] - m[n - 1, 1],
 
                        -(m[n - 2, 0] - m[n - 1, 0])))
 

	
 
    nx, ny = inward_norm(nx, ny, m[n - 1, 1])
 
    data[n - 1, 0] = nx
 
    data[n - 1, 1] = ny
 

	
 
    print data
 
    return data
 

	
 
def normal((x0, y0), (x1, y1)):
 
    print x0, y0, x1, y1
 
    
 
    if y0 == 0 and y1 == 0:
 
        return (0.0, 1.0)
 
    
 
    return ((-(x1**2 * y0) + (x0**2 + y0 * (y0 - y1)) * y1)
 
            / (x1 * y0 - x0 * y1),
 
            (x0**2 * x1 + x1 * y0**2 - x0 * (x1**2 + y1**2))
 
            / (-(x1 * y0) + x0 * y1))
 

	
 

	
 
def normalize((x, y)):
 
    a = sqrt(x**2 + y**2)
 
    return (x / a, y / a)
 

	
 

	
 
def quadratic_max(y):
 
    """ Finds a parabola defined by the points (dx i, y[i]) {i = 1, 2, 3}
 
    and returns its maximum. """
 
    
 
    try:
 
        f1 = (y[2] - 2 * y[1] + y[0])
 
        f2 = (y[2] - y[0])
 
    except IndexError:
 
        return num.NAN, num.NAN
 
    
 
    if f1 == 0:
 
        # The points are aligned: no parabola fit
 
        return 0, max(*y)
 
    
 
    return (- f2 / (2 * f1)), y[1] - f2**2 / f1 / 8.0
 
    
 

	
 
def var_interpolate(var, grid, r_staggered=False, z_staggered=False, **kwargs):
 
    lr, lz, var = load_var(var, grid, **kwargs)
 
    
 
    lr = num.array(lr)
 
    lz = num.array(lz)
 

	
 
    if r_staggered:
 
        lr = lr + 0.5 * (lr[1] - lr[0])
 

	
 
    if z_staggered:
 
        lz = lz + 0.5 * (lz[1] - lz[0])
 

	
 
    var = num.array(var)
 

	
 
    lr, var = mirror_var(lr, var)
 
    
 
    f = Interpolation.InterpolatingFunction([lr, lz], var, default=1e20)
 
    return f
 

	
 
def find_front_iter(var, grid, n, **kwargs):
 
    """ After finding a front, improves it by n iterations. """
 
    print kwargs
 
    
 
    m = find_front(var, grid, **kwargs)
 
    m = num.array(m)
 

	
 
    interp = var_interpolate(var, grid, path=kwargs.get('path', None))
 

	
 
    for i in xrange(n):
 
        norm = normal_to_curve(m)    
 
        m = iter_front(m, norm, interp)
 

	
 
    return m
 

	
 

	
 
def iter_front(m, norm, interp, max_step=80, ds=0.05):
 
    mnext = zeros(m.shape, typecode='fd')
 
                  
 
    for p in xrange(m.shape[0]):
 
        z0, r0 = m[p, 0], m[p, 1]
 
        nz, nr = norm[p, 0], norm[p, 1]
 

	
 
        zs = [z0 + ds * i * nz for i in xrange(-max_step, max_step)]
 
        rs = [r0 + ds * i * nr for i in xrange(-max_step, max_step)]
 

	
 
        interp_zr = [interp(r, z) for r, z in zip(rs, zs)]
 

	
 
        max_i, max_e = -1, 0
 

	
 
        for i, e in enumerate(interp_zr):
 
            print i, rs[i], zs[i], e
 
            if e > max_e or max_i < 0:
 
                max_i = i
 
                max_e = e
 
                
 
        mnext[p, 0] = zs[max_i]
 
        mnext[p, 1] = rs[max_i]
 

	
 
    return mnext
 

	
 
def fit_saffman(id, grid, var='eabs', width=64.0, fix_vertex=False, rmax=300):
 
    path = os.path.join(BASE_PATH, id)
 

	
 
    front = find_front(var, grid,
 
                       path=path, fmin=0.0001, rmin=1.1, rmax=rmax,
 
                       zmin=0, zmax=4000, shift_origin=False,
 
                       pre_margin=100, post_margin=5)
 
    front = num.array(front)
 

	
 
    #front = find_envelope(grid, path=path)
 
    #print front
 

	
 
    vert_x, vert_y, _ = front[num.abs(front).argmax(0)[0]]
 
    #vert_x, vert_y = front[num.abs(front).argmax(0)[0]]
 
    
 
    print "Vertex found at (%f, %f)" % (vert_x, vert_y)
 

	
 
    L = 2 * width
 
    
 
    def chi2((lam, C)):
 
        #A = 2 * width / math.pi
 
        #y = A * num.arccos(num.exp(((front[:, 0] - C) / A)).clip(-10000.0, 1.0))
 
        x = front[:, 0]
 
        y = (lam * L / (2 * math.pi)
 
             * num.arccos((2 * num.exp(2 * math.pi * x
 
                                      / (L * (1 - lam)) - C) - 1.0)
 
                          .clip(-1.0, 1.0)))
 
                        
 
        return ((y - front[:, 1]) ** 2).sum()
 

	
 
    if not fix_vertex:
 
        lam, C =  fmin(chi2, (0.5, vert_x), maxfun=100)
 
    else:
 
        def chi2_partial(lam):
 
            return chi2((lam, vert_x))
 

	
 
        lam, C =  fmin(chi2_partial, (0.5,), maxfun=100), vert_x
 
        
 

	
 
    print "\tlambda = %f\n\tx0 = %f" % (lam, C)
 
    return lam, C
 

	
 

	
 
def find_emax(rid, grid, zinterval=None, path=BASE_PATH):
 
    var = rid + ':eabs'
 
    lr, lz, m = [num.array(a) for a in load_var(var, grid, path=path)]
 

	
 
    z0 = lz[0]
 
    dz = lz[1] - z0
 

	
 
    if zinterval is not None:
 
        zmin, zmax = zinterval
 
        izmin, izmax = [int((z - z0) / dz) for z in zmin, zmax]
 

	
 
        izmin = max(0, izmin)
 
        izmax = min(m.shape[1] - 1, izmax)
 

	
 
        m = m[:, izmin:izmax]
 
        
 
    return max(m.flat)
 

	
 
def find_diameters(var, grid, neg_z, pos_z, path=BASE_PATH):
 
    lr, lz, m = load_var(var, grid, path=path)
 
            
 
    m = m * m
 
    diameters = num.dot(lr, m) / num.dot(ones(lr.shape), m)
 
    table = num.concatenate((lz.reshape(len(lz), 1),
 
                             diameters.reshape(len(lz), 1)), axis=1)
 

	
 
    table = table.compress(num.logical_and(num.greater(lz, pos_z),
 
                                           num.less(lz, neg_z)), axis=0)
 
    #ipshell()
 
    dfile = os.path.join(path, 'diameters.%s.tsv' % grid)
 
    try:
 
        neg_diam = num.max(table[-40:, 1])
 
    except ValueError:
 
        neg_diam = num.NAN
 

	
 
    try:
 
        pos_diam = num.max(table[:40, 1])
 
    except ValueError:
 
        pos_diam = num.NAN
 

	
 
    num.savetxt(os.path.join(path, 'diameters.%s.tsv' % grid), table)
 

	
 
    return neg_diam, pos_diam
 

	
 
def posneg_fronts(grid, path=BASE_PATH, var='charge', zmax=None):
 
    lz, m = load_axis(var, grid, path=path)
 
    neg_z, neg_max, neg_min = max_in_axis(lz, m, -1, zmax=zmax)
 
    pos_z, pos_max, pos_min = max_in_axis(lz, m, 1, zmax=zmax)
 
    return [neg_z, pos_z, neg_max, pos_max]
 

	
 

	
 
def integrate(grid, path=BASE_PATH, var='charge', sign=1.0, threshold=None,
 
              zrange=None):
 
    lr, lz, m = load_var(var, grid, path=path)
 
    if zrange is not None:
 
        z0, z1 = zrange
 
        if not num.isnan(z0) and not num.isnan(z1):
 
            lzmask = num.less(lz, z1) * num.greater(lz, z0)
 
            m = m.compress(lzmask, axis=1)
 
            lz = lz.compress(lzmask)
 
        
 
    if threshold is not None:
 
        if sign > 0:
 
            thres = num.greater
 
            mmax = max(num.ravel(m))
 
        else:
 
            thres = num.less
 
            mmax = min(num.ravel(m))
 
        
 
        m = m * thres(m, threshold * mmax)
 
        
 
    dr = lr[1] - lr[0]
 
    dz = lz[1] - lz[0]
 
    
 
    dV = lr.reshape((lr.size, 1)) * dr * dz
 
    m = m * dV
 
    return sum(m.flat)
 
    
 
    
 
# Commands
 
def cmd_fields(*args, **kwargs):
 
    try:
 
        path = os.path.join(BASE_PATH, args[0])
 
    except IndexError:
 
        path = '.'
 

	
 
    ofile = os.path.join(path, 'field.dat')
 
    
 
    fp = open(ofile, "w")
 
            
 
    pattern = kwargs.get('pattern', 'C???')
 
    subs = kwargs.get('subs', '')
 
    last = kwargs.get('last', False)
 

	
 
    get_grids = all_grids
 
    
 
    if last:
 
        get_grids = last_grid
 

	
 
    
 
    for grid in get_grids(pattern=pattern, path=path):
 
        print path, grid + subs
 
        lz, m = load_axis('ez', grid + subs, path=path)
 
        z, field_max, field_min = max_in_axis(lz, m, -1)
 
        if grid != 'C000':
 
            fp.write("%d\t%f\t%.18e\t%.18e\n" % (int(grid[1:]), z,
 
                                                 field_max, field_min))
 
        
 
            print int(grid[1:]), z, field_max, field_min
 
            fp.flush()
 

	
 
    fp.close()
 

	
 
def cmd_fields_last(*args, **kwargs):
 
    _d = kwargs.copy()
 
    _d['last'] = True
 
    
 
    cmd_fields(*args, **_d)
 

	
 
def cmd_posneg(*args, **kwargs):
 
    try:
 
        path = os.path.join(BASE_PATH, args[0])
 
    except IndexError:
 
        path = '.'
 

	
 
    ofile = os.path.join(path, 'posneg.dat')
 
    
 
    fp = open(ofile, "w")
 
    pattern = kwargs.get('pattern', 'C???')
 
    for grid in all_grids(pattern=pattern, path=path):
 
        print path, grid
 
        try:
 
            neg_z, pos_z, neg_max, pos_max = posneg_fronts(grid, path=path)
 
        
 
            if grid != 'C000':
 
                fp.write("%d\t%f\t%f\t%f\t%f\n"
 
                         % (int(grid[1:]), neg_z, pos_z, neg_max, pos_max))
 
        
 
                print int(grid[1:]), neg_z, pos_z, neg_max, pos_max
 
                fp.flush()
 
        except IOError:
 
            print >> sys.stderr, "IOError exception in %s" % grid
 

	
 
    fp.close()
 

	
 

	
 
def cmd_complete(*args, **kwargs):
 
    try:
 
        path = os.path.join(BASE_PATH, args[0])
 
    except IndexError:
 
        path = '.'
 

	
 
    ofile = os.path.join(path, kwargs.get('ofile', 'analysis.dat'))
 
    
 
    fp = open(ofile, "w")
 
    pattern = kwargs.get('pattern', 'C???')
 
    sub = kwargs.get('sub', '')
 
    zmax = kwargs.get('zmax', None)
 
    dt = kwargs.get('dt', 1.0)
 
    
 
    try:
 
        set_def_rescaling(kwargs['rescaling'])
 
    except KeyError:
 
        pass
 
    
 
                          
 
    
 
    fmt_str = None
 

	
 
    varnames = ['t',
 
                'Zc1', 'Zc2', 'Qmin', 'Qmax',
 
                'Ze1', 'Ze2', 'Emin', 'Emax',
 
                'V1', 'V2', 'R1', 'R2',
 
                'Q', 'Q1', 'Q2']
 
                
 
    fp.write("\t".join(varnames) + "\n")
 
    
 
    old_c_fronts = None
 
    pos_z0, neg_z0 = None, None
 
    
 
    for grid in all_grids(pattern=pattern, path=path):
 
        sgrid = grid + sub
 
        clear_load_var_memoize()
 
        print path, sgrid
 
        try:
 
            it = int(grid[1:]) * dt
 
            itms = [it]
 
            c_fronts = posneg_fronts(sgrid, path=path,
 
                                     var='charge', zmax=zmax)
 

	
 
            if pos_z0 is None or neg_z0 is None:
 
                neg_z0, pos_z0 = [c_fronts[i] for i in 0, 1]
 

	
 
            c_fronts[0] -= neg_z0
 
            c_fronts[1] -= pos_z0
 
                
 
            if old_c_fronts:
 
                vminus, vplus = [(c_fronts[i] - old_c_fronts[i]) / (it - old_it)
 
                                 for i in 0, 1]
 
            else:
 
                vminus, vplus = num.NAN, num.NAN
 

	
 
            old_it = it
 
            old_c_fronts = c_fronts
 
            
 
            neg_r, pos_r = [fit_radius('charge', sgrid, s,
 
                                             path=path, zmax=zmax)
 
                                  for s in -1, 1]
 
            
 
            #neg_diam, pos_diam = find_diameters('electrons', grid,
 
            #                                    c_fronts[0], c_fronts[1],
 
            #                                    path=path)
 

	
 
            itms.extend(c_fronts)
 

	
 
            e_fronts = posneg_fronts(sgrid, path=path, var='ez', zmax=zmax)
 
            e_fronts[0] -= neg_z0
 
            e_fronts[1] -= pos_z0
 

	
 
            itms.extend(e_fronts)
 
            itms.extend((vminus, vplus, neg_r, pos_r))
 
            itms.append(integrate(sgrid, path=path, var='charge'))
 

	
 
            # This was the old way of calculating the integrated charge:
 
            # neg_zrange=(neg_z0 - 2*neg_r, neg_z0 + neg_r)
 
            # pos_zrange=(pos_z0 - 2*pos_r, pos_z0 + pos_r)
 
            # 
 
            # itms.append(integrate(sgrid, path=path, var='charge',
 
            #                      sign=-1.0, threshold=0.5, zrange=neg_zrange))
 
            # itms.append(integrate(sgrid, path=path, var='charge',
 
            #                       sign=1.0, threshold=0.5, zrange=pos_zrange))
 

	
 
            neg_zrange=(neg_z0 - neg_r, neg_z0 + 0.5 * neg_r)
 
            pos_zrange=(pos_z0 - pos_r, pos_z0 + 0.5 * pos_r)
 
             
 
            itms.append(integrate(sgrid, path=path, var='charge',
 
                                  zrange=neg_zrange))
 
            itms.append(integrate(sgrid, path=path, var='charge',
 
                                  zrange=pos_zrange))
 
            
 
            if fmt_str == None:
 
                fmt_str = "\t".join("%g" for x in itms) + "\n"
 
                
 
            if int(grid[1:4]) >= 2:
 
                fp.write(fmt_str % tuple(itms))
 
                fp.flush()
 
        
 
                print "\t".join(str(x) for x in itms)
 
        except IOError:
 
            print >> sys.stderr, "IOError exception in %s" % grid
 

	
 
    fp.close()
 

	
 

	
 
    
 
def cmd_diameters(runid, var, **kwargs):
 
    try:
 
        path = os.path.join(BASE_PATH, runid)
 
    except IndexError:
 
        path = '.'
 

	
 
    ofile = os.path.join(path, 'diameters.dat')
 
    fp = open(ofile, "w")
 

	
 
    pattern = kwargs.get('pattern', 'C???')
 
    for grid in all_grids(pattern=pattern, path=path):
 
        try:
 

	
 
            neg_z, pos_z, neg_max, pos_max = posneg_fronts(grid, path=path)
 
            neg_diam, pos_diam = find_diameters(var, grid, neg_z, pos_z, path=path)
 
            print grid
 

	
 
            if grid != 'C000':
 
                fp.write("%d\t%f\t%f\t%f\t%f\t%f\t%f\n"
 
                         % (int(grid[1:]), neg_z, pos_z,
 
                            neg_max, pos_max, neg_diam, pos_diam))
 
        
 
                print int(grid[1:]), neg_z, pos_z, neg_max, pos_max, neg_diam,\
 
                      pos_diam
 
                fp.flush()
 
        except IOError:
 
            print >> sys.stderr, "IOError exception in %s" % grid
 

	
 
    fp.close()
 

	
 

	
 
def cmd_voltage(*args, **kwargs):
 
    try:
 
        path = os.path.join(BASE_PATH, args[0])
 
    except IndexError:
 
        path = '.'
 

	
 
    ofile = os.path.join(path, 'voltage.dat')
 
    
 
    fp = open(ofile, "w")
 
    pattern = kwargs.get('pattern', 'C???')
 
    for grid in all_grids(pattern=pattern, path=path):
 
        try:
 
            lz, m = load_axis('ez', grid, path=path)
 
            voltage = - sum(m) * (lz[1] - lz[0])
 
            s = "%d\t%f\n" % (int(grid[1:]), voltage)
 
            fp.write(s)
 
            fp.flush()
 
            print s
 
            
 
        except IOError:
 
            print >> sys.stderr, "IOError exception in %s" % grid
 

	
 

	
 
def cmd_front(grid, ofile):
 
    fp = open(ofile, "w")
 

	
 
    id, grid  = grid.split(':')
 
    path = os.path.join(BASE_PATH, id)
 

	
 
    print grid, path
 
    
 
    front = find_front('charge', grid, path=path, fmin=0.0001, rmin=2.1, rmax=300,
 
                       zmin=0, zmax=2000, shift_origin=False,
 
                       pre_margin=200, post_margin=20)
 
    for z, r, f in front:
 
        fp.write("%f\t%f\t%f\n" % (z, r, f))
 
    
 
    fp.close()
 

	
 
def cmd_fit_saffman(grid):
 
    id, grid  = grid.split(':')
 
    return fit_saffman(id, grid)
 
    
 
def cmd_envelope(grid, ofile):
 
    id, grid  = grid.split(':')
 
    path = os.path.join(BASE_PATH, id)
 

	
 
    front = find_envelope(grid, path=path)
 

	
 
    fp = open(ofile, "w")
 
    for r, z in front:
 
        fp.write("%f\t%f\n" % (r, z))
 
    fp.close()
 
    
 
    return front
 

	
 
def cmd_simple_front(var, grid):
 
    id, grid  = grid.split(':')
 
    path = os.path.join(BASE_PATH, id)
 

	
 
    front = simple_find_front(var, grid, path=path)
 
    for r, z in front:
 
        print r, z
 
        
 
    return front
 

	
 
def cmd_iterfront(grid, n, ofile):
 
    id, grid  = grid.split(':')
 
    path = os.path.join(BASE_PATH, id)
 
    
 
    fp = open(ofile, "w")
 
    n = int(n)
 
    
 
    front = find_front_iter('eabs', grid, n, path=path,
 
                            fmin=0.001, rmin=1, rmax=300,
 
                            zmin=300, zmax=600, shift_origin=False,
 
                            pre_margin=250, post_margin=20)
 
    for z, r, f in front:
 
        fp.write("%f\t%f\t%f\n" % (z, r, f))
 
    
 
    fp.close()
 

	
 

	
 
def cmd_max_off_axis(runid, var, sgn, **kwargs):
 
    try:
 
        path = os.path.join(BASE_PATH, runid)
 
    except IndexError:
 
        path = '.'
 

	
 
    ofile = os.path.join(path, 'max-off-%s.dat' % var)
 
    fp = open(ofile, "w")
 

	
 
    pattern = kwargs.get('pattern', 'C???')
 
    sgn = int(sgn)
 
    for grid in all_grids(pattern=pattern, path=path):
 
        try:
 
            lr, lz, m = load_var(var, grid, path=path)
 
        except IOError:
 
            print >> sys.stderr, "IOError exception in %s" % grid
 
            continue
 

	
 
        r, z = max_off_axis(lr, lz, m, sgn)
 
        if grid != 'C000':
 
            fp.write("%d\t%f\t%f\n" % (int(grid[1:]), r, z))
 
        
 
            print int(grid[1:]), r, z
 
            fp.flush()
 

	
 
    fp.close()
 

	
 
def cmd_shell():
 
    from IPython.Shell import IPShellEmbed
 
    ipshell = IPShellEmbed([]) 
 
    ipshell()
 

	
 
def check_eval_float(option, opt, value):
 
    eval_value = eval(value)
 
    try:
 
        return float(eval_value)
 
    except ValueError:
 
        raise OptionValueError(
 
            "option %s: invalid float value: '%s' -> %r"
 
            % (opt, value, eval_value))
 

	
 

	
 

	
 
class ExtOption(Option):
 
    TYPES = Option.TYPES + ("eval_float",)
 
    TYPE_CHECKER = Option.TYPE_CHECKER.copy()
 
    TYPE_CHECKER["eval_float"] = check_eval_float
 

	
 
    
 
    
 
def main():
 
    parser = OptionParser("", option_class=ExtOption)
 

	
 
    parser.add_option('--dt', '', action='store',
 
                      dest='dt',
 
                      type='eval_float',
 
                      default=1.0,
 
                      help="dt")
 

	
 
    parser.add_option('--with-units', '', action='store',
 
                      dest='units',
 
                      default=None,
 
                      help="Normalize with units")
 

	
 

	
 
    parser.add_option('--pattern', '', action='store',
 
                      dest='pattern', type='str',
 
                      default='C???',
 
                      help='Pattern of the analyzed grids')
 
    
 

	
 
    parser.add_option('--sub', '', action='store',
 
                      dest='sub', type='str',
 
                      default="",
 
                      help='Sub-grid postfix')
 

	
 
    parser.add_option('--ofile', '', action='store',
 
                      dest='sub', type='str',
 
                      help='Output file')
 

	
 

	
 
    parser.add_option('--zmax', '', action='store',
 
                      dest='zmax', type='eval_float',
 
                      default=None,
 
                      help='Sub-grid postfix')
 

	
 
    parser.add_option('--rescaling', '', action='store',
 
                      dest='rescaling', type='str',
 
                      default=None,
 
                      help='Rescaling to use')
 

	
 

	
 
    main = sys.modules[__name__]
 

	
 
    opts, args = parser.parse_args()
 
    
 
    try:
 
        cmd = sys.argv[1].replace('-', '_')
 
        process_method = getattr(main, 'cmd_%s' % cmd)
 
    except AttributeError:
 
        print "Dont know how to process command `%s`" % args[0]
 
        sys.exit(-1)
 
        
 
    process_method(*args[1:], **opts.__dict__)
 

	
 

	
 

	
 
if __name__ == "__main__":
 
    main()