path: root/helpers.py

"""
Various Helper Functions.

Copied from Viz2. Fat trimmed.

Volker Hoffmann <volker@cheleb.net>
19 April 2013
"""

from math import pi
# from time import gmtime, strftime
import numpy as np
import os

def prof1d(x, data, xbins):
    """Compute 1D Profile.

    Initially, xbins contains bin edges.
    After binning, we compute bin centres and return those in xbins.
    """

    # Flat Arrays?
    if x.ndim > 1 or data.ndim > 1 or xbins.ndim > 1:
        raise Exception("Non-Flat Input Arrays for Binning. Terminating.")

    # Allocate Memory
    bval = np.zeros(xbins.shape[0])
    bcount = np.zeros_like(bval)

    # Binning
    # print "(%s UTC) Binning 1D Profile.." % strftime("%H:%M:%S", gmtime())
    for idata in range(data.shape[0]):
        for ibin in range(xbins.shape[0] - 1):
            if x[idata] > xbins[ibin] and x[idata] <= xbins[ibin + 1]:
                bval[ibin] += data[idata]
                bcount[ibin] += 1
    # print "(%s UTC) Done Binning." % strftime("%H:%M:%S", gmtime())

    # Drop Last Entry
    bval = bval[:-1]
    bcount = bcount[:-1]

    # Transpose Array
    bval = bval.T
    bcount = bcount.T

    # Normalize
    bcount[bcount == 0] = 1
    bval /= bcount

    return bval

def mkvec(x, y, z):
    """Takes Three Vectors.
    Builds Three Vectors Combining All Possible Combinations.

    Example In:
    - x = [ 0 1 ]
    - y = [ 2 3 ]
    - z = [ 4 5 ]

    Example Out:
    - xx = [ 0 1 0 1 0 1 0 1 ]
    - yy = [ 2 2 3 3 2 2 3 3 ]
    - zz = [ 4 4 4 4 5 5 5 5 ]
    """

    # Length Counters
    xlen = len(x)
    ylen = len(y)
    zlen = len(z)

    xylen = xlen * ylen
    xyzlen = xylen * zlen

    # Allocate Output Arrays
    xx = np.zeros(xyzlen)
    yy = np.zeros_like(xx)
    zz = np.zeros_like(xx)

    # Run Through Arrays
    xcount = 0
    ycount = 0
    zcount = 0
    for ii in range(xyzlen):

        # Assign Values
        xx[ii] = x[xcount]
        yy[ii] = y[ycount]
        zz[ii] = z[zcount]

        # Increment or Reset Counters
        if xcount == xlen - 1:
            xcount = 0
            if ycount == ylen - 1:
                ycount = 0
                zcount += 1
            else:
                ycount += 1
        else:
            xcount += 1

    return xx, yy, zz

def mkpoints_xyz(center=0.5, radius=0.4, thickness=0.2, nx=128, ny=128, nz=64):
    """Generate Sampling Points. Sampling Defined in (x,y,z) Coordinates."""

    # Compute Characters
    cchar = str(center).translate(None, ".")
    rchar = str(radius).translate(None, ".")
    tchar = str(thickness).translate(None, ".")
    nxchar = str(nx)
    nychar = str(ny)
    nzchar = str(nz)

    # Build File Name, File Location
    fname = "points_c%s_r%s_t%s_nx%s_ny%s_nz%s.npz" % \
            ( cchar, rchar, tchar, nxchar, nychar, nzchar )
    floc = "%s/%s" % \
           ( os.path.dirname(os.path.abspath(__file__)), fname )

    # File Exist?
    if os.path.isfile(floc):

        # Load Points
        npz = np.load(floc)
        points_xyz = npz["points_xyz"]
        points_xy  = npz["points_xy"]
        x = npz["x"]
        y = npz["y"]
        z = npz["z"]
        dl = npz["dl"]
        dl = dl[()]     # Recover dict from 0d-array
                        # http://stackoverflow.com/questions/8361561/recover-dict-from-0-d-numpy-array
        npz.close()

    else:

        print ""
        print "*** Generating Sampling Points."
        print ""

        # Generate Points
        x, dx = np.linspace(center - radius, center + radius, nx, retstep=True)
        y, dy = np.linspace(center - radius, center + radius, ny, retstep=True)
        z, dz = np.linspace(center - thickness/2., center + thickness/2., nz, retstep=True)

        # Point Spacing
        dl = {"dx": dx, "dy": dy, "dz": dz}

        # Reshape
        points_xyz = mkpoints_pymses(x, y, z)
        points_xy  = mkpoints_pymses(x, y, np.array([center]))

        # Save
        np.savez(floc, \
                 points_xyz=points_xyz, points_xy=points_xy, \
                 x=x, y=y, z=z, dl=dl)

    return points_xyz, points_xy, x, y, z, dl

