source: sasmodels/sasmodels/jitter.py @ 4e28511

#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Jitter Explorer
===============

Application to explore orientation angle and angular dispersity.

From the command line::

    # Show docs
    python -m sasmodels.jitter --help

    # Guyou projection jitter, uniform over 20 degree theta and 10 in phi
    python -m sasmodels.jitter --projection=guyou --dist=uniform --jitter=20,10,0

From a jupyter cell::

    import ipyvolume as ipv
    from sasmodels import jitter
    import importlib; importlib.reload(jitter)
    jitter.set_plotter("ipv")

    size = (10, 40, 100)
    view = (20, 0, 0)

    #size = (15, 15, 100)
    #view = (60, 60, 0)

    dview = (0, 0, 0)
    #dview = (5, 5, 0)
    #dview = (15, 180, 0)
    #dview = (180, 15, 0)

    projection = 'equirectangular'
    #projection = 'azimuthal_equidistance'
    #projection = 'guyou'
    #projection = 'sinusoidal'
    #projection = 'azimuthal_equal_area'

    dist = 'uniform'
    #dist = 'gaussian'

    jitter.run(size=size, view=view, jitter=dview, dist=dist, projection=projection)
    #filename = projection+('_theta' if dview[0] == 180 else '_phi' if dview[1] == 180 else '')
    #ipv.savefig(filename+'.png')
