Changeset 8678a34 in sasmodels

Timestamp:

Apr 9, 2017 5:45:20 PM (8 years ago)

Author:

Paul Kienzle <pkienzle@…>

Branches:

master, core_shell_microgels, magnetic_model, ticket-1257-vesicle-product, ticket_1156, ticket_1265_superball, ticket_822_more_unit_tests

Children:

a0ebc96

Parents:

85190c2

Message:

code cleanup on jitter visualization

File:

• explore/jitter.py (modified) (11 diffs)

explore/jitter.py

@@ -3 +3 @@
 dispersity and possible replacement algorithms.
 """
+import sys
+
 import mpl_toolkits.mplot3d   # Adds projection='3d' option to subplot
 import matplotlib.pyplot as plt
 from matplotlib.widgets import Slider, CheckButtons
 from matplotlib import cm
-
 import numpy as np
 from numpy import pi, cos, sin, sqrt, exp, degrees, radians
 
-def draw_beam(ax):
+def draw_beam(ax, view=(0, 0)):
     #ax.plot([0,0],[0,0],[1,-1])
     #ax.scatter([0]*100,[0]*100,np.linspace(1, -1, 100), alpha=0.8)
@@ -22 +23 @@
     x = r*np.outer(np.cos(u), np.ones_like(v))
     y = r*np.outer(np.sin(u), np.ones_like(v))
-    z = np.outer(np.ones_like(u), 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]
 
     ax.plot_surface(x, y, z, rstride=4, cstride=4, color='y', alpha=0.5)
 
-def draw_shimmy(ax, theta, phi, psi, dtheta, dphi, dpsi):
-    size=[0.1, 0.4, 1.0]
-    view=[theta, phi, psi]
-    shimmy=[0,0,0]
-    #draw_shape = draw_parallelepiped
-    draw_shape = draw_ellipsoid
+def draw_jitter(ax, view, jitter):
+    size = [0.1, 0.4, 1.0]
+    draw_shape = draw_parallelepiped
+    #draw_shape = draw_ellipsoid
 
     #np.random.seed(10)
@@ -64 +69 @@
         [ 1,  1,  1],
     ]
+    dtheta, dphi, dpsi = jitter
     if dtheta == 0:
         cloud = [v for v in cloud if v[0] == 0]
@@ -70 +76 @@
     if dpsi == 0:
         cloud = [v for v in cloud if v[2] == 0]
-    draw_shape(ax, size, view, shimmy, steps=100, alpha=0.8)
+    draw_shape(ax, size, view, [0, 0, 0], steps=100, alpha=0.8)
     for point in cloud:
-        shimmy=[dtheta*point[0], dphi*point[1], dpsi*point[2]]
-        draw_shape(ax, size, view, shimmy, alpha=0.8)
+        delta = [dtheta*point[0], dphi*point[1], dpsi*point[2]]
+        draw_shape(ax, size, view, delta, alpha=0.8)
     for v in 'xyz':
         a, b, c = size
@@ -80 +86 @@
         getattr(ax, v+'axis').label.set_text(v)
 
-def draw_ellipsoid(ax, size, view, shimmy, steps=25, alpha=1):
+def draw_ellipsoid(ax, size, view, jitter, steps=25, alpha=1):
     a,b,c = size
-    theta, phi, psi = view
-    dtheta, dphi, dpsi = shimmy
-
     u = np.linspace(0, 2 * np.pi, steps)
     v = np.linspace(0, np.pi, steps)
@@ -90 +93 @@
     y = b*np.outer(np.sin(u), np.sin(v))
     z = c*np.outer(np.ones_like(u), np.cos(v))
-
-    shape = x.shape
-    points = np.matrix([x.flatten(),y.flatten(),z.flatten()])
-    points = Rz(dpsi)*Ry(dtheta)*Rx(dphi)*points
-    points = Rz(phi)*Ry(theta)*Rz(psi)*points
-    x,y,z = [v.reshape(shape) for v in points]
+    x, y, z = transform_xyz(view, jitter, x, y, z)
 
     ax.plot_surface(x, y, z, rstride=4, cstride=4, color='w', alpha=alpha)
 
-def draw_parallelepiped(ax, size, view, shimmy, alpha=1):
+    draw_labels(ax, 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_parallelepiped(ax, size, view, jitter, steps=None, alpha=1):
     a,b,c = size
-    theta, phi, psi = view
-    dtheta, dphi, dpsi = shimmy
-
     x = a*np.array([ 1,-1, 1,-1, 1,-1, 1,-1])
     y = b*np.array([ 1, 1,-1,-1, 1, 1,-1,-1])
@@ -118 +122 @@
     ])
 
-    points = np.matrix([x,y,z])
-    points = Rz(dpsi)*Ry(dtheta)*Rx(dphi)*points
-    points = Rz(phi)*Ry(theta)*Rz(psi)*points
-
-    x,y,z = [np.array(v).flatten() for v in points]
+    x, y, z = transform_xyz(view, jitter, x, y, z)
     ax.plot_trisurf(x, y, triangles=tri, Z=z, color='w', alpha=alpha)
 
-def draw_sphere(ax, radius=10., steps=100):
-    u = np.linspace(0, 2 * np.pi, steps)
-    v = np.linspace(0, np.pi, steps)
-
-    x = radius * np.outer(np.cos(u), np.sin(v))
-    y = radius * np.outer(np.sin(u), np.sin(v))
-    z = radius * np.outer(np.ones(np.size(u)), np.cos(v))
-    ax.plot_surface(x, y, z, rstride=4, cstride=4, color='w')
-
-def draw_mesh_new(ax, theta, dtheta, phi, dphi, flow, radius=10., dist='gauss'):
-    theta_center = radians(theta)
-    phi_center = radians(phi)
-    flow_center = radians(flow)
-    dtheta = radians(dtheta)
-    dphi = radians(dphi)
-
-    # 10 point 3-sigma gaussian weights
-    t = np.linspace(-3., 3., 11)
+    draw_labels(ax, 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_mesh(ax, view, jitter, radius=1.2, n=11, dist='gauss'):
+    theta, phi, psi = view
+    dtheta, dphi, dpsi = jitter
     if dist == 'gauss':
+        t = np.linspace(-3, 3, n)
         weights = exp(-0.5*t**2)
     elif dist == 'rect':
+        t = np.linspace(0, 1, n)
         weights = np.ones_like(t)
     else:
         raise ValueError("expected dist to be 'gauss' or 'rect'")
-    theta = dtheta*t
-    phi = dphi*t
-
-    x = radius * np.outer(cos(phi), cos(theta))
-    y = radius * np.outer(sin(phi), cos(theta))
-    z = radius * np.outer(np.ones_like(phi), sin(theta))
-    #w = np.outer(weights, weights*abs(cos(dtheta*t)))
-    w = np.outer(weights, weights*abs(cos(theta)))
-
-    x, y, z, w = [v.flatten() for v in (x,y,z,w)]
-    x, y, z = rotate(x, y, z, phi_center, theta_center, flow_center)
-
-    ax.scatter(x, y, z, c=w, marker='o', vmin=0., vmax=1.)
-
-def rotate(x, y, z, phi, theta, psi):
-    R = Rz(phi)*Ry(theta)*Rz(psi)
-    p = np.vstack([x,y,z])
-    return R*p
+
+    # mesh in theta, phi formed by rotating z
+    z = np.matrix([[0], [0], [radius]])
+    points = np.hstack([Rx(phi_i)*Ry(theta_i)*z
+                        for theta_i in dtheta*t
+                        for phi_i in dphi*t])
+    # rotate relative to beam
+    points = orient_relative_to_beam(view, points)
+
+    w = np.outer(weights, weights)
+
+    x, y, z = [np.array(v).flatten() for v in points]
+    ax.scatter(x, y, z, c=w.flatten(), marker='o', vmin=0., vmax=1.)
 
 def Rx(angle):
@@ -188 +179 @@
          [0., 0., 1.]]
     return np.matrix(R)
+
+def transform_xyz(view, jitter, x, y, z):
+    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):
+    dtheta, dphi, dpsi = jitter
+    points = Rz(dpsi)*Ry(dtheta)*Rx(dphi)*points
+    return points
+
+def orient_relative_to_beam(view, points):
+    theta, phi, psi = view
+    points = Rz(phi)*Ry(theta)*Rz(psi)*points
+    return points
+
+def draw_labels(ax, view, jitter, text):
+    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)
+
+    for label, p, zdir in zip(labels, zip(px, py, pz), zip(dx, dy, dz)):
+        zdir = np.asarray(zdir).flatten()
+        ax.text(p[0], p[1], p[2], label, zdir=zdir)
+
+def draw_sphere(ax, radius=10., steps=100):
+    u = np.linspace(0, 2 * np.pi, steps)
+    v = np.linspace(0, np.pi, steps)
+
+    x = radius * np.outer(np.cos(u), np.sin(v))
+    y = radius * np.outer(np.sin(u), np.sin(v))
+    z = radius * np.outer(np.ones(np.size(u)), np.cos(v))
+    ax.plot_surface(x, y, z, rstride=4, cstride=4, color='w')
 
 def main():
@@ -206 +237 @@
 
     axcolor = 'lightgoldenrodyellow'
+
     axtheta  = plt.axes([0.1, 0.15, 0.45, 0.04], axisbg=axcolor)
     axphi = plt.axes([0.1, 0.1, 0.45, 0.04], axisbg=axcolor)
@@ -212 +244 @@
     sphi = Slider(axphi, 'Phi', -180, 180, valinit=phi)
     spsi = Slider(axpsi, 'Psi', -180, 180, valinit=psi)
+
     axdtheta  = plt.axes([0.75, 0.15, 0.15, 0.04], axisbg=axcolor)
     axdphi = plt.axes([0.75, 0.1, 0.15, 0.04], axisbg=axcolor)
@@ -217 +250 @@
     sdtheta = Slider(axdtheta, 'dTheta', 0, 30, valinit=dtheta)
     sdphi = Slider(axdphi, 'dPhi', 0, 30, valinit=dphi)
-    sdpsi = Slider(axdpsi, 'dPsi', 0, 30, valinit=dphi)
+    sdpsi = Slider(axdpsi, 'dPsi', 0, 30, valinit=dpsi)
 
     def update(val, axis=None):
-        theta, phi, psi = stheta.val, sphi.val, spsi.val
-        dtheta, dphi, dpsi = sdtheta.val, sdphi.val, sdpsi.val
+        view = stheta.val, sphi.val, spsi.val
+        jitter = sdtheta.val, sdphi.val, sdpsi.val
         ax.cla()
-        draw_beam(ax)
-        draw_shimmy(ax, theta, phi, psi, dtheta, dphi, dpsi)
-        #if not axis.startswith('d'):
-        #    ax.view_init(elev=theta, azim=phi)
+        draw_beam(ax, (0, 0))
+        if 0:
+            draw_jitter(ax, view, jitter)
+        else:
+            draw_jitter(ax, view, (0,0,0))
+            draw_mesh(ax, view, jitter)
         plt.gcf().canvas.draw()