def mkpoints_rtz(center=0.5, radius=0.4, thickness=0.2, \
                 nr=128, ntheta=128, nz=64):
    """Generate Sampling Points. Defined in (r,theta,z) Coordinates.
    >>> BEWARE: THE PYMSES SAMPLING MATRIX IS SORTED R,Z,THETA <<<
    """

    # Compute Characters
    cchar = str(center).translate(None, ".")
    rchar = str(radius).translate(None, ".")
    tchar = str(thickness).translate(None, ".")
    nrchar = str(nr)
    nthetachar = str(ntheta)
    nzchar = str(nz)

    # Build File Name, File Location
    fname = "points_c%s_r%s_t%s_nr%s_ntheta%s_nz%s.npz" % \
            ( cchar, rchar, tchar, nrchar, nthetachar, nzchar )
    floc = "%s/%s" % \
           ( os.path.dirname(os.path.abspath(__file__)), fname )

    # File Exist?
    if os.path.isfile(floc):

        # Load Points
        npz = np.load(floc)
        points_rzt = npz["points_rzt"]
        points_rz  = npz["points_rz"]
        points_xyz = npz["points_xyz"]
        r = npz["r"]
        theta = npz["theta"]
        z = npz["z"]
        dl = npz["dl"]
        dl = dl[()]     # Recover dict from 0d-array
                        # http://stackoverflow.com/questions/8361561/recover-dict-from-0-d-numpy-array
        npz.close()

    else:

        print ""
        print "*** Generating Sampling Points."
        print ""

        # Generate Points
        # We generally don't want r=0
        r, dr = np.linspace(radius, 0, nr, retstep=True, endpoint=False)
        r = r[::-1]; dr = -dr
        theta, dtheta = np.linspace(0, 2 * pi, \
                                    ntheta, retstep=True, endpoint=False)
        z, dz = np.linspace(- thickness/2., thickness / 2., \
                            nz, retstep=True)

        # Reshape
        points_rzt = mkpoints_pymses(r, z, theta)
        points_rz  = mkpoints_pymses(r, z, np.array([0.0]))

        # Generate XYZ
        points_xyz = np.zeros_like(points_rzt)
        points_xyz[:,0] = points_rzt[:,0] * np.cos(points_rzt[:,2])
        points_xyz[:,1] = points_rzt[:,0] * np.sin(points_rzt[:,2])
        points_xyz[:,2] = points_rzt[:,1]
        points_xyz += center

        # Point Spacing
        dl = {"dr": dr, "dtheta": dtheta, "dz": dz}

        # Save
        np.savez(floc, \
                 points_rzt = points_rzt, points_rz = points_rz, \
                 points_xyz = points_xyz, \
                 r = r, theta = theta, z = z, dl = dl)

    return points_rzt, points_rz, points_xyz, r, theta, z, dl

def mkpoints_pymses(x, y, z):
    """Generate Points Array for Pymses."""

    zz, yy, xx = mkvec(z, y, x)
    points = np.zeros((xx.shape[0], 3))

    for ii in range(xx.shape[0]):
        points[ii,:] = np.array([xx[ii], yy[ii], zz[ii]])

    return points