"""
from __future__ import division, print_function

import argparse

import numpy as np
from numpy import pi, cos, sin, sqrt, exp, log, degrees, radians, arccos, arctan2

# Too many complaints about variable names from pylint:
#    a, b, c, u, v, x, y, z, dx, dy, dz, px, py, pz, R, Rx, Ry, Rz, ...
# pylint: disable=invalid-name

def draw_beam(axes, view=(0, 0), alpha=0.5, steps=25):
    """
    Draw the beam going from source at (0, 0, 1) to detector at (0, 0, -1)
    """
    #axes.plot([0,0],[0,0],[1,-1])
    #axes.scatter([0]*100,[0]*100,np.linspace(1, -1, 100), alpha=alpha)

    u = np.linspace(0, 2 * pi, steps)
    v = np.linspace(-1, 1, 2)

    r = 0.02
    x = r*np.outer(cos(u), np.ones_like(v))
    y = r*np.outer(sin(u), np.ones_like(v))
    z = 1.3*np.outer(np.ones_like(u), v)

    theta, phi = view
    shape = x.shape
    points = np.matrix([x.flatten(), y.flatten(), z.flatten()])
    points = Rz(phi)*Ry(theta)*points
    x, y, z = [v.reshape(shape) for v in points]
    axes.plot_surface(x, y, z, color='yellow', alpha=alpha)

    # TODO: draw endcaps on beam
    ## Drawing tiny balls on the end will work
    #draw_sphere(axes, radius=0.02, center=(0, 0, 1.3), color='yellow', alpha=alpha)
    #draw_sphere(axes, radius=0.02, center=(0, 0, -1.3), color='yellow', alpha=alpha)
    ## The following does not work
    #triangles = [(0, i+1, i+2) for i in range(steps-2)]
    #x_cap, y_cap = x[:, 0], y[:, 0]
    #for z_cap in z[:, 0], z[:, -1]:
    #    axes.plot_trisurf(x_cap, y_cap, z_cap, triangles,
    #                      color='yellow', alpha=alpha)


def draw_ellipsoid(axes, size, view, jitter, steps=25, alpha=1):
    """Draw an ellipsoid."""
    a, b, c = size
    u = np.linspace(0, 2 * pi, steps)
    v = np.linspace(0, pi, steps)
    x = a*np.outer(cos(u), sin(v))
    y = b*np.outer(sin(u), sin(v))
    z = c*np.outer(np.ones_like(u), cos(v))
    x, y, z = transform_xyz(view, jitter, x, y, z)

    axes.plot_surface(x, y, z, color='w', alpha=alpha)

    draw_labels(axes, view, jitter, [
        ('c+', [+0, +0, +c], [+1, +0, +0]),
        ('c-', [+0, +0, -c], [+0, +0, -1]),
        ('a+', [+a, +0, +0], [+0, +0, +1]),
        ('a-', [-a, +0, +0], [+0, +0, -1]),
        ('b+', [+0, +b, +0], [-1, +0, +0]),
        ('b-', [+0, -b, +0], [-1, +0, +0]),
    ])

def draw_sc(axes, size, view, jitter, steps=None, alpha=1):
    """Draw points for simple cubic paracrystal"""
    # pylint: disable=unused-argument
    atoms = _build_sc()
    _draw_crystal(axes, size, view, jitter, atoms=atoms)

def draw_fcc(axes, size, view, jitter, steps=None, alpha=1):
    """Draw points for face-centered cubic paracrystal"""
    # pylint: disable=unused-argument
    # Build the simple cubic crystal
    atoms = _build_sc()
    # Define the centers for each face
    # x planes at -1, 0, 1 have four centers per plane, at +/- 0.5 in y and z
    x, y, z = (
        [-1]*4 + [0]*4 + [1]*4,
        ([-0.5]*2 + [0.5]*2)*3,
        [-0.5, 0.5]*12,
    )
    # y and z planes can be generated by substituting x for y and z respectively
    atoms.extend(zip(x+y+z, y+z+x, z+x+y))
    _draw_crystal(axes, size, view, jitter, atoms=atoms)

def draw_bcc(axes, size, view, jitter, steps=None, alpha=1):
    """Draw points for body-centered cubic paracrystal"""
    # pylint: disable=unused-argument
    # Build the simple cubic crystal
    atoms = _build_sc()
    # Define the centers for each octant
    # x plane at +/- 0.5 have four centers per plane at +/- 0.5 in y and z
    x, y, z = (
        [-0.5]*4 + [0.5]*4,
        ([-0.5]*2 + [0.5]*2)*2,
        [-0.5, 0.5]*8,
    )
    atoms.extend(zip(x, y, z))
    _draw_crystal(axes, size, view, jitter, atoms=atoms)

def _draw_crystal(axes, size, view, jitter, atoms=None):
    atoms, size = np.asarray(atoms, 'd').T, np.asarray(size, 'd')
    x, y, z = atoms*size[:, None]
    x, y, z = transform_xyz(view, jitter, x, y, z)
    axes.scatter([x[0]], [y[0]], [z[0]], c='yellow', marker='o')
    axes.scatter(x[1:], y[1:], z[1:], c='r', marker='o')

def _build_sc():
    # three planes of 9 dots for x at -1, 0 and 1
    x, y, z = (
        [-1]*9 + [0]*9 + [1]*9,
        ([-1]*3 + [0]*3 + [1]*3)*3,
        [-1, 0, 1]*9,
    )
    atoms = list(zip(x, y, z))
    #print(list(enumerate(atoms)))
    # Pull the dot at (0, 0, 1) to the front of the list
    # It will be highlighted in the view
    index = 14
    highlight = atoms[index]
    del atoms[index]
    atoms.insert(0, highlight)
    return atoms

def draw_box(axes, size, view):
    """Draw a wireframe box at a particular view."""
    a, b, c = size
    x = a*np.array([+1, -1, +1, -1, +1, -1, +1, -1])
    y = b*np.array([+1, +1, -1, -1, +1, +1, -1, -1])
    z = c*np.array([+1, +1, +1, +1, -1, -1, -1, -1])
    x, y, z = transform_xyz(view, None, x, y, z)
    def _draw(i, j):
        axes.plot([x[i], x[j]], [y[i], y[j]], [z[i], z[j]], color='black')
    _draw(0, 1)
    _draw(0, 2)
    _draw(0, 3)
    _draw(7, 4)
    _draw(7, 5)
    _draw(7, 6)

def draw_parallelepiped(axes, size, view, jitter, steps=None,
                        color=(0.6, 1.0, 0.6), alpha=1):
    """Draw a parallelepiped surface, with view and jitter."""
    # pylint: disable=unused-argument
    a, b, c = size
    x = a*np.array([+1, -1, +1, -1, +1, -1, +1, -1])
    y = b*np.array([+1, +1, -1, -1, +1, +1, -1, -1])
    z = c*np.array([+1, +1, +1, +1, -1, -1, -1, -1])
    tri = np.array([
        # counter clockwise triangles
        # z: up/down, x: right/left, y: front/back
        [0, 1, 2], [3, 2, 1], # top face
        [6, 5, 4], [5, 6, 7], # bottom face
        [0, 2, 6], [6, 4, 0], # right face
        [1, 5, 7], [7, 3, 1], # left face
        [2, 3, 6], [7, 6, 3], # front face
        [4, 1, 0], [5, 1, 4], # back face
    ])

    x, y, z = transform_xyz(view, jitter, x, y, z)
    axes.plot_trisurf(x, y, triangles=tri, Z=z, color=color, alpha=alpha,
                      linewidth=0)

    # Colour the c+ face of the box.
    # Since I can't control face color, instead draw a thin box situated just
    # in front of the "c+" face.  Use the c face so that rotations about psi
    # rotate that face.
    if 0: # pylint: disable=using-constant-test
        color = (1, 0.6, 0.6)  # pink
        x = a*np.array([+1, -1, +1, -1, +1, -1, +1, -1])
        y = b*np.array([+1, +1, -1, -1, +1, +1, -1, -1])
        z = c*np.array([+1, +1, +1, +1, -1, -1, -1, -1])
        x, y, z = transform_xyz(view, jitter, x, y, abs(z)+0.001)
        axes.plot_trisurf(x, y, triangles=tri, Z=z, color=color, alpha=alpha)

    draw_labels(axes, view, jitter, [
        ('c+', [+0, +0, +c], [+1, +0, +0]),
        ('c-', [+0, +0, -c], [+0, +0, -1]),
        ('a+', [+a, +0, +0], [+0, +0, +1]),
        ('a-', [-a, +0, +0], [+0, +0, -1]),
        ('b+', [+0, +b, +0], [-1, +0, +0]),
        ('b-', [+0, -b, +0], [-1, +0, +0]),
    ])

def draw_sphere(axes, radius=1.0, steps=25,
                center=(0, 0, 0), color='w', alpha=1.):
    """Draw a sphere"""
    u = np.linspace(0, 2 * pi, steps)
    v = np.linspace(0, pi, steps)

    x = radius * np.outer(cos(u), sin(v)) + center[0]
    y = radius * np.outer(sin(u), sin(v)) + center[1]
    z = radius * np.outer(np.ones(np.size(u)), cos(v)) + center[2]
    axes.plot_surface(x, y, z, color=color, alpha=alpha)
    #axes.plot_wireframe(x, y, z)

def draw_axes(axes, origin=(-1, -1, -1), length=(2, 2, 2)):
    """Draw wireframe axes lines, with given origin and length"""
    x, y, z = origin
    dx, dy, dz = length
    axes.plot([x, x+dx], [y, y], [z, z], color='black')
    axes.plot([x, x], [y, y+dy], [z, z], color='black')
    axes.plot([x, x], [y, y], [z, z+dz], color='black')

def draw_person_on_sphere(axes, view, height=0.5, radius=1.0):
    """
    Draw a person on the surface of a sphere.

    *view* indicates (latitude, longitude, orientation)
    """
    limb_offset = height * 0.05
    head_radius = height * 0.10
    head_height = height - head_radius
    neck_length = head_radius * 0.50
    shoulder_height = height - 2*head_radius - neck_length
    torso_length = shoulder_height * 0.55
    torso_radius = torso_length * 0.30
    leg_length = shoulder_height - torso_length
    arm_length = torso_length * 0.90

    def _draw_part(y, z):
        x = np.zeros_like(y)
        xp, yp, zp = transform_xyz(view, None, x, y, z + radius)
        axes.plot(xp, yp, zp, color='k')

    # circle for head
    u = np.linspace(0, 2 * pi, 40)
    y = head_radius * cos(u)
    z = head_radius * sin(u) + head_height
    _draw_part(y, z)

    # rectangle for body
    y = np.array([-torso_radius, torso_radius, torso_radius, -torso_radius, -torso_radius])
    z = np.array([0., 0, torso_length, torso_length, 0]) + leg_length
    _draw_part(y, z)

    # arms
    y = np.array([-torso_radius - limb_offset, -torso_radius - limb_offset, -torso_radius])
    z = np.array([shoulder_height - arm_length, shoulder_height, shoulder_height])
    _draw_part(y, z)
    _draw_part(-y, z)  # pylint: disable=invalid-unary-operand-type

    # legs
    y = np.array([-torso_radius + limb_offset, -torso_radius + limb_offset])
    z = np.array([0, leg_length])
    _draw_part(y, z)
    _draw_part(-y, z)  # pylint: disable=invalid-unary-operand-type

    limits = [-radius-height, radius+height]
    axes.set_xlim(limits)
    axes.set_ylim(limits)
    axes.set_zlim(limits)
    axes.set_axis_off()

def draw_jitter(axes, view, jitter, dist='gaussian',
                size=(0.1, 0.4, 1.0),
                draw_shape=draw_parallelepiped,
                projection='equirectangular',
                alpha=0.8,
                views=None):
    """
    Represent jitter as a set of shapes at different orientations.
    """
    project, project_weight = get_projection(projection)

    # set max diagonal to 0.95
    scale = 0.95/sqrt(sum(v**2 for v in size))
    size = tuple(scale*v for v in size)

    dtheta, dphi, dpsi = jitter
    base = {'gaussian':3, 'rectangle':sqrt(3), 'uniform':1}[dist]
    def _steps(delta):
        if views is None:
            n = max(3, min(25, 2*int(base*delta/5)))
        else:
            n = views
        return base*delta*np.linspace(-1, 1, n) if delta > 0 else [0.]
    for theta in _steps(dtheta):
        for phi in _steps(dphi):
            for psi in _steps(dpsi):
                w = project_weight(theta, phi, 1.0, 1.0)
                if w > 0:
                    dview = project(theta, phi, psi)
                    draw_shape(axes, size, view, dview, alpha=alpha)
    for v in 'xyz':
        a, b, c = size
        lim = sqrt(a**2 + b**2 + c**2)
        getattr(axes, 'set_'+v+'lim')([-lim, lim])
        #getattr(axes, v+'axis').label.set_text(v)

PROJECTIONS = [
    # in order of PROJECTION number; do not change without updating the
    # constants in kernel_iq.c
    'equirectangular', 'sinusoidal', 'guyou', 'azimuthal_equidistance',
    'azimuthal_equal_area',
]
def get_projection(projection):

    """
    jitter projections
    <https://en.wikipedia.org/wiki/List_of_map_projections>

    equirectangular (standard latitude-longitude mesh)
        <https://en.wikipedia.org/wiki/Equirectangular_projection>
        Allows free movement in phi (around the equator), but theta is
        limited to +/- 90, and points are cos-weighted. Jitter in phi is
        uniform in weight along a line of latitude.  With small theta and
        phi ranging over +/- 180 this forms a wobbling disk.  With small
        phi and theta ranging over +/- 90 this forms a wedge like a slice
        of an orange.
    azimuthal_equidistance (Postel)
        <https://en.wikipedia.org/wiki/Azimuthal_equidistant_projection>
        Preserves distance from center, and so is an excellent map for
        representing a bivariate gaussian on the surface.  Theta and phi
        operate identically, cutting wegdes from the antipode of the viewing
        angle.  This unfortunately does not allow free movement in either
        theta or phi since the orthogonal wobble decreases to 0 as the body
        rotates through 180 degrees.
    sinusoidal (Sanson-Flamsteed, Mercator equal-area)
        <https://en.wikipedia.org/wiki/Sinusoidal_projection>
        Preserves arc length with latitude, giving bad behaviour at
        theta near +/- 90.  Theta and phi operate somewhat differently,
        so a system with a-b-c dtheta-dphi-dpsi will not give the same
        value as one with b-a-c dphi-dtheta-dpsi, as would be the case
        for azimuthal equidistance.  Free movement using theta or phi
        uniform over +/- 180 will work, but not as well as equirectangular
        phi, with theta being slightly worse.  Computationally it is much
        cheaper for wide theta-phi meshes since it excludes points which
        lie outside the sinusoid near theta +/- 90 rather than packing
        them close together as in equirectangle.  Note that the poles
        will be slightly overweighted for theta > 90 with the circle
        from theta at 90+dt winding backwards around the pole, overlapping
        the circle from theta at 90-dt.
    Guyou (hemisphere-in-a-square) **not weighted**
        <https://en.wikipedia.org/wiki/Guyou_hemisphere-in-a-square_projection>
        With tiling, allows rotation in phi or theta through +/- 180, with
        uniform spacing.  Both theta and phi allow free rotation, with wobble
        in the orthogonal direction reasonably well behaved (though not as
        good as equirectangular phi). The forward/reverse transformations
        relies on elliptic integrals that are somewhat expensive, so the
        behaviour has to be very good to justify the cost and complexity.
        The weighting function for each point has not yet been computed.
        Note: run the module *guyou.py* directly and it will show the forward
        and reverse mappings.
    azimuthal_equal_area  **incomplete**
        <https://en.wikipedia.org/wiki/Lambert_azimuthal_equal-area_projection>
        Preserves the relative density of the surface patches.  Not that
        useful and not completely implemented
    Gauss-Kreuger **not implemented**
        <https://en.wikipedia.org/wiki/Transverse_Mercator_projection#Ellipsoidal_transverse_Mercator>
        Should allow free movement in theta, but phi is distorted.
    """
    # pylint: disable=unused-argument
    # TODO: try Kent distribution instead of a gaussian warped by projection

    if projection == 'equirectangular':  #define PROJECTION 1
        def _project(theta_i, phi_j, psi):
            latitude, longitude = theta_i, phi_j
            return latitude, longitude, psi, 'xyz'
            #return Rx(phi_j)*Ry(theta_i)
        def _weight(theta_i, phi_j, w_i, w_j):
            return w_i*w_j*abs(cos(radians(theta_i)))
    elif projection == 'sinusoidal':  #define PROJECTION 2
        def _project(theta_i, phi_j, psi):
            latitude = theta_i
            scale = cos(radians(latitude))
            longitude = phi_j/scale if abs(phi_j) < abs(scale)*180 else 0
            #print("(%+7.2f, %+7.2f) => (%+7.2f, %+7.2f)"%(theta_i, phi_j, latitude, longitude))
            return latitude, longitude, psi, 'xyz'
            #return Rx(longitude)*Ry(latitude)
        def _project(theta_i, phi_j, w_i, w_j):
            latitude = theta_i
            scale = cos(radians(latitude))
            active = 1 if abs(phi_j) < abs(scale)*180 else 0
            return active*w_i*w_j
    elif projection == 'guyou':  #define PROJECTION 3  (eventually?)
        def _project(theta_i, phi_j, psi):
            from .guyou import guyou_invert
            #latitude, longitude = guyou_invert([theta_i], [phi_j])
            longitude, latitude = guyou_invert([phi_j], [theta_i])
            return latitude, longitude, psi, 'xyz'
            #return Rx(longitude[0])*Ry(latitude[0])
        def _weight(theta_i, phi_j, w_i, w_j):
            return w_i*w_j
    elif projection == 'azimuthal_equidistance':
        # Note that calculates angles for Rz Ry rather than Rx Ry
        def _project(theta_i, phi_j, psi):
            latitude = sqrt(theta_i**2 + phi_j**2)
            longitude = degrees(arctan2(phi_j, theta_i))
            #print("(%+7.2f, %+7.2f) => (%+7.2f, %+7.2f)"%(theta_i, phi_j, latitude, longitude))
            return latitude, longitude, psi-longitude, 'zyz'
            #R = Rz(longitude)*Ry(latitude)*Rz(psi)
            #return R_to_xyz(R)
            #return Rz(longitude)*Ry(latitude)
        def _weight(theta_i, phi_j, w_i, w_j):
            # Weighting for each point comes from the integral:
            #     \int\int I(q, lat, log) sin(lat) dlat dlog
            # We are doing a conformal mapping from disk to sphere, so we need
            # a change of variables g(theta, phi) -> (lat, long):
            #     lat, long = sqrt(theta^2 + phi^2), arctan(phi/theta)
            # giving:
            #     dtheta dphi = det(J) dlat dlong
            # where J is the jacobian from the partials of g. Using
            #     R = sqrt(theta^2 + phi^2),
            # then
            #     J = [[x/R, Y/R], -y/R^2, x/R^2]]
            # and
            #     det(J) = 1/R
            # with the final integral being:
            #    \int\int I(q, theta, phi) sin(R)/R dtheta dphi
            #
            # This does approximately the right thing, decreasing the weight
            # of each point as you go farther out on the disk, but it hasn't
            # yet been checked against the 1D integral results. Prior
            # to declaring this "good enough" and checking that integrals
            # work in practice, we will examine alternative mappings.
            #
            # The issue is that the mapping does not support the case of free
            # rotation about a single axis correctly, with a small deviation
            # in the orthogonal axis independent of the first axis.  Like the
            # usual polar coordiates integration, the integrated sections
            # form wedges, though at least in this case the wedge cuts through
            # the entire sphere, and treats theta and phi identically.
            latitude = sqrt(theta_i**2 + phi_j**2)
            weight = sin(radians(latitude))/latitude if latitude != 0 else 1
            return weight*w_i*w_j if latitude < 180 else 0
    elif projection == 'azimuthal_equal_area':
        # Note that calculates angles for Rz Ry rather than Rx Ry
        def _project(theta_i, phi_j, psi):
            radius = min(1, sqrt(theta_i**2 + phi_j**2)/180)
            latitude = 180-degrees(2*arccos(radius))
            longitude = degrees(arctan2(phi_j, theta_i))
            #print("(%+7.2f, %+7.2f) => (%+7.2f, %+7.2f)"%(theta_i, phi_j, latitude, longitude))
            return latitude, longitude, psi, 'zyz'
            #R = Rz(longitude)*Ry(latitude)*Rz(psi)
            #return R_to_xyz(R)
            #return Rz(longitude)*Ry(latitude)
        def _weight(theta_i, phi_j, w_i, w_j):
            latitude = sqrt(theta_i**2 + phi_j**2)
            weight = sin(radians(latitude))/latitude if latitude != 0 else 1
            return weight*w_i*w_j if latitude < 180 else 0
    else:
        raise ValueError("unknown projection %r"%projection)

    return _project, _weight

def R_to_xyz(R):
    """
    Return phi, theta, psi Tait-Bryan angles corresponding to the given rotation matrix.

    Extracting Euler Angles from a Rotation Matrix
    Mike Day, Insomniac Games
    https://d3cw3dd2w32x2b.cloudfront.net/wp-content/uploads/2012/07/euler-angles1.pdf
    Based on: Shoemake’s "Euler Angle Conversion", Graphics Gems IV, pp.  222-229
    """
    phi = arctan2(R[1, 2], R[2, 2])
    theta = arctan2(-R[0, 2], sqrt(R[0, 0]**2 + R[0, 1]**2))
    psi = arctan2(R[0, 1], R[0, 0])
    return degrees(phi), degrees(theta), degrees(psi)

def draw_mesh(axes, view, jitter, radius=1.2, n=11, dist='gaussian',
              projection='equirectangular'):
    """
    Draw the dispersion mesh showing the theta-phi orientations at which
    the model will be evaluated.
    """

    _project, _weight = get_projection(projection)
    def _rotate(theta, phi, z):
        dview = _project(theta, phi, 0.)
        if dview[3] == 'zyz':
            return Rz(dview[1])*Ry(dview[0])*z
        else:  # dview[3] == 'xyz':
            return Rx(dview[1])*Ry(dview[0])*z


    dist_x = np.linspace(-1, 1, n)
    weights = np.ones_like(dist_x)
    if dist == 'gaussian':
        dist_x *= 3
        weights = exp(-0.5*dist_x**2)
    elif dist == 'rectangle':
        # Note: uses sasmodels ridiculous definition of rectangle width
        dist_x *= sqrt(3)
    elif dist == 'uniform':
        pass
    else:
        raise ValueError("expected dist to be gaussian, rectangle or uniform")

    # mesh in theta, phi formed by rotating z
    dtheta, dphi, dpsi = jitter  # pylint: disable=unused-variable
    z = np.matrix([[0], [0], [radius]])
    points = np.hstack([_rotate(theta_i, phi_j, z)
                        for theta_i in dtheta*dist_x
                        for phi_j in dphi*dist_x])
    dist_w = np.array([_weight(theta_i, phi_j, w_i, w_j)
                       for w_i, theta_i in zip(weights, dtheta*dist_x)
                       for w_j, phi_j in zip(weights, dphi*dist_x)])
    #print(max(dist_w), min(dist_w), min(dist_w[dist_w > 0]))
    points = points[:, dist_w > 0]
    dist_w = dist_w[dist_w > 0]
    dist_w /= max(dist_w)

    # rotate relative to beam
    points = orient_relative_to_beam(view, points)

    x, y, z = [np.array(v).flatten() for v in points]
    #plt.figure(2); plt.clf(); plt.hist(z, bins=np.linspace(-1, 1, 51))
    axes.scatter(x, y, z, c=dist_w, marker='o', vmin=0., vmax=1.)

def draw_labels(axes, view, jitter, text):
    """
    Draw text at a particular location.
    """
    labels, locations, orientations = zip(*text)
    px, py, pz = zip(*locations)
    dx, dy, dz = zip(*orientations)

    px, py, pz = transform_xyz(view, jitter, px, py, pz)
    dx, dy, dz = transform_xyz(view, jitter, dx, dy, dz)

    # TODO: zdir for labels is broken, and labels aren't appearing.
    for label, p, zdir in zip(labels, zip(px, py, pz), zip(dx, dy, dz)):
        zdir = np.asarray(zdir).flatten()
        axes.text(p[0], p[1], p[2], label, zdir=zdir)

# Definition of rotation matrices comes from wikipedia:
#    https://en.wikipedia.org/wiki/Rotation_matrix#Basic_rotations
def Rx(angle):
    """Construct a matrix to rotate points about *x* by *angle* degrees."""
    angle = radians(angle)
    rot = [[1, 0, 0],
           [0, +cos(angle), -sin(angle)],
           [0, +sin(angle), +cos(angle)]]
    return np.matrix(rot)

def Ry(angle):
    """Construct a matrix to rotate points about *y* by *angle* degrees."""
    angle = radians(angle)
    rot = [[+cos(angle), 0, +sin(angle)],
           [0, 1, 0],
           [-sin(angle), 0, +cos(angle)]]
    return np.matrix(rot)

def Rz(angle):
    """Construct a matrix to rotate points about *z* by *angle* degrees."""
    angle = radians(angle)
    rot = [[+cos(angle), -sin(angle), 0],
           [+sin(angle), +cos(angle), 0],
           [0, 0, 1]]
    return np.matrix(rot)

def transform_xyz(view, jitter, x, y, z):
    """
    Send a set of (x,y,z) points through the jitter and view transforms.
    """
    x, y, z = [np.asarray(v) for v in (x, y, z)]
    shape = x.shape
    points = np.matrix([x.flatten(), y.flatten(), z.flatten()])
    points = apply_jitter(jitter, points)
    points = orient_relative_to_beam(view, points)
    x, y, z = [np.array(v).reshape(shape) for v in points]
    return x, y, z

def apply_jitter(jitter, points):
    """
    Apply the jitter transform to a set of points.

    Points are stored in a 3 x n numpy matrix, not a numpy array or tuple.
    """
    if jitter is None:
        return points
    # Hack to deal with the fact that azimuthal_equidistance uses euler angles
    if len(jitter) == 4:
        dtheta, dphi, dpsi, _ = jitter
        points = Rz(dphi)*Ry(dtheta)*Rz(dpsi)*points
    else:
        dtheta, dphi, dpsi = jitter
        points = Rx(dphi)*Ry(dtheta)*Rz(dpsi)*points
    return points

def orient_relative_to_beam(view, points):
    """
    Apply the view transform to a set of points.

    Points are stored in a 3 x n numpy matrix, not a numpy array or tuple.
    """
    theta, phi, psi = view
    points = Rz(phi)*Ry(theta)*Rz(psi)*points # viewing angle
    #points = Rz(psi)*Ry(pi/2-theta)*Rz(phi)*points # 1-D integration angles
    #points = Rx(phi)*Ry(theta)*Rz(psi)*points  # angular dispersion angle
    return points

def orient_relative_to_beam_quaternion(view, points):
    """
    Apply the view transform to a set of points.

    Points are stored in a 3 x n numpy matrix, not a numpy array or tuple.

    This variant uses quaternions rather than rotation matrices for the
    computation.  It works but it is not used because it doesn't solve
    any problems.  The challenge of mapping theta/phi/psi to SO(3) does
    not disappear by calculating the transform differently.
    """
    theta, phi, psi = view
    x, y, z = [1, 0, 0], [0, 1, 0], [0, 0, 1]
    q = Quaternion(1, [0, 0, 0])
    ## Compose a rotation about the three axes by rotating
    ## the unit vectors before applying the rotation.
    #q = Quaternion.from_angle_axis(theta, q.rot(x)) * q
    #q = Quaternion.from_angle_axis(phi, q.rot(y)) * q
    #q = Quaternion.from_angle_axis(psi, q.rot(z)) * q
    ## The above turns out to be equivalent to reversing
    ## the order of application, so ignore it and use below.
    q = q * Quaternion.from_angle_axis(theta, x)
    q = q * Quaternion.from_angle_axis(phi, y)
    q = q * Quaternion.from_angle_axis(psi, z)
    ## Reverse the order by post-multiply rather than pre-multiply
    #q = Quaternion.from_angle_axis(theta, x) * q
    #q = Quaternion.from_angle_axis(phi, y) * q
    #q = Quaternion.from_angle_axis(psi, z) * q
    #print("axes psi", q.rot(np.matrix([x, y, z]).T))
    return q.rot(points)
#orient_relative_to_beam = orient_relative_to_beam_quaternion

# === Quaterion class definition === BEGIN
# Simple stand-alone quaternion class

# Note: this code works but isn't unused since quaternions didn't solve the
# representation problem.  Leave it here in case we want to revisit this later.

#import numpy as np
class Quaternion(object):
    r"""
    Quaternion(w, r) = w + ir[0] + jr[1] + kr[2]

    Quaternion.from_angle_axis(theta, r) for a rotation of angle theta about
    an axis oriented toward the direction r.  This defines a unit quaternion,
    normalizing $r$ to the unit vector $\hat r$, and setting quaternion
    $Q = \cos \theta + \sin \theta \hat r$

    Quaternion objects can be multiplied, which applies a rotation about the
    given axis, allowing composition of rotations without risk of gimbal lock.
    The resulting quaternion is applied to a set of points using *Q.rot(v)*.
    """
    def __init__(self, w, r):
        self.w = w
        self.r = np.asarray(r, dtype='d')

    @staticmethod
    def from_angle_axis(theta, r):
        """Build quaternion as rotation theta about axis r"""
        theta = np.radians(theta)/2
        r = np.asarray(r)
        w = np.cos(theta)
        r = np.sin(theta)*r/np.dot(r, r)
        return Quaternion(w, r)

    def __mul__(self, other):
        """Multiply quaterions"""
        if isinstance(other, Quaternion):
            w = self.w*other.w - np.dot(self.r, other.r)
            r = self.w*other.r + other.w*self.r + np.cross(self.r, other.r)
            return Quaternion(w, r)
        raise NotImplementedError("Quaternion * non-quaternion not implemented")

    def rot(self, v):
        """Transform point *v* by quaternion"""
        v = np.asarray(v).T
        use_transpose = (v.shape[-1] != 3)
        if use_transpose:
            v = v.T
        v = v + np.cross(2*self.r, np.cross(self.r, v) + self.w*v)
        #v = v + 2*self.w*np.cross(self.r, v) + np.cross(2*self.r, np.cross(self.r, v))
        if use_transpose:
            v = v.T
        return v.T

    def conj(self):
        """Conjugate quaternion"""
        return Quaternion(self.w, -self.r)

    def inv(self):
        """Inverse quaternion"""
        return self.conj()/self.norm()**2

    def norm(self):
        """Quaternion length"""
        return np.sqrt(self.w**2 + np.sum(self.r**2))

    def __str__(self):
        return "%g%+gi%+gj%+gk"%(self.w, self.r[0], self.r[1], self.r[2])

def test_qrot():
    """Quaternion checks"""
    # Define rotation of 60 degrees around an axis in y-z that is 60 degrees
    # from y.  The rotation axis is determined by rotating the point [0, 1, 0]
    # about x.
    ax = Quaternion.from_angle_axis(60, [1, 0, 0]).rot([0, 1, 0])
    q = Quaternion.from_angle_axis(60, ax)
    # Set the point to be rotated, and its expected rotated position.
    p = [1, -1, 2]
    target = [(10+4*np.sqrt(3))/8, (1+2*np.sqrt(3))/8, (14-3*np.sqrt(3))/8]
    #print(q, q.rot(p) - target)
    assert max(abs(q.rot(p) - target)) < 1e-14
#test_qrot()
#import sys; sys.exit()
# === Quaterion class definition === END

# translate between number of dimension of dispersity and the number of
# points along each dimension.
PD_N_TABLE = {
    (0, 0, 0): (0, 0, 0),     # 0
    (1, 0, 0): (100, 0, 0),   # 100
    (0, 1, 0): (0, 100, 0),
    (0, 0, 1): (0, 0, 100),
    (1, 1, 0): (30, 30, 0),   # 900
    (1, 0, 1): (30, 0, 30),
    (0, 1, 1): (0, 30, 30),
    (1, 1, 1): (15, 15, 15),  # 3375
}

def clipped_range(data, portion=1.0, mode='central'):
    """
    Determine range from data.

    If *portion* is 1, use full range, otherwise use the center of the range
    or the top of the range, depending on whether *mode* is 'central' or 'top'.
    """
    if portion < 1.0:
        if mode == 'central':
            data = np.sort(data.flatten())
            offset = int(portion*len(data)/2 + 0.5)
            return data[offset], data[-offset]
        if mode == 'top':
            data = np.sort(data.flatten())
            offset = int(portion*len(data) + 0.5)
            return data[offset], data[-1]
    # Default: full range
    return data.min(), data.max()

def draw_scattering(calculator, axes, view, jitter, dist='gaussian'):
    """
    Plot the scattering for the particular view.

    *calculator* is returned from :func:`build_model`.  *axes* are the 3D axes
    on which the data will be plotted.  *view* and *jitter* are the current
    orientation and orientation dispersity.  *dist* is one of the sasmodels
    weight distributions.
    """
    if dist == 'uniform':  # uniform is not yet in this branch
        dist, scale = 'rectangle', 1/sqrt(3)
    else:
        scale = 1

    # add the orientation parameters to the model parameters
    theta, phi, psi = view
    theta_pd, phi_pd, psi_pd = [scale*v for v in jitter]
    theta_pd_n, phi_pd_n, psi_pd_n = PD_N_TABLE[(theta_pd > 0, phi_pd > 0, psi_pd > 0)]
    ## increase pd_n for testing jitter integration rather than simple viz
    #theta_pd_n, phi_pd_n, psi_pd_n = [5*v for v in (theta_pd_n, phi_pd_n, psi_pd_n)]

    pars = dict(
        theta=theta, theta_pd=theta_pd, theta_pd_type=dist, theta_pd_n=theta_pd_n,
        phi=phi, phi_pd=phi_pd, phi_pd_type=dist, phi_pd_n=phi_pd_n,
        psi=psi, psi_pd=psi_pd, psi_pd_type=dist, psi_pd_n=psi_pd_n,
    )
    pars.update(calculator.pars)

    # compute the pattern
    qx, qy = calculator.qxqy
    Iqxy = calculator(**pars).reshape(len(qx), len(qy))

    # scale it and draw it
    Iqxy = log(Iqxy)
    if calculator.limits:
        # use limits from orientation (0,0,0)
        vmin, vmax = calculator.limits
    else:
        vmax = Iqxy.max()
        vmin = vmax*10**-7
        #vmin, vmax = clipped_range(Iqxy, portion=portion, mode='top')
    #vmin, vmax = Iqxy.min(), Iqxy.max()
    #print("range",(vmin,vmax))
    #qx, qy = np.meshgrid(qx, qy)
    if 0:  # pylint: disable=using-constant-test
        level = np.asarray(255*(Iqxy - vmin)/(vmax - vmin), 'i')
        level[level < 0] = 0
        from matplotlib import pylab as plt
        colors = plt.get_cmap()(level)
        #from matplotlib import cm
        #colors = cm.coolwarm(level)
        #colors = cm.gist_yarg(level)
        #colors = cm.Wistia(level)
        colors[level <= 0, 3] = 0.  # set floor to transparent
        x, y = np.meshgrid(qx/qx.max(), qy/qy.max())
        axes.plot_surface(x, y, -1.1*np.ones_like(x), facecolors=colors)
    elif 1:  # pylint: disable=using-constant-test
        axes.contourf(qx/qx.max(), qy/qy.max(), Iqxy, zdir='z', offset=-1.1,
                      levels=np.linspace(vmin, vmax, 24))
    else:
        axes.pcolormesh(qx, qy, Iqxy)

def build_model(model_name, n=150, qmax=0.5, **pars):
    """
    Build a calculator for the given shape.

    *model_name* is any sasmodels model.  *n* and *qmax* define an n x n mesh
    on which to evaluate the model.  The remaining parameters are stored in
    the returned calculator as *calculator.pars*.  They are used by
    :func:`draw_scattering` to set the non-orientation parameters in the
    calculation.

    Returns a *calculator* function which takes a dictionary or parameters and
    produces Iqxy.  The Iqxy value needs to be reshaped to an n x n matrix
    for plotting.  See the :class:`sasmodels.direct_model.DirectModel` class
    for details.
    """
    from sasmodels.core import load_model_info, build_model as build_sasmodel
    from sasmodels.data import empty_data2D
    from sasmodels.direct_model import DirectModel

    model_info = load_model_info(model_name)
    model = build_sasmodel(model_info) #, dtype='double!')
    q = np.linspace(-qmax, qmax, n)
    data = empty_data2D(q, q)
    calculator = DirectModel(data, model)

    # Remember the data axes so we can plot the results
    calculator.qxqy = (q, q)

    # stuff the values for non-orientation parameters into the calculator
    calculator.pars = pars.copy()
    calculator.pars.setdefault('backgound', 1e-3)

    # fix the data limits so that we can see if the pattern fades
    # under rotation or angular dispersion
    Iqxy = calculator(theta=0, phi=0, psi=0, **calculator.pars)
    Iqxy = log(Iqxy)
    vmin, vmax = clipped_range(Iqxy, 0.95, mode='top')
    calculator.limits = vmin, vmax+1

    return calculator

def select_calculator(model_name, n=150, size=(10, 40, 100)):
    """
    Create a model calculator for the given shape.

    *model_name* is one of sphere, cylinder, ellipsoid, triaxial_ellipsoid,
    parallelepiped or bcc_paracrystal. *n* is the number of points to use
    in the q range.  *qmax* is chosen based on model parameters for the
    given model to show something intersting.

    Returns *calculator* and tuple *size* (a,b,c) giving minor and major
    equitorial axes and polar axis respectively.  See :func:`build_model`
    for details on the returned calculator.
    """
    a, b, c = size
    d_factor = 0.06  # for paracrystal models
    if model_name == 'sphere':
        calculator = build_model('sphere', n=n, radius=c)
        a = b = c
    elif model_name == 'sc_paracrystal':
        a = b = c
        dnn = c
        radius = 0.5*c
        calculator = build_model('sc_paracrystal', n=n, dnn=dnn,
                                 d_factor=d_factor, radius=(1-d_factor)*radius,
                                 background=0)
    elif model_name == 'fcc_paracrystal':
        a = b = c
        # nearest neigbour distance dnn should be 2 radius, but I think the
        # model uses lattice spacing rather than dnn in its calculations
        dnn = 0.5*c
        radius = sqrt(2)/4 * c
        calculator = build_model('fcc_paracrystal', n=n, dnn=dnn,
                                 d_factor=d_factor, radius=(1-d_factor)*radius,
                                 background=0)
    elif model_name == 'bcc_paracrystal':
        a = b = c
        # nearest neigbour distance dnn should be 2 radius, but I think the
        # model uses lattice spacing rather than dnn in its calculations
        dnn = 0.5*c
        radius = sqrt(3)/2 * c
        calculator = build_model('bcc_paracrystal', n=n, dnn=dnn,
                                 d_factor=d_factor, radius=(1-d_factor)*radius,
                                 background=0)
    elif model_name == 'cylinder':
        calculator = build_model('cylinder', n=n, qmax=0.3, radius=b, length=c)
        a = b
    elif model_name == 'ellipsoid':
        calculator = build_model('ellipsoid', n=n, qmax=1.0,
                                 radius_polar=c, radius_equatorial=b)
        a = b
    elif model_name == 'triaxial_ellipsoid':
        calculator = build_model('triaxial_ellipsoid', n=n, qmax=0.5,
                                 radius_equat_minor=a,
                                 radius_equat_major=b,
                                 radius_polar=c)
    elif model_name == 'parallelepiped':
        calculator = build_model('parallelepiped', n=n, a=a, b=b, c=c)
    else:
        raise ValueError("unknown model %s"%model_name)

    return calculator, (a, b, c)

SHAPES = [
    'parallelepiped',
    'sphere', 'ellipsoid', 'triaxial_ellipsoid',
    'cylinder',
    'fcc_paracrystal', 'bcc_paracrystal', 'sc_paracrystal',
]

DRAW_SHAPES = {
    'fcc_paracrystal': draw_fcc,
    'bcc_paracrystal': draw_bcc,
    'sc_paracrystal': draw_sc,
    'parallelepiped': draw_parallelepiped,
}

DISTRIBUTIONS = [
    'gaussian', 'rectangle', 'uniform',
]
DIST_LIMITS = {
    'gaussian': 30,
    'rectangle': 90/sqrt(3),
    'uniform': 90,
}


def run(model_name='parallelepiped', size=(10, 40, 100),
        view=(0, 0, 0), jitter=(0, 0, 0),
        dist='gaussian', mesh=30,
        projection='equirectangular'):
    """
    Show an interactive orientation and jitter demo.

    *model_name* is one of: sphere, ellipsoid, triaxial_ellipsoid,
    parallelepiped, cylinder, or sc/fcc/bcc_paracrystal

    *size* gives the dimensions (a, b, c) of the shape.

    *view* gives the initial view (theta, phi, psi) of the shape.

    *view* gives the initial jitter (dtheta, dphi, dpsi) of the shape.

    *dist* is the type of dispersition: gaussian, rectangle, or uniform.

    *mesh* is the number of points in the dispersion mesh.

    *projection* is the map projection to use for the mesh: equirectangular,
    sinusoidal, guyou, azimuthal_equidistance, or azimuthal_equal_area.
    """
    # projection number according to 1-order position in list, but
    # only 1 and 2 are implemented so far.
    from sasmodels import generate
    generate.PROJECTION = PROJECTIONS.index(projection) + 1
    if generate.PROJECTION > 2:
        print("*** PROJECTION %s not implemented in scattering function ***"%projection)
        generate.PROJECTION = 2

    # set up calculator
    calculator, size = select_calculator(model_name, n=150, size=size)
    draw_shape = DRAW_SHAPES.get(model_name, draw_parallelepiped)
    #draw_shape = draw_fcc

    ## uncomment to set an independent the colour range for every view
    ## If left commented, the colour range is fixed for all views
    calculator.limits = None

    PLOT_ENGINE(calculator, draw_shape, size, view, jitter, dist, mesh, projection)

def _mpl_plot(calculator, draw_shape, size, view, jitter, dist, mesh, projection):
    # Note: travis-ci does not support mpl_toolkits.mplot3d, but this shouldn't be
    # an issue since we are lazy-loading the package on a path that isn't tested.
    # Importing mplot3d adds projection='3d' option to subplot
    import mpl_toolkits.mplot3d  # pylint: disable=unused-variable
    import matplotlib as mpl
    import matplotlib.pyplot as plt
    from matplotlib.widgets import Slider

    ## create the plot window
    #plt.hold(True)
    plt.subplots(num=None, figsize=(5.5, 5.5))
    plt.set_cmap('gist_earth')
    plt.clf()
    plt.gcf().canvas.set_window_title(projection)
    #gs = gridspec.GridSpec(2,1,height_ratios=[4,1])
    #axes = plt.subplot(gs[0], projection='3d')
    axes = plt.axes([0.0, 0.2, 1.0, 0.8], projection='3d')
    try:  # CRUFT: not all versions of matplotlib accept 'square' 3d projection
        axes.axis('square')
    except Exception:
        pass

    # CRUFT: use axisbg instead of facecolor for matplotlib<2
    facecolor_prop = 'facecolor' if mpl.__version__ > '2' else 'axisbg'
    props = {facecolor_prop: 'lightgoldenrodyellow'}

    ## add control widgets to plot
    axes_theta = plt.axes([0.05, 0.15, 0.50, 0.04], **props)
    axes_phi = plt.axes([0.05, 0.10, 0.50, 0.04], **props)
    axes_psi = plt.axes([0.05, 0.05, 0.50, 0.04], **props)
    stheta = Slider(axes_theta, u'Ξ', -90, 90, valinit=0)
    sphi = Slider(axes_phi, u'φ', -180, 180, valinit=0)
    spsi = Slider(axes_psi, u'ψ', -180, 180, valinit=0)

    axes_dtheta = plt.axes([0.70, 0.15, 0.20, 0.04], **props)
    axes_dphi = plt.axes([0.70, 0.1, 0.20, 0.04], **props)
    axes_dpsi = plt.axes([0.70, 0.05, 0.20, 0.04], **props)

    # Note: using ridiculous definition of rectangle distribution, whose width
    # in sasmodels is sqrt(3) times the given width.  Divide by sqrt(3) to keep
    # the maximum width to 90.
    dlimit = DIST_LIMITS[dist]
    sdtheta = Slider(axes_dtheta, u'Δξ', 0, 2*dlimit, valinit=0)
    sdphi = Slider(axes_dphi, u'Δφ', 0, 2*dlimit, valinit=0)
    sdpsi = Slider(axes_dpsi, u'Δψ', 0, 2*dlimit, valinit=0)

    ## initial view and jitter
    theta, phi, psi = view
    stheta.set_val(theta)
    sphi.set_val(phi)
    spsi.set_val(psi)
    dtheta, dphi, dpsi = jitter
    sdtheta.set_val(dtheta)
    sdphi.set_val(dphi)
    sdpsi.set_val(dpsi)

    ## callback to draw the new view
    def _update(val, axis=None):
        # pylint: disable=unused-argument
        view = stheta.val, sphi.val, spsi.val
        jitter = sdtheta.val, sdphi.val, sdpsi.val
        # set small jitter as 0 if multiple pd dims
        dims = sum(v > 0 for v in jitter)
        limit = [0, 0.5, 5, 5][dims]
        jitter = [0 if v < limit else v for v in jitter]
        axes.cla()

        ## Visualize as person on globe
        #draw_sphere(axes, radius=0.5)
        #draw_person_on_sphere(axes, view, radius=0.5)

        ## Move beam instead of shape
        #draw_beam(axes, -view[:2])
        #draw_jitter(axes, (0,0,0), (0,0,0), views=3)

        ## Move shape and draw scattering
        draw_beam(axes, (0, 0), alpha=1.)
        #draw_person_on_sphere(axes, view, radius=1.2, height=0.5)
        draw_jitter(axes, view, jitter, dist=dist, size=size, alpha=1.,
                    draw_shape=draw_shape, projection=projection, views=3)
        draw_mesh(axes, view, jitter, dist=dist, n=mesh, projection=projection)
        draw_scattering(calculator, axes, view, jitter, dist=dist)

        plt.gcf().canvas.draw()

    ## bind control widgets to view updater
    stheta.on_changed(lambda v: _update(v, 'theta'))
    sphi.on_changed(lambda v: _update(v, 'phi'))
    spsi.on_changed(lambda v: _update(v, 'psi'))
    sdtheta.on_changed(lambda v: _update(v, 'dtheta'))
    sdphi.on_changed(lambda v: _update(v, 'dphi'))
    sdpsi.on_changed(lambda v: _update(v, 'dpsi'))

    ## initialize view
    _update(None, 'phi')

    ## go interactive
    plt.show()


def map_colors(z, kw):
    """
    Process matplotlib-style colour arguments.

    Pulls 'cmap', 'alpha', 'vmin', and 'vmax' from th *kw* dictionary, setting
    the *kw['color']* to an RGB array.  These are ignored if 'c' or 'color' are
    set inside *kw*.
    """
    from matplotlib import cm

    cmap = kw.pop('cmap', cm.coolwarm)
    alpha = kw.pop('alpha', None)
    vmin = kw.pop('vmin', z.min())
    vmax = kw.pop('vmax', z.max())
    c = kw.pop('c', None)
    color = kw.pop('color', c)
    if color is None:
        znorm = ((z - vmin) / (vmax - vmin)).clip(0, 1)
        color = cmap(znorm)
    elif isinstance(color, np.ndarray) and color.shape == z.shape:
        color = cmap(color)
    if alpha is None:
        if isinstance(color, np.ndarray):
            color = color[..., :3]
    else:
        color[..., 3] = alpha
    kw['color'] = color

def make_vec(*args):
    """Turn all elements of *args* into numpy arrays"""
    #return [np.asarray(v, 'd').flatten() for v in args]
    return [np.asarray(v, 'd') for v in args]

def make_image(z, kw):
    """Convert numpy array *z* into a *PIL* RGB image."""
    import PIL.Image
    from matplotlib import cm

    cmap = kw.pop('cmap', cm.coolwarm)

    znorm = (z-z.min())/z.ptp()
    c = cmap(znorm)
    c = c[..., :3]
    rgb = np.asarray(c*255, 'u1')
    image = PIL.Image.fromarray(rgb, mode='RGB')
    return image


_IPV_MARKERS = {
    'o': 'sphere',
}
_IPV_COLORS = {
    'w': 'white',
    'k': 'black',
    'c': 'cyan',
    'm': 'magenta',
    'y': 'yellow',
    'r': 'red',
    'g': 'green',
    'b': 'blue',
}
def _ipv_fix_color(kw):
    alpha = kw.pop('alpha', None)
    color = kw.get('color', None)
    if isinstance(color, str):
        color = _IPV_COLORS.get(color, color)
        kw['color'] = color
    if alpha is not None:
        color = kw['color']
        #TODO: convert color to [r, g, b, a] if not already
        if isinstance(color, (tuple, list)):
            if len(color) == 3:
                color = (color[0], color[1], color[2], alpha)
            else:
                color = (color[0], color[1], color[2], alpha*color[3])
            color = np.array(color)
        if isinstance(color, np.ndarray) and color.shape[-1] == 4:
            color[..., 3] = alpha
            kw['color'] = color

def _ipv_set_transparency(kw, obj):
    color = kw.get('color', None)
    if (isinstance(color, np.ndarray)
            and color.shape[-1] == 4
            and (color[..., 3] != 1.0).any()):
        obj.material.transparent = True
        obj.material.side = "FrontSide"

def ipv_axes():
    """
    Build a matplotlib style Axes interface for ipyvolume
    """
    import ipyvolume as ipv

    class Axes(object):
        """
        Matplotlib Axes3D style interface to ipyvolume renderer.
        """
        # pylint: disable=no-self-use,no-init
        # transparency can be achieved by setting the following:
        #    mesh.color = [r, g, b, alpha]
        #    mesh.material.transparent = True
        #    mesh.material.side = "FrontSide"
        # smooth(ish) rotation can be achieved by setting:
        #    slide.continuous_update = True
        #    figure.animation = 0.
        #    mesh.material.x = x
        #    mesh.material.y = y
        #    mesh.material.z = z
        # maybe need to synchronize update of x/y/z to avoid shimmy when moving
        def plot(self, x, y, z, **kw):
            """mpl style plot interface for ipyvolume"""
            _ipv_fix_color(kw)
            x, y, z = make_vec(x, y, z)
            ipv.plot(x, y, z, **kw)
        def plot_surface(self, x, y, z, **kw):
            """mpl style plot_surface interface for ipyvolume"""
            facecolors = kw.pop('facecolors', None)
            if facecolors is not None:
                kw['color'] = facecolors
            _ipv_fix_color(kw)
            x, y, z = make_vec(x, y, z)
            h = ipv.plot_surface(x, y, z, **kw)
            _ipv_set_transparency(kw, h)
            #h.material.side = "DoubleSide"
            return h
        def plot_trisurf(self, x, y, triangles=None, Z=None, **kw):
            """mpl style plot_trisurf interface for ipyvolume"""
            kw.pop('linewidth', None)
            _ipv_fix_color(kw)
            x, y, z = make_vec(x, y, Z)
            if triangles is not None:
                triangles = np.asarray(triangles)
            h = ipv.plot_trisurf(x, y, z, triangles=triangles, **kw)
            _ipv_set_transparency(kw, h)
            return h
        def scatter(self, x, y, z, **kw):
            """mpl style scatter interface for ipyvolume"""
            x, y, z = make_vec(x, y, z)
            map_colors(z, kw)
            marker = kw.get('marker', None)
            kw['marker'] = _IPV_MARKERS.get(marker, marker)
            h = ipv.scatter(x, y, z, **kw)
            _ipv_set_transparency(kw, h)
            return h
        def contourf(self, x, y, v, zdir='z', offset=0, levels=None, **kw):
            """mpl style contourf interface for ipyvolume"""
            # pylint: disable=unused-argument
            # Don't use contour for now (although we might want to later)
            self.pcolor(x, y, v, zdir='z', offset=offset, **kw)
        def pcolor(self, x, y, v, zdir='z', offset=0, **kw):
            """mpl style pcolor interface for ipyvolume"""
            # pylint: disable=unused-argument
            x, y, v = make_vec(x, y, v)
            image = make_image(v, kw)
            xmin, xmax = x.min(), x.max()
            ymin, ymax = y.min(), y.max()
            x = np.array([[xmin, xmax], [xmin, xmax]])
            y = np.array([[ymin, ymin], [ymax, ymax]])
            z = x*0 + offset
            u = np.array([[0., 1], [0, 1]])
            v = np.array([[0., 0], [1, 1]])
            h = ipv.plot_mesh(x, y, z, u=u, v=v, texture=image, wireframe=False)
            _ipv_set_transparency(kw, h)
            h.material.side = "DoubleSide"
            return h
        def text(self, *args, **kw):
            """mpl style text interface for ipyvolume"""
            pass
        def set_xlim(self, limits):
            """mpl style set_xlim interface for ipyvolume"""
            ipv.xlim(*limits)
        def set_ylim(self, limits):
            """mpl style set_ylim interface for ipyvolume"""
            ipv.ylim(*limits)
        def set_zlim(self, limits):
            """mpl style set_zlim interface for ipyvolume"""
            ipv.zlim(*limits)
        def set_axes_on(self):
            """mpl style set_axes_on interface for ipyvolume"""
            ipv.style.axis_on()
        def set_axis_off(self):
            """mpl style set_axes_off interface for ipyvolume"""
            ipv.style.axes_off()
    return Axes()

def _ipv_plot(calculator, draw_shape, size, view, jitter, dist, mesh, projection):
    from IPython.display import display
    import ipywidgets as widgets
    import ipyvolume as ipv

    axes = ipv_axes()

    def _draw(view, jitter):
        camera = ipv.gcf().camera
        #print(ipv.gcf().__dict__.keys())
        #print(dir(ipv.gcf()))
        ipv.figure(animation=0.)  # no animation when updating object mesh

        # set small jitter as 0 if multiple pd dims
        dims = sum(v > 0 for v in jitter)
        limit = [0, 0.5, 5, 5][dims]
        jitter = [0 if v < limit else v for v in jitter]

        ## Visualize as person on globe
        #draw_beam(axes, (0, 0))
        #draw_sphere(axes, radius=0.5)
        #draw_person_on_sphere(axes, view, radius=0.5)

        ## Move beam instead of shape
        #draw_beam(axes, view=(-view[0], -view[1]))
        #draw_jitter(axes, view=(0,0,0), jitter=(0,0,0))

        ## Move shape and draw scattering
        draw_beam(axes, (0, 0), steps=25)
        draw_jitter(axes, view, jitter, dist=dist, size=size, alpha=1.0,
                    draw_shape=draw_shape, projection=projection)
        draw_mesh(axes, view, jitter, dist=dist, n=mesh, radius=0.95,
                  projection=projection)
        draw_scattering(calculator, axes, view, jitter, dist=dist)

        draw_axes(axes, origin=(-1, -1, -1.1))
        ipv.style.box_off()
        ipv.style.axes_off()
        ipv.xyzlabel(" ", " ", " ")

        ipv.gcf().camera = camera
        ipv.show()


    trange, prange = (-180., 180., 1.), (-180., 180., 1.)
    dtrange, dprange = (0., 180., 1.), (0., 180., 1.)

    ## Super simple interfaca, but uses non-ascii variable namese
    # Ξ φ ψ Δξ Δφ Δψ
    #def update(**kw):
    #    view = kw['Ξ'], kw['φ'], kw['ψ']
    #    jitter = kw['Δξ'], kw['Δφ'], kw['Δψ']
    #    draw(view, jitter)
    #widgets.interact(update, Ξ=trange, φ=prange, ψ=prange, Δξ=dtrange, Δφ=dprange, Δψ=dprange)

    def _update(theta, phi, psi, dtheta, dphi, dpsi):
        _draw(view=(theta, phi, psi), jitter=(dtheta, dphi, dpsi))

    def _slider(name, slice, init=0.):
        return widgets.FloatSlider(
            value=init,
            min=slice[0],
            max=slice[1],
            step=slice[2],
            description=name,
            disabled=False,
            #continuous_update=True,
            continuous_update=False,
            orientation='horizontal',
            readout=True,
            readout_format='.1f',
            )
    theta = _slider(u'Ξ', trange, view[0])
    phi = _slider(u'φ', prange, view[1])
    psi = _slider(u'ψ', prange, view[2])
    dtheta = _slider(u'Δξ', dtrange, jitter[0])
    dphi = _slider(u'Δφ', dprange, jitter[1])
    dpsi = _slider(u'Δψ', dprange, jitter[2])
    fields = {
        'theta': theta, 'phi': phi, 'psi': psi,
        'dtheta': dtheta, 'dphi': dphi, 'dpsi': dpsi,
    }
    ui = widgets.HBox([
        widgets.VBox([theta, phi, psi]),
        widgets.VBox([dtheta, dphi, dpsi])
    ])

    out = widgets.interactive_output(_update, fields)
    display(ui, out)


_ENGINES = {
    "matplotlib": _mpl_plot,
    "mpl": _mpl_plot,
    #"plotly": _plotly_plot,
    "ipvolume": _ipv_plot,
    "ipv": _ipv_plot,
}
PLOT_ENGINE = _ENGINES["matplotlib"]
def set_plotter(name):
    """
    Setting the plotting engine to matplotlib/ipyvolume or equivalently mpl/ipv.
    """
    global PLOT_ENGINE
    PLOT_ENGINE = _ENGINES[name]

def main():
    """
    Command line interface to the jitter viewer.
    """
    parser = argparse.ArgumentParser(
        description="Display jitter",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
        )
    parser.add_argument('-p', '--projection', choices=PROJECTIONS,
                        default=PROJECTIONS[0],
                        help='coordinate projection')
    parser.add_argument('-s', '--size', type=str, default='10,40,100',
                        help='a,b,c lengths')
    parser.add_argument('-v', '--view', type=str, default='0,0,0',
                        help='initial view angles')
    parser.add_argument('-j', '--jitter', type=str, default='0,0,0',
                        help='initial angular dispersion')
    parser.add_argument('-d', '--distribution', choices=DISTRIBUTIONS,
                        default=DISTRIBUTIONS[0],
                        help='jitter distribution')
    parser.add_argument('-m', '--mesh', type=int, default=30,
                        help='#points in theta-phi mesh')
    parser.add_argument('shape', choices=SHAPES, nargs='?', default=SHAPES[0],
                        help='oriented shape')
    opts = parser.parse_args()
    size = tuple(float(v) for v in opts.size.split(','))
    view = tuple(float(v) for v in opts.view.split(','))
    jitter = tuple(float(v) for v in opts.jitter.split(','))
    run(opts.shape, size=size, view=view, jitter=jitter,
        mesh=opts.mesh, dist=opts.distribution,
        projection=opts.projection)

if __name__ == "__main__":
    main()