Read an xsil field file and plot at Bloch vectors using mlab and as mollweide plot

mollweide_vis_win.py

#!/usr/bin/env python
from __future__ import with_statement

import sys
import numpy as np
from numpy import array, fromfile, dstack, exp, abs, sin, cos, angle, dot, arccos, pi
from enthought.mayavi import mlab, modules
from lxml import etree, objectify
import ConfigParser
import matplotlib.pyplot as plt
from mpl_toolkits.basemap import Basemap
from matplotlib.patches import Polygon
from scipy.ndimage import maximum_filter, map_coordinates

__rev__ = "$Revision$"


def pointsOnSphere(coords = 'cartesian', radius=1.0):
    x,y,z = np.loadtxt('ico6.txt.gz').T
    if coords == 'cartesian':
        return x*radius,y*radius,z*radius
    elif coords == 'spherical':
        return 180/pi*np.arcsin(z), 180/pi*np.arctan2(y,x)
    else:
        raise 'coord type not implemented'


def interpolate(pt, u, v, w):
    ''' x ranges from 1:xlen but may have missing points because I
        mask them if the prob dens is too low
    '''
    # get interpolated u components
    su = map_coordinates(u, pt, order=1)
    sv = map_coordinates(v, pt, order=1)
    sw = map_coordinates(w, pt, order=1)
    return su, sv, sw


def show_bloch_vectors(x,y,z,u,v,w,norm,N,volume):
    nmask = maximum_filter(norm, size=5) > float(N)/volume
    xm = x[nmask]
    ym = y[nmask]
    zm = z[nmask]
    u = u[nmask]
    v = v[nmask]
    w = w[nmask]

    FACTOR = 1
    #~ OPACITY = 1
    mode = ['2darrow', 'arrow'][0]
    colormap = ['jet', 'Blues', 'winter', 'copper', 'cool', 'blue-red'][0]
    #~ q = mlab.quiver3d(xm, ym, zm, u, v, w, mode=mode, colormap=colormap, opacity=OPACITY, scale_factor=FACTOR, mask_points=FACTOR)
    q = mlab.quiver3d(xm, ym, zm, u, v, w, mode=mode, colormap=colormap, opacity=0.3, scale_factor=0.5, mask_points=FACTOR)
    q.glyph.glyph_source.glyph_position = 'center'


def texture(frame_number, xsil_fname):
    xml = open(xsil_fname)
    xsil = objectify.parse(xml).getroot()
    ulong_len = xsil.XSIL.Array[1].Stream.Metalink.get("UnsignedLong")
    precision = xsil.XSIL.Array[1].Stream.Metalink.get("precision")
    xml.close()
    ulong_type = {'uint32':'<u4','uint64':'<u8'}[ulong_len]
    float_type = {'single':'<f4','double':'<f8'}[precision]

    with open('field%03d.dat'%frame_number,'rb') as f:
        xlen = fromfile(f, ulong_type, 1)[0]
        xs   = fromfile(f, float_type, xlen)
        ylen = fromfile(f, ulong_type, 1)[0]
        ys   = fromfile(f, float_type, ylen)
        zlen = fromfile(f, ulong_type, 1)[0]
        zs   = fromfile(f, float_type, zlen)

        rlen1  = fromfile(f, ulong_type, 1)[0]
        reals1 = fromfile(f, float_type, rlen1)
        ilen1  = fromfile(f, ulong_type, 1)[0]
        imags1 = fromfile(f, float_type, ilen1)
        psi1 = (reals1+1j*imags1).astype(np.complex64).reshape(xlen,ylen,zlen)
        del reals1, imags1

        rlen2  = fromfile(f, ulong_type, 1)[0]
        reals2 = fromfile(f, float_type, rlen2)
        ilen2  = fromfile(f, ulong_type, 1)[0]
        imags2 = fromfile(f, float_type, ilen2)
        psi2 = (reals2+1j*imags2).astype(np.complex64).reshape(xlen,ylen,zlen)
        del reals2, imags2

    print xlen,ylen,zlen,rlen1,ilen1,rlen2,ilen2

    # normalise psi
    norm = np.hypot(abs(psi1), abs(psi2))
    N = norm.sum()
    psi1, psi2 =  (psi1, psi2) / norm

    xi = angle(psi1)                        # find the local phase xi so we can remove it
    psi1, psi2 = exp(-1j*xi) * (psi1, psi2) # now rotate psi0 and psi1 by -xi to make the first coefficient real

    theta = 2*arccos(psi1.real)             # this gives us theta
    phi = angle(psi2/sin(theta/2))          # which gives us phi
    phi[np.isnan(phi)] = 0                  # deal with the poles i.e. where theta=0 or pi

    # Now we've got our angles, get the Bloch sphere arrow components
    u = sin(theta) * cos(phi)
    v = sin(theta) * sin(phi)
    w = cos(theta)

    mlab.figure(bgcolor=(1,1,1))
    x, y, z = np.mgrid[1:xlen:xlen*1j, 1:ylen:ylen*1j, 1:zlen:zlen*1j]
    show_bloch_vectors(x,y,z,u,v,w,norm,N,xlen*ylen*zlen)

    radius_s = xlen/8
    sx,sy,sz = pointsOnSphere(coords = 'cartesian', radius=radius_s)
    sx += xlen/2
    sy += ylen/2
    sz += zlen/2

    sphere = mlab.triangular_mesh(sx, sy, sz, np.arange(len(sx)).reshape(-1,3), representation = 'wireframe')

    # get interpolated vector components at sphere vertices
    su, sv, sw = interpolate(np.vstack((sx, sy, sz)), u, v, w)

    q = mlab.quiver3d(sx,sy,sz,su,sv,sw)
    mlab.show()

    su.shape = (-1,3)
    sv.shape = (-1,3)
    sw.shape = (-1,3)
           
    s_theta = 180/pi*np.arcsin(sw)
    s_phi = 180/pi*np.arctan2(sv,su)

    # Mollweide plot
    # Make Mollweide plot
    m = Basemap(projection='moll', lon_0=0, resolution='c')

    # draw the edge of the map projection region (the projection limb)
    m.drawmapboundary()
    # draw lat/lon grid lines every 30 degrees.
    m.drawmeridians(np.arange(0,360,30), dashes=[10,0])
    m.drawparallels(np.arange(-90,90,30), dashes=[10,0])

    ax = plt.gca()                          # get current axes instance
    for i in range(s_theta.shape[0]):
        pts = np.vstack((s_theta[i], s_phi[i])).T
        phi1, phi2, phi3 = pts[:,1]
        if abs(phi1-phi2)>180 or abs(phi2-phi3)>180 or abs(phi3-phi1)>180:
            continue            # skip triangles that cross the map boundary
        polypts = [m(pt[1], pt[0]) for pt in pts]
        c = tuple(np.random.random(3))
        #~ poly = Polygon(polypts, facecolor='b', edgecolor=c) #edgecolor="None")
        poly = Polygon(polypts, facecolor='b', edgecolor="None")
        ax.add_patch(poly)

    plt.show()


if __name__ == "__main__":
    from optparse import OptionParser

    parser = OptionParser()
    parser.add_option("-f", "--frame", type="int",    dest="frame_number", default=8)
    parser.add_option("-x", "--xsil",  type="string", dest="xsil_fname",   default="field008.xsil")
    (options, args) = parser.parse_args()

    texture(frame_number=options.frame_number,
              xsil_fname=options.xsil_fname)