Changeset 802c412 in sasmodels


Ignore:
Timestamp:
Mar 26, 2018 4:55:17 PM (4 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:
b3703f5
Parents:
d86f0fc
Message:

lint reduction

File:
1 edited

Legend:

Unmodified
Added
Removed
  • sasmodels/jitter.py

    rd86f0fc r802c412  
    1616 
    1717import matplotlib.pyplot as plt 
    18 from matplotlib.widgets import Slider, CheckButtons 
    19 from matplotlib import cm 
     18from matplotlib.widgets import Slider 
    2019import numpy as np 
    2120from numpy import pi, cos, sin, sqrt, exp, degrees, radians 
    2221 
    23 def draw_beam(ax, view=(0, 0)): 
     22def draw_beam(axes, view=(0, 0)): 
    2423    """ 
    2524    Draw the beam going from source at (0, 0, 1) to detector at (0, 0, -1) 
    2625    """ 
    27     #ax.plot([0,0],[0,0],[1,-1]) 
    28     #ax.scatter([0]*100,[0]*100,np.linspace(1, -1, 100), alpha=0.8) 
     26    #axes.plot([0,0],[0,0],[1,-1]) 
     27    #axes.scatter([0]*100,[0]*100,np.linspace(1, -1, 100), alpha=0.8) 
    2928 
    3029    steps = 25 
     
    4342    x, y, z = [v.reshape(shape) for v in points] 
    4443 
    45     ax.plot_surface(x, y, z, rstride=4, cstride=4, color='y', alpha=0.5) 
    46  
    47 def draw_ellipsoid(ax, size, view, jitter, steps=25, alpha=1): 
     44    axes.plot_surface(x, y, z, rstride=4, cstride=4, color='y', alpha=0.5) 
     45 
     46def draw_ellipsoid(axes, size, view, jitter, steps=25, alpha=1): 
    4847    """Draw an ellipsoid.""" 
    4948    a, b, c = size 
     
    5554    x, y, z = transform_xyz(view, jitter, x, y, z) 
    5655 
    57     ax.plot_surface(x, y, z, rstride=4, cstride=4, color='w', alpha=alpha) 
    58  
    59     draw_labels(ax, view, jitter, [ 
     56    axes.plot_surface(x, y, z, rstride=4, cstride=4, color='w', alpha=alpha) 
     57 
     58    draw_labels(axes, view, jitter, [ 
    6059        ('c+', [+0, +0, +c], [+1, +0, +0]), 
    6160        ('c-', [+0, +0, -c], [+0, +0, -1]), 
     
    6665    ]) 
    6766 
    68 def draw_sc(ax, size, view, jitter, steps=None, alpha=1): 
     67def draw_sc(axes, size, view, jitter, steps=None, alpha=1): 
    6968    """Draw points for simple cubic paracrystal""" 
    7069    atoms = _build_sc() 
    71     _draw_crystal(ax, size, view, jitter, atoms=atoms) 
    72  
    73 def draw_fcc(ax, size, view, jitter, steps=None, alpha=1): 
     70    _draw_crystal(axes, size, view, jitter, atoms=atoms) 
     71 
     72def draw_fcc(axes, size, view, jitter, steps=None, alpha=1): 
    7473    """Draw points for face-centered cubic paracrystal""" 
    7574    # Build the simple cubic crystal 
     
    8483    # y and z planes can be generated by substituting x for y and z respectively 
    8584    atoms.extend(zip(x+y+z, y+z+x, z+x+y)) 
    86     _draw_crystal(ax, size, view, jitter, atoms=atoms) 
    87  
    88 def draw_bcc(ax, size, view, jitter, steps=None, alpha=1): 
     85    _draw_crystal(axes, size, view, jitter, atoms=atoms) 
     86 
     87def draw_bcc(axes, size, view, jitter, steps=None, alpha=1): 
    8988    """Draw points for body-centered cubic paracrystal""" 
    9089    # Build the simple cubic crystal 
     
    9897    ) 
    9998    atoms.extend(zip(x, y, z)) 
    100     _draw_crystal(ax, size, view, jitter, atoms=atoms) 
    101  
    102 def _draw_crystal(ax, size, view, jitter, steps=None, alpha=1, atoms=None): 
     99    _draw_crystal(axes, size, view, jitter, atoms=atoms) 
     100 
     101def _draw_crystal(axes, size, view, jitter, steps=None, alpha=1, atoms=None): 
    103102    atoms, size = np.asarray(atoms, 'd').T, np.asarray(size, 'd') 
    104103    x, y, z = atoms*size[:, None] 
    105104    x, y, z = transform_xyz(view, jitter, x, y, z) 
    106     ax.scatter([x[0]], [y[0]], [z[0]], c='yellow', marker='o') 
    107     ax.scatter(x[1:], y[1:], z[1:], c='r', marker='o') 
     105    axes.scatter([x[0]], [y[0]], [z[0]], c='yellow', marker='o') 
     106    axes.scatter(x[1:], y[1:], z[1:], c='r', marker='o') 
    108107 
    109108def _build_sc(): 
     
    124123    return atoms 
    125124 
    126 def draw_parallelepiped(ax, size, view, jitter, steps=None, alpha=1): 
     125def draw_parallelepiped(axes, size, view, jitter, steps=None, alpha=1): 
    127126    """Draw a parallelepiped.""" 
    128127    a, b, c = size 
     
    142141 
    143142    x, y, z = transform_xyz(view, jitter, x, y, z) 
    144     ax.plot_trisurf(x, y, triangles=tri, Z=z, color='w', alpha=alpha) 
     143    axes.plot_trisurf(x, y, triangles=tri, Z=z, color='w', alpha=alpha) 
    145144 
    146145    # Draw pink face on box. 
     
    153152        z = c*np.array([+1, +1, +1, +1, -1, -1, -1, -1]) 
    154153        x, y, z = transform_xyz(view, jitter, x, y, abs(z)+0.001) 
    155         ax.plot_trisurf(x, y, triangles=tri, Z=z, color=[1, 0.6, 0.6], alpha=alpha) 
    156  
    157     draw_labels(ax, view, jitter, [ 
     154        axes.plot_trisurf(x, y, triangles=tri, Z=z, color=[1, 0.6, 0.6], alpha=alpha) 
     155 
     156    draw_labels(axes, view, jitter, [ 
    158157        ('c+', [+0, +0, +c], [+1, +0, +0]), 
    159158        ('c-', [+0, +0, -c], [+0, +0, -1]), 
     
    164163    ]) 
    165164 
    166 def draw_sphere(ax, radius=10., steps=100): 
     165def draw_sphere(axes, radius=10., steps=100): 
    167166    """Draw a sphere""" 
    168167    u = np.linspace(0, 2 * np.pi, steps) 
     
    172171    y = radius * np.outer(np.sin(u), np.sin(v)) 
    173172    z = radius * np.outer(np.ones(np.size(u)), np.cos(v)) 
    174     ax.plot_surface(x, y, z, rstride=4, cstride=4, color='w') 
    175  
    176 def draw_jitter(ax, view, jitter, dist='gaussian', size=(0.1, 0.4, 1.0), 
     173    axes.plot_surface(x, y, z, rstride=4, cstride=4, color='w') 
     174 
     175def draw_jitter(axes, view, jitter, dist='gaussian', size=(0.1, 0.4, 1.0), 
    177176                draw_shape=draw_parallelepiped): 
    178177    """ 
     
    221220    if dpsi == 0: 
    222221        cloud = [v for v in cloud if v[2] == 0] 
    223     draw_shape(ax, size, view, [0, 0, 0], steps=100, alpha=0.8) 
     222    draw_shape(axes, size, view, [0, 0, 0], steps=100, alpha=0.8) 
    224223    scale = {'gaussian':1, 'rectangle':1/sqrt(3), 'uniform':1/3}[dist] 
    225224    for point in cloud: 
    226225        delta = [scale*dtheta*point[0], scale*dphi*point[1], scale*dpsi*point[2]] 
    227         draw_shape(ax, size, view, delta, alpha=0.8) 
     226        draw_shape(axes, size, view, delta, alpha=0.8) 
    228227    for v in 'xyz': 
    229228        a, b, c = size 
    230229        lim = np.sqrt(a**2 + b**2 + c**2) 
    231         getattr(ax, 'set_'+v+'lim')([-lim, lim]) 
    232         getattr(ax, v+'axis').label.set_text(v) 
     230        getattr(axes, 'set_'+v+'lim')([-lim, lim]) 
     231        getattr(axes, v+'axis').label.set_text(v) 
    233232 
    234233PROJECTIONS = [ 
     
    238237    'azimuthal_equal_area', 
    239238] 
    240 def draw_mesh(ax, view, jitter, radius=1.2, n=11, dist='gaussian', 
     239def draw_mesh(axes, view, jitter, radius=1.2, n=11, dist='gaussian', 
    241240              projection='equirectangular'): 
    242241    """ 
     
    297296        Should allow free movement in theta, but phi is distorted. 
    298297    """ 
     298    # TODO: try Kent distribution instead of a gaussian warped by projection 
     299 
    299300    t = np.linspace(-1, 1, n) 
    300301    weights = np.ones_like(t) 
     
    313314        def _rotate(theta_i, phi_j): 
    314315            return Rx(phi_j)*Ry(theta_i) 
    315         def _weight(theta_i, phi_j, wi, wj): 
    316             return wi*wj*abs(cos(radians(theta_i))) 
     316        def _weight(theta_i, phi_j, w_i, w_j): 
     317            return w_i*w_j*abs(cos(radians(theta_i))) 
    317318    elif projection == 'sinusoidal':  #define PROJECTION 2 
    318319        def _rotate(theta_i, phi_j): 
     
    322323            #print("(%+7.2f, %+7.2f) => (%+7.2f, %+7.2f)"%(theta_i, phi_j, latitude, longitude)) 
    323324            return Rx(longitude)*Ry(latitude) 
    324         def _weight(theta_i, phi_j, wi, wj): 
     325        def _weight(theta_i, phi_j, w_i, w_j): 
    325326            latitude = theta_i 
    326327            scale = cos(radians(latitude)) 
    327             w = 1 if abs(phi_j) < abs(scale)*180 else 0 
    328             return w*wi*wj 
     328            active = 1 if abs(phi_j) < abs(scale)*180 else 0 
     329            return active*w_i*w_j 
    329330    elif projection == 'guyou':  #define PROJECTION 3  (eventually?) 
    330331        def _rotate(theta_i, phi_j): 
    331             from guyou import guyou_invert 
     332            from .guyou import guyou_invert 
    332333            #latitude, longitude = guyou_invert([theta_i], [phi_j]) 
    333334            longitude, latitude = guyou_invert([phi_j], [theta_i]) 
    334335            return Rx(longitude[0])*Ry(latitude[0]) 
    335         def _weight(theta_i, phi_j, wi, wj): 
    336             return wi*wj 
     336        def _weight(theta_i, phi_j, w_i, w_j): 
     337            return w_i*w_j 
    337338    elif projection == 'azimuthal_equidistance':  # Note: Rz Ry, not Rx Ry 
    338339        def _rotate(theta_i, phi_j): 
     
    341342            #print("(%+7.2f, %+7.2f) => (%+7.2f, %+7.2f)"%(theta_i, phi_j, latitude, longitude)) 
    342343            return Rz(longitude)*Ry(latitude) 
    343         def _weight(theta_i, phi_j, wi, wj): 
     344        def _weight(theta_i, phi_j, w_i, w_j): 
    344345            # Weighting for each point comes from the integral: 
    345346            #     \int\int I(q, lat, log) sin(lat) dlat dlog 
     
    371372            # the entire sphere, and treats theta and phi identically. 
    372373            latitude = sqrt(theta_i**2 + phi_j**2) 
    373             w = sin(radians(latitude))/latitude if latitude != 0 else 1 
    374             return w*wi*wj if latitude < 180 else 0 
     374            weight = sin(radians(latitude))/latitude if latitude != 0 else 1 
     375            return weight*w_i*w_j if latitude < 180 else 0 
    375376    elif projection == 'azimuthal_equal_area': 
    376377        def _rotate(theta_i, phi_j): 
    377             R = min(1, sqrt(theta_i**2 + phi_j**2)/180) 
    378             latitude = 180-degrees(2*np.arccos(R)) 
     378            radius = min(1, sqrt(theta_i**2 + phi_j**2)/180) 
     379            latitude = 180-degrees(2*np.arccos(radius)) 
    379380            longitude = degrees(np.arctan2(phi_j, theta_i)) 
    380381            #print("(%+7.2f, %+7.2f) => (%+7.2f, %+7.2f)"%(theta_i, phi_j, latitude, longitude)) 
    381382            return Rz(longitude)*Ry(latitude) 
    382         def _weight(theta_i, phi_j, wi, wj): 
     383        def _weight(theta_i, phi_j, w_i, w_j): 
    383384            latitude = sqrt(theta_i**2 + phi_j**2) 
    384             w = sin(radians(latitude))/latitude if latitude != 0 else 1 
    385             return w*wi*wj if latitude < 180 else 0 
     385            weight = sin(radians(latitude))/latitude if latitude != 0 else 1 
     386            return weight*w_i*w_j if latitude < 180 else 0 
    386387    else: 
    387388        raise ValueError("unknown projection %r"%projection) 
     
    393394                        for theta_i in dtheta*t 
    394395                        for phi_j in dphi*t]) 
    395     # select just the active points (i.e., those with phi < 180 
    396     w = np.array([_weight(theta_i, phi_j, wi, wj) 
    397                   for wi, theta_i in zip(weights, dtheta*t) 
    398                   for wj, phi_j in zip(weights, dphi*t)]) 
     396    w = np.array([_weight(theta_i, phi_j, w_i, w_j) 
     397                  for w_i, theta_i in zip(weights, dtheta*t) 
     398                  for w_j, phi_j in zip(weights, dphi*t)]) 
    399399    #print(max(w), min(w), min(w[w>0])) 
    400400    points = points[:, w > 0] 
     
    402402    w /= max(w) 
    403403 
    404     if 0: # Kent distribution 
    405         points = np.hstack([Rx(phi_j)*Ry(theta_i)*z for theta_i in 30*t for phi_j in 60*t]) 
    406         xp, yp, zp = [np.array(v).flatten() for v in points] 
    407         kappa = max(1e6, radians(dtheta)/(2*pi)) 
    408         beta = 1/max(1e-6, radians(dphi)/(2*pi))/kappa 
    409         w = exp(kappa*zp) #+ beta*(xp**2 + yp**2) 
    410         print(kappa, dtheta, radians(dtheta), min(w), max(w), sum(w)) 
    411         #w /= abs(cos(radians( 
    412         #w /= sum(w) 
    413  
    414404    # rotate relative to beam 
    415405    points = orient_relative_to_beam(view, points) 
     
    417407    x, y, z = [np.array(v).flatten() for v in points] 
    418408    #plt.figure(2); plt.clf(); plt.hist(z, bins=np.linspace(-1, 1, 51)) 
    419     ax.scatter(x, y, z, c=w, marker='o', vmin=0., vmax=1.) 
    420  
    421 def draw_labels(ax, view, jitter, text): 
     409    axes.scatter(x, y, z, c=w, marker='o', vmin=0., vmax=1.) 
     410 
     411def draw_labels(axes, view, jitter, text): 
    422412    """ 
    423413    Draw text at a particular location. 
     
    433423    for label, p, zdir in zip(labels, zip(px, py, pz), zip(dx, dy, dz)): 
    434424        zdir = np.asarray(zdir).flatten() 
    435         ax.text(p[0], p[1], p[2], label, zdir=zdir) 
     425        axes.text(p[0], p[1], p[2], label, zdir=zdir) 
    436426 
    437427# Definition of rotation matrices comes from wikipedia: 
     
    439429def Rx(angle): 
    440430    """Construct a matrix to rotate points about *x* by *angle* degrees.""" 
    441     a = radians(angle) 
     431    angle = radians(angle) 
    442432    R = [[1, 0, 0], 
    443          [0, +cos(a), -sin(a)], 
    444          [0, +sin(a), +cos(a)]] 
     433         [0, +cos(angle), -sin(angle)], 
     434         [0, +sin(angle), +cos(angle)]] 
    445435    return np.matrix(R) 
    446436 
    447437def Ry(angle): 
    448438    """Construct a matrix to rotate points about *y* by *angle* degrees.""" 
    449     a = radians(angle) 
    450     R = [[+cos(a), 0, +sin(a)], 
     439    angle = radians(angle) 
     440    R = [[+cos(angle), 0, +sin(angle)], 
    451441         [0, 1, 0], 
    452          [-sin(a), 0, +cos(a)]] 
     442         [-sin(angle), 0, +cos(angle)]] 
    453443    return np.matrix(R) 
    454444 
    455445def Rz(angle): 
    456446    """Construct a matrix to rotate points about *z* by *angle* degrees.""" 
    457     a = radians(angle) 
    458     R = [[+cos(a), -sin(a), 0], 
    459          [+sin(a), +cos(a), 0], 
     447    angle = radians(angle) 
     448    R = [[+cos(angle), -sin(angle), 0], 
     449         [+sin(angle), +cos(angle), 0], 
    460450         [0, 0, 1]] 
    461451    return np.matrix(R) 
     
    524514        return data[offset], data[-1] 
    525515 
    526 def draw_scattering(calculator, ax, view, jitter, dist='gaussian'): 
     516def draw_scattering(calculator, axes, view, jitter, dist='gaussian'): 
    527517    """ 
    528518    Plot the scattering for the particular view. 
    529519 
    530     *calculator* is returned from :func:`build_model`.  *ax* are the 3D axes 
     520    *calculator* is returned from :func:`build_model`.  *axes* are the 3D axes 
    531521    on which the data will be plotted.  *view* and *jitter* are the current 
    532522    orientation and orientation dispersity.  *dist* is one of the sasmodels 
     
    571561        level[level < 0] = 0 
    572562        colors = plt.get_cmap()(level) 
    573         ax.plot_surface(qx, qy, -1.1, rstride=1, cstride=1, facecolors=colors) 
     563        axes.plot_surface(qx, qy, -1.1, rstride=1, cstride=1, facecolors=colors) 
    574564    elif 1: 
    575         ax.contourf(qx/qx.max(), qy/qy.max(), Iqxy, zdir='z', offset=-1.1, 
     565        axes.contourf(qx/qx.max(), qy/qy.max(), Iqxy, zdir='z', offset=-1.1, 
    576566                    levels=np.linspace(vmin, vmax, 24)) 
    577567    else: 
    578         ax.pcolormesh(qx, qy, Iqxy) 
     568        axes.pcolormesh(qx, qy, Iqxy) 
    579569 
    580570def build_model(model_name, n=150, qmax=0.5, **pars): 
     
    749739    plt.gcf().canvas.set_window_title(projection) 
    750740    #gs = gridspec.GridSpec(2,1,height_ratios=[4,1]) 
    751     #ax = plt.subplot(gs[0], projection='3d') 
    752     ax = plt.axes([0.0, 0.2, 1.0, 0.8], projection='3d') 
     741    #axes = plt.subplot(gs[0], projection='3d') 
     742    axes = plt.axes([0.0, 0.2, 1.0, 0.8], projection='3d') 
    753743    try:  # CRUFT: not all versions of matplotlib accept 'square' 3d projection 
    754         ax.axis('square') 
     744        axes.axis('square') 
    755745    except Exception: 
    756746        pass 
     
    759749 
    760750    ## add control widgets to plot 
    761     axtheta = plt.axes([0.1, 0.15, 0.45, 0.04], axisbg=axcolor) 
    762     axphi = plt.axes([0.1, 0.1, 0.45, 0.04], axisbg=axcolor) 
    763     axpsi = plt.axes([0.1, 0.05, 0.45, 0.04], axisbg=axcolor) 
    764     stheta = Slider(axtheta, 'Theta', -90, 90, valinit=theta) 
    765     sphi = Slider(axphi, 'Phi', -180, 180, valinit=phi) 
    766     spsi = Slider(axpsi, 'Psi', -180, 180, valinit=psi) 
    767  
    768     axdtheta = plt.axes([0.75, 0.15, 0.15, 0.04], axisbg=axcolor) 
    769     axdphi = plt.axes([0.75, 0.1, 0.15, 0.04], axisbg=axcolor) 
    770     axdpsi = plt.axes([0.75, 0.05, 0.15, 0.04], axisbg=axcolor) 
     751    axes_theta = plt.axes([0.1, 0.15, 0.45, 0.04], axisbg=axcolor) 
     752    axes_phi = plt.axes([0.1, 0.1, 0.45, 0.04], axisbg=axcolor) 
     753    axes_psi = plt.axes([0.1, 0.05, 0.45, 0.04], axisbg=axcolor) 
     754    stheta = Slider(axes_theta, 'Theta', -90, 90, valinit=theta) 
     755    sphi = Slider(axes_phi, 'Phi', -180, 180, valinit=phi) 
     756    spsi = Slider(axes_psi, 'Psi', -180, 180, valinit=psi) 
     757 
     758    axes_dtheta = plt.axes([0.75, 0.15, 0.15, 0.04], axisbg=axcolor) 
     759    axes_dphi = plt.axes([0.75, 0.1, 0.15, 0.04], axisbg=axcolor) 
     760    axes_dpsi = plt.axes([0.75, 0.05, 0.15, 0.04], axisbg=axcolor) 
    771761    # Note: using ridiculous definition of rectangle distribution, whose width 
    772762    # in sasmodels is sqrt(3) times the given width.  Divide by sqrt(3) to keep 
    773763    # the maximum width to 90. 
    774764    dlimit = DIST_LIMITS[dist] 
    775     sdtheta = Slider(axdtheta, 'dTheta', 0, 2*dlimit, valinit=dtheta) 
    776     sdphi = Slider(axdphi, 'dPhi', 0, 2*dlimit, valinit=dphi) 
    777     sdpsi = Slider(axdpsi, 'dPsi', 0, 2*dlimit, valinit=dpsi) 
     765    sdtheta = Slider(axes_dtheta, 'dTheta', 0, 2*dlimit, valinit=dtheta) 
     766    sdphi = Slider(axes_dphi, 'dPhi', 0, 2*dlimit, valinit=dphi) 
     767    sdpsi = Slider(axes_dpsi, 'dPsi', 0, 2*dlimit, valinit=dpsi) 
    778768 
    779769 
     
    786776        limit = [0, 0, 0.5, 5][dims] 
    787777        jitter = [0 if v < limit else v for v in jitter] 
    788         ax.cla() 
    789         draw_beam(ax, (0, 0)) 
    790         draw_jitter(ax, view, jitter, dist=dist, size=size, draw_shape=draw_shape) 
    791         #draw_jitter(ax, view, (0,0,0)) 
    792         draw_mesh(ax, view, jitter, dist=dist, n=mesh, projection=projection) 
    793         draw_scattering(calculator, ax, view, jitter, dist=dist) 
     778        axes.cla() 
     779        draw_beam(axes, (0, 0)) 
     780        draw_jitter(axes, view, jitter, dist=dist, size=size, draw_shape=draw_shape) 
     781        #draw_jitter(axes, view, (0,0,0)) 
     782        draw_mesh(axes, view, jitter, dist=dist, n=mesh, projection=projection) 
     783        draw_scattering(calculator, axes, view, jitter, dist=dist) 
    794784        plt.gcf().canvas.draw() 
    795785 
Note: See TracChangeset for help on using the changeset viewer.