Esempio n. 1
0
def plot_reference_sphere(theta=0,
                          phi=0,
                          isometric=True,
                          ax=None,
                          azimuth=None,
                          elevation=None,
                          degrees=True,
                          labels=True,
                          light_ambient=.6,
                          zoom=1.4,
                          arrow_color=(.25, .95, .8),
                          close_figure=False):

    if azimuth is None:
        azimuth = DEFAULT_AZIMUTH
    if elevation is None:
        elevation = DEFAULT_ELEVATION

    import visvis as vv

    if ax is None:
        app = vv.use()
        fig = vv.figure()
        ax = vv.subplot(111)

    else:
        app = None
        fig = ax.GetFigure()

    ax.axis.visible = 0
    ax.axis.xLabel = 'x'
    ax.axis.yLabel = 'y'
    ax.axis.zLabel = 'z'

    # coordinate system
    length = 1.4
    cyl_diameter = .05
    for i in range(3):
        direction = np.zeros((3, ))
        direction[i] = 1
        cyl = vv.solidCylinder(translation=(0, 0, 0),
                               scaling=(cyl_diameter, cyl_diameter, length),
                               direction=tuple(direction),
                               axesAdjust=False,
                               axes=ax)
        cyl.faceColor = (.75, .75, .75)

        translation = np.zeros((3, ))
        translation[i] = length
        cone = vv.solidCone(translation=tuple(translation),
                            scaling=(.1, .1, .2),
                            direction=tuple(direction),
                            axesAdjust=False,
                            axes=ax)
        cone.faceColor = (.75, .75, .75)

    # example direction
    length = 1
    shrink = .825
    direction = sph2cart(1, theta, phi, degrees=degrees).ravel()
    cyl = vv.solidCylinder(translation=(0, 0, 0),
                           scaling=(.05, .05, shrink * length),
                           direction=tuple(direction),
                           axesAdjust=False,
                           axes=ax)
    cyl.faceColor = arrow_color

    translation = direction * shrink
    cone = vv.solidCone(translation=tuple(translation),
                        scaling=(.1, .1, .2),
                        direction=tuple(direction),
                        axesAdjust=False,
                        axes=ax)
    cone.faceColor = arrow_color

    # indicate unit sphere
    sphere = vv.solidSphere((0, 0, 0), N=100, M=100)
    sphere.faceColor = (.5, .5, .5, .25)

    # some lines on sphere indicating main axes
    phi = np.linspace(0, 2 * np.pi, 100)
    theta = np.ones_like(phi) * np.pi / 2.
    r = np.ones_like(phi)
    xyz = sph2cart(r, theta, phi, degrees=False)
    vv.plot(xyz[:, 0],
            xyz[:, 1],
            xyz[:, 2],
            lw=1,
            lc=(.5, .5, .5),
            ls="-",
            mc='b',
            axesAdjust=False,
            axes=ax)

    theta = np.linspace(-np.pi, np.pi, 100)
    phi = np.ones_like(theta) * 0
    r = np.ones_like(phi)
    xyz = sph2cart(r, theta, phi, degrees=False)
    vv.plot(xyz[:, 0],
            xyz[:, 1],
            xyz[:, 2],
            lw=1,
            lc=(.5, .5, .5),
            ls="-",
            mc='b',
            axesAdjust=False,
            axes=ax)

    theta = np.linspace(-np.pi, np.pi, 100)
    phi = np.ones_like(theta) * np.pi / 2.
    r = np.ones_like(phi)
    xyz = sph2cart(r, theta, phi, degrees=False)
    vv.plot(xyz[:, 0],
            xyz[:, 1],
            xyz[:, 2],
            lw=1,
            lc=(.5, .5, .5),
            ls="-",
            mc='b',
            axesAdjust=False,
            axes=ax)

    # add pitch and roll axes
    aa = np.deg2rad(np.linspace(45, 315, 100) - 90)
    r = .25 + .025
    d = 1.25 + .05
    color = (.25, .25, .25)

    # pitch rotation (in x/z plane, i.e. around y-axis)
    xx = r * np.cos(aa)
    zz = r * np.sin(aa)
    yy = d * np.ones_like(xx)
    vv.plot(xx, yy, zz, lw=5, lc=color, ls="-", axesAdjust=False, axes=ax)

    translation = (xx[0], yy[0], zz[0])
    direction = (xx[0] - xx[1], yy[0] - yy[1], zz[0] - zz[1])
    cone = vv.solidCone(translation=translation,
                        scaling=(.05, .05, .1),
                        direction=direction,
                        axesAdjust=False,
                        axes=ax)
    cone.faceColor = color

    if labels:
        vv.Text(ax, 'Pitch', x=0, y=1.25 * d, z=.25, fontSize=28, color=color)

    # roll rotation (in y/z plane, i.e. around x-axis)
    yy = r * np.cos(aa)
    zz = r * np.sin(aa)
    xx = d * np.ones_like(xx)
    vv.plot(xx, yy, zz, lw=5, lc=color, ls="-", axesAdjust=False, axes=ax)

    translation = (xx[-1], yy[-1], zz[-1])
    direction = (xx[-1] - xx[-2], yy[-1] - yy[-2], zz[-1] - zz[-2])
    cone = vv.solidCone(translation=translation,
                        scaling=(.05, .05, .1),
                        direction=direction,
                        axesAdjust=False,
                        axes=ax)
    cone.faceColor = color

    if labels:
        vv.Text(ax, 'Roll', x=1.25 * d, y=-.8, z=0, fontSize=28, color=color)

    # set camera view
    zoom_ = vv.view()['zoom']
    if isometric:
        vv.view(dict(azimuth=90 + ISO_AZIMUTH, elevation=ISO_ELEVATION),
                zoom=zoom * zoom_,
                axes=ax)
    else:
        vv.view(dict(azimuth=90 + azimuth, elevation=elevation),
                zoom=zoom * zoom_,
                _axes=ax)

    ax.light0.ambient = light_ambient
    ax.light0.specular = .5

    fig.DrawNow()

    if app is not None:
        app.ProcessEvents()

    img = vv.screenshot(None, ob=ax, sf=2, bg=None, format=None)

    if close_figure:
        vv.close(fig)
        fig = None

    return fig, img
Esempio n. 2
0
def plot_mouse_head(show_now=False,
                    size=(500, 500),
                    ax=None,
                    theta=0,
                    phi=0,
                    pitch=None,
                    roll=None,
                    isometric=True,
                    azimuth=None,
                    elevation=None,
                    zoom=1.5,
                    gravity_axes=True,
                    head_axes=True,
                    close_figure=False,
                    color_mouse=(.5, .5, .5),
                    color_head_axes=(.25, .95, .8),
                    color_gravity_axes=(.75, .75, .75),
                    backend='visvis',
                    light_ambient=.6,
                    light_specular=.5,
                    arrow_kw=dict(length_cone=.1,
                                  radius_cone=.05,
                                  radius_shaft=.025)):

    if azimuth is None:
        azimuth = DEFAULT_AZIMUTH
    if elevation is None:
        elevation = DEFAULT_ELEVATION

    t = np.arange(0, 1 * np.pi - 2 * np.pi / 10, np.pi / 10)
    r = 2 + np.cos(t)
    n = 20

    if backend in ['mayavi', 'mlab']:

        from mayavi import mlab
        from tvtk.tools import visual

        fig = mlab.figure(bgcolor=(1, 1, 1), size=size)
        fig.scene.parallel_projection = True

        # head
        c = -3
        [X, Y, Z] = cylinder(r, n=n)
        head = mlab.mesh(X, Y, c + Z * 6, color=(.5, .5, .5))

        # nose
        [x, y, z] = sphere(20)
        nose = mlab.mesh(x * 1.5, y * 1.5, c + z * 1.5 + 6, color=(.5, .5, .5))

        # EARS
        color = (.5, .5, .5)
        hsE1 = mlab.mesh((x * 1.0) - 2.427, (y * 1.8) - 1.763,
                         c + (z * 0.4) + 0,
                         color=color)
        hsE2 = mlab.mesh((x * 1.0) + 2.427, (y * 1.8) - 1.763,
                         c + (z * 0.4) + 0,
                         color=color)

        # EYES
        [x, y, z] = sphere(10)
        color = (.9, .9, .9)
        hsEYE1 = mlab.mesh((x * .8) - 1.2, (y * .8) - 1.6,
                           c + (z * .8) + 3.5,
                           color=color)
        hsEYE2 = mlab.mesh((x * .8) + 1.2, (y * .8) - 1.6,
                           c + (z * .8) + 3.5,
                           color=color)

        # pupils
        hsPUP1 = mlab.mesh((x * .3) - 1.476, (y * .3) - 1.98,
                           c + (z * .3) + 4.147,
                           color=(0.1, 0.1, 0.1))
        hsPUP2 = mlab.mesh((x * .3) + 1.476, (y * .3) - 1.98,
                           c + (z * .3) + 4.147,
                           color=(0.1, 0.1, 0.1))

        # whiskers
        Px = np.array([-1.154, -1.214, -1.154, +1.154, +1.214, +1.154])
        Py = np.array([-0.375, 0, 0.375, -0.375, 0, 0.375])
        Pz = np.ones(np.size(Px)) * 3.882
        WL = np.array([[0.5, 2], [0.05, 0.4]])  # whisker length

        hsWHISK = []
        for W in range(0, len(Px)):  # W=0

            if Px[W] < 0:
                XSIGN = -1
            else:
                XSIGN = +1

            if Py[W] < 0:
                YSIGN = -1
            elif Py[W] == 0:
                YSIGN = 0
            else:
                YSIGN = +1

            x = np.append(Px[W], Px[W] + XSIGN * WL[0, :])
            y = np.append(Py[W], Py[W] + YSIGN * WL[1, :])
            z = np.append(Pz[W], Pz[W])

            tck, u = interpolate.splprep([x, y], k=2)
            xi, yi = interpolate.splev(np.linspace(0, 1, 100), tck, der=0)

            zi = np.ones(np.size(yi)) * Pz[W]
            hsWHISK.append(
                mlab.plot3d(xi, yi, zi, color=(0, 0, 0), line_width=2))

        # rotate all objects
        angle_x = -90
        angle_y = 0
        angle_z = -90

        for actor in [head, nose, hsE1, hsE2, hsEYE1, hsEYE2, hsPUP1, hsPUP2]:
            actor.actor.actor.rotate_z(angle_z)
            actor.actor.actor.rotate_x(angle_x)
            actor.actor.actor.rotate_y(angle_y)

            #        actor.actor.actor.rotate_y(phi)
            #        actor.actor.actor.rotate_z(theta)

            actor.actor.actor.rotate_x(theta)
            actor.actor.actor.rotate_y(phi)

        for i in range(0, len(hsWHISK)):
            hsWHISK[i].actor.actor.rotate_z(angle_z)
            hsWHISK[i].actor.actor.rotate_x(angle_x)
            hsWHISK[i].actor.actor.rotate_y(angle_y)

            hsWHISK[i].actor.actor.rotate_x(theta)
            hsWHISK[i].actor.actor.rotate_y(phi)

        if head_axes:
            # axes of head-centered coordinate system
            visual.set_viewer(fig)

            arr1 = Arrow3D(0,
                           0,
                           0,
                           0,
                           -6,
                           0,
                           color=color_head_axes,
                           **arrow_kw)

            for arr in [arr1]:  # , arr2, arr3]:

                arr.actor.rotate_z(angle_z)
                arr.actor.rotate_x(angle_x)
                arr.actor.rotate_y(angle_y)

                arr.actor.rotate_x(theta)
                arr.actor.rotate_y(phi)

        if gravity_axes:

            # gravitational coordinate system
            visual.set_viewer(fig)

            arr1 = Arrow3D(0, 0, 0, -6, 0, 0, color=(.8, .8, .8), **arrow_kw)
            arr2 = Arrow3D(0, 0, 0, 0, -6, 0, color=(.5, .5, .5), **arrow_kw)
            arr3 = Arrow3D(0, 0, 0, 0, 0, 6, color=(.25, .25, .25), **arrow_kw)

            for arr in [arr1, arr2, arr3]:

                arr.actor.rotate_z(angle_z)
                arr.actor.rotate_x(angle_x)
                arr.actor.rotate_y(angle_y)

        if isometric:
            fig.scene.isometric_view()
        else:
            mlab.view(azimuth=azimuth, elevation=elevation, figure=fig)
        fig.scene.camera.zoom(zoom)

        if show_now:
            mlab.show()

        img = mlab.screenshot(figure=fig, mode='rgb', antialiased=False)

    elif backend in ['visvis']:

        import visvis as vv

        if ax is None:
            app = vv.use()
            fig = vv.figure()
            ax = vv.subplot(111)

        else:
            app = None
            fig = ax.GetFigure()

        ax.bgcolor = (1, 1, 1)

        ax.axis.visible = 0
        ax.axis.xLabel = 'x'
        ax.axis.yLabel = 'y'
        ax.axis.zLabel = 'z'

        # head
        c = -3
        [X, Y, Z] = cylinder(r, n=n)
        head = vv.surf(c + 6 * Z, X, Y)
        head.faceColor = color_mouse

        # nose
        [x, y, z] = sphere(20)
        nose = vv.surf(c + z * 1.5 + 6, x * 1.5, y * 1.5)
        nose.faceColor = color_mouse

        # ears
        color = (.5, .5, .5)
        ear1 = vv.surf(c + (z * 0.4) + 0, (x * 1.0) - 2.427,
                       -1 * ((y * 1.8) - 1.763))
        ear1.faceColor = color
        ear2 = vv.surf(c + (z * 0.4) + 0, (x * 1.0) + 2.427,
                       -1 * ((y * 1.8) - 1.763))
        ear2.faceColor = color_mouse

        # eyes
        [x, y, z] = sphere(10)
        color = (.9, .9, .9)
        eye1 = vv.surf(c + (z * .8) + 3.5, (x * .8) - 1.2,
                       -1 * ((y * .8) - 1.6))
        eye2 = vv.surf(c + (z * .8) + 3.5, (x * .8) + 1.2,
                       -1 * ((y * .8) - 1.6))
        [setattr(eye, 'faceColor', color) for eye in [eye1, eye2]]

        # pupils
        color = (.1, .1, .1)
        pupil1 = vv.surf(c + (z * .3) + 4.147 - .5, (x * .3) - 1.476 - .2,
                         -1 * ((y * .3) - 1.98))
        pupil2 = vv.surf(c + (z * .3) + 4.147 - .5, (x * .3) + 1.476 + .2,
                         -1 * ((y * .3) - 1.98))
        [setattr(pupil, 'faceColor', color) for pupil in [pupil1, pupil2]]

        # whiskers
        Px = np.array([-1.154, -1.214, -1.154, +1.154, +1.214, +1.154])
        Py = np.array([-0.375, 0, 0.375, -0.375, 0, 0.375])
        Pz = np.ones(np.size(Px)) * 3.882
        WL = np.array([[0.5, 2], [0.05, 0.4]])  # whisker length

        whiskers = []
        for W in range(0, len(Px)):  # W=0

            if Px[W] < 0:
                XSIGN = -1
            else:
                XSIGN = +1

            if Py[W] < 0:
                YSIGN = -1
            elif Py[W] == 0:
                YSIGN = 0
            else:
                YSIGN = +1

            x = np.append(Px[W], Px[W] + XSIGN * WL[0, :])
            y = np.append(Py[W], Py[W] + YSIGN * WL[1, :])
            z = np.append(Pz[W], Pz[W])

            tck, u = interpolate.splprep([x, y], k=2)
            xi, yi = interpolate.splev(np.linspace(0, 1, 100), tck, der=0)

            zi = np.ones(np.size(yi)) * Pz[W]
            whisker = vv.plot(zi, xi, -1 * yi, lw=2, lc=(0, 0, 0))
            whiskers.append(whisker)

        if head_axes:
            # show vector indicating orientation of head
            length = 1
            shrink = .825

            if pitch is not None and roll is not None:
                # convert pitch/roll to spherical coordinates by rotating
                # a reference point around cartesian axes; note that x and
                # y are swapped in visvis
                p0 = np.array([0, 0, 1])

                a = np.deg2rad(roll)
                Rx = np.array([[1, 0, 0], [0, np.cos(a), -np.sin(a)],
                               [0, np.sin(a), np.cos(a)]])

                b = np.deg2rad(pitch)
                Ry = np.array([[np.cos(b), 0, np.sin(b)], [0, 1, 0],
                               [-np.sin(b), 0, np.cos(b)]])

                direction = np.dot(Ry, np.dot(Rx, p0)) * length
            else:
                direction = sph2cart(length, theta, phi, degrees=True).ravel()
            cyl = vv.solidCylinder(translation=(0, 0, .25),
                                   scaling=(.05, .05, shrink * length),
                                   direction=tuple(direction),
                                   axesAdjust=False,
                                   axes=ax)
            cyl.faceColor = (.25, .95, .8)
            cyl.do_not_rotate = True
            cyl.do_not_scale = True

            translation = direction * shrink + np.array([0, 0, .25])
            cone = vv.solidCone(translation=tuple(translation),
                                scaling=(.1, .1, .2),
                                direction=tuple(direction),
                                axesAdjust=False,
                                axes=ax)
            cone.faceColor = (.25, .95, .8)
            cone.do_not_rotate = True
            cone.do_not_scale = True

        if gravity_axes:
            # show vector indicating (negative) orientation of gravity
            length = 1
            shrink = .825
            color = (.75, .75, .75)

            direction = np.array([0, 0, 1])
            cyl = vv.solidCylinder(translation=(0, 0, .25),
                                   scaling=(.05, .05, shrink * length),
                                   direction=tuple(direction),
                                   axesAdjust=False,
                                   axes=ax)
            cyl.faceColor = color
            cyl.do_not_rotate = True
            cyl.do_not_scale = True

            translation = direction * shrink + np.array([0, 0, .25])
            cone = vv.solidCone(translation=tuple(translation),
                                scaling=(.1, .1, .2),
                                direction=tuple(direction),
                                axesAdjust=False,
                                axes=ax)
            cone.faceColor = color
            cone.do_not_rotate = True
            cone.do_not_scale = True

        # compute pitch and/or roll if not given
        if pitch is None or roll is None:

            # we have to find rotations (=angles) between the unit vector
            # in spherical coordinates (1, theta, phi) and the planes
            # for pitch (y/z) and roll (x/z); note that xyz is already
            # normaizeed (norm = 1)
            xyz = sph2cart(1, theta, phi, degrees=True).ravel()

            if pitch is None:
                # pitch (y/z plane with normal vector (1, 0, 0))
                pitch = 90 - np.degrees(np.arccos(np.dot(xyz, [1, 0, 0])))

            if roll is None:
                # roll (x/z plane with normal vector (0, -1, 0))
                roll = 90 - np.degrees(np.arccos(np.dot(xyz, [0, -1, 0])))

        for obj in ax.wobjects:

            if not hasattr(obj, 'do_not_scale'):
                rot = vv.Transform_Scale(sx=.25, sy=.25, sz=.25)
                obj.transformations.append(rot)

            if obj.__class__ != vv.core.axises.CartesianAxis and\
                    not hasattr(obj, 'do_not_rotate'):

                # note that x and y are swapped
                rot = vv.Transform_Rotate(pitch, ax=0, ay=1, az=0)
                obj.transformations.append(rot)

                rot = vv.Transform_Rotate(roll, ax=1, ay=0, az=0)
                obj.transformations.append(rot)

        zoom = vv.view()['zoom']
        if isometric:
            vv.view(dict(azimuth=90 + ISO_AZIMUTH, elevation=ISO_ELEVATION),
                    zoom=4 * zoom,
                    axes=ax)
        else:
            vv.view(dict(azimuth=90 + azimuth, elevation=elevation),
                    zoom=4 * zoom,
                    axes=ax)

        ax.light0.ambient = light_ambient
        ax.light0.specular = light_specular

        fig.DrawNow()

        if app is not None:
            app.ProcessEvents()

        img = vv.screenshot(None, ob=ax, sf=2, bg=None, format=None)

        if close_figure:
            vv.close(fig)
            fig = None

    return fig, img
Esempio n. 3
0
    
    Parameters
    ----------
    viewparams : dict
        View parameters to set.
    axes : Axes instance
        The axes the view parameters are for.  Uses the current axes by default.
    keyword pairs
        View parameters to set.  These take precidence.
    
    If neither viewparams or any keyword pairs are given, returns the current
    view parameters (as a dict).  Otherwise, sets the view parameters given.
    """
    
    if axes is None:
        axes = vv.gca()
    
    if viewparams or kw:
        axes.SetView(viewparams, **kw)
    else:
        return axes.GetView()


if __name__=='__main__':
    a = vv.gca()
    vv.solidTeapot()
    v = vv.view()
    # ... rotate the figure a bit and then call:
    vv.view(v)
    # Note that a.GetView() and a.SetView(v) are equivalent
Esempio n. 4
0
def plot_density_sphere(xyz_centers,
                        H,
                        ax=None,
                        method='hist',
                        azimuth=None,
                        elevation=None,
                        backend='visvis',
                        oversample=1,
                        smoothing=None,
                        zoom=1.7,
                        cmap='jet',
                        scaling='log',
                        log_offset=-.1,
                        clim=None,
                        isometric=False,
                        light_ambient=.6,
                        light_specular=.5,
                        avg_direction=None,
                        pitch_label=True,
                        roll_label=True,
                        pitch_arrow=True,
                        roll_arrow=True,
                        arrow_color=(.25, .95, .8),
                        close_figure=False):

    if azimuth is None:
        azimuth = DEFAULT_AZIMUTH
    if elevation is None:
        elevation = DEFAULT_ELEVATION

    scaling = scaling.lower()
    if scaling.startswith('log'):
        H = H / float(H.sum())
        v = H > 0
        if scaling == 'log':
            H[v] = np.log(H[v])
        elif scaling == 'log2':
            H[v] = np.log2(H[v])
        H[~v] = H[v].min() + log_offset

    if smoothing is not None and smoothing > 1:
        x = np.linspace(-3., 3., smoothing)
        xx, yy = np.meshgrid(x, x)
        G = np.exp((-xx**2 - yy**2))
        H = signal.convolve2d(H, G / G.sum(), mode='same', boundary='wrap')

    if backend in ['mpl', 'matplotlib']:

        if ax is None:
            fig = plt.figure()
            ax = fig.add_subplot(111, projection='3d')

        from matplotlib.colors import LightSource
        ls = LightSource(azdeg=315, altdeg=45)
        cm = getattr(plt.cm, cmap)
        colors = ls.shade(H, cmap=cm, blend_mode='soft', vert_exag=1)
        ax.plot_surface(xyz_centers[:, 0].reshape(H),
                        xyz_centers[:, 1].reshape(H),
                        xyz_centers[:, 2].reshape(H),
                        cstride=1,
                        rstride=1,
                        facecolors=colors,
                        linewidth=0,
                        antialiased=False,
                        shade=False)

        # add arrows in front of sphere
        import mousecam_helpers as helpers
        arrow_props = dict(arrowstyle='simple',
                           mutation_scale=15,
                           mutation_aspect=None,
                           linewidth=.5,
                           facecolor=3 * [.75],
                           edgecolor=3 * [.1])
        length = .6

        arr = helpers.Arrow3D((1, 1 + length), (0, 0), (0, 0), **arrow_props)
        ax.add_artist(arr)

        arr = helpers.Arrow3D((0, 0), (1, 1 + length), (0, 0), **arrow_props)
        ax.add_artist(arr)

        arr = helpers.Arrow3D((0, 0), (0, 0), (1, 1 + length), **arrow_props)
        ax.add_artist(arr)

        extents = 3 * [[-.6, .6]]
        ax.auto_scale_xyz(*extents)
        ax.set_aspect('equal')
        ax.axis('off')

        ax.view_init(elev=elevation, azim=360 - azimuth)

    elif backend in ['mayavi', 'mlab']:

        from mayavi import mlab

        fig = mlab.figure(bgcolor=(1, 1, 1), size=(1000, 1000))
        fig.scene.parallel_projection = True

        mlab.mesh(xyz_centers[:, 0].reshape(H.shape),
                  xyz_centers[:, 1].reshape(H.shape),
                  xyz_centers[:, 2].reshape(H.shape),
                  scalars=H,
                  colormap='jet')

        from tvtk.tools import visual
        visual.set_viewer(fig)
        arrow_kw = dict(color=(.75, .75, .75),
                        length_cone=.1,
                        radius_cone=.05,
                        radius_shaft=.025)
        Arrow3D(0, 0, 0, 1.4, 0, 0, **arrow_kw)
        Arrow3D(0, 0, 0, 0, 1.4, 0, **arrow_kw)
        Arrow3D(0, 0, 0, 0, 0, 1.4, **arrow_kw)

        if isometric:
            fig.scene.isometric_view()
        else:
            mlab.view(azimuth=-azimuth, elevation=elevation, figure=fig)
        fig.scene.camera.zoom(zoom)

        ax = mlab.screenshot(figure=fig, mode='rgb', antialiased=False)

    elif backend in ['visvis']:

        import visvis as vv

        app = vv.use()

        fig = vv.figure()
        ax = vv.subplot(111)

        ax.axis.visible = 0
        ax.axis.xLabel = 'x'
        ax.axis.yLabel = 'y'
        ax.axis.zLabel = 'z'

        length = 1.4

        for i in range(3):
            direction = np.zeros((3, ))
            direction[i] = 1
            cyl = vv.solidCylinder(translation=(0, 0, 0),
                                   scaling=(.05, .05, length),
                                   direction=tuple(direction),
                                   axesAdjust=False,
                                   axes=ax)
            cyl.faceColor = (.75, .75, .75)

            translation = np.zeros((3, ))
            translation[i] = length
            cone = vv.solidCone(translation=tuple(translation),
                                scaling=(.1, .1, .2),
                                direction=tuple(direction),
                                axesAdjust=False,
                                axes=ax)
            cone.faceColor = (.75, .75, .75)

        surf = vv.surf(xyz_centers[:, 0].reshape(H.shape),
                       xyz_centers[:, 1].reshape(H.shape),
                       xyz_centers[:, 2].reshape(H.shape),
                       H,
                       axes=ax)

        # set colormap
        cmap = cmap.lower()
        if cmap.endswith('_r'):
            cm = vv.colormaps[cmap[:-2]]
            cm = cm[::-1]
        else:
            cm = vv.colormaps[cmap]
        surf.colormap = cm

        if clim is not None:
            # apply colormap limits
            surf.clim = clim


#            ax.Draw()

# add labels to indicate pitch and/or roll axes
        aa = np.deg2rad(np.linspace(45, 315, 100) - 90)
        r = .25 + .05
        d = 1.25 + .25
        color = (.25, .25, .25)

        if pitch_arrow:
            # pitch rotation (in x/z plane, i.e. around y-axis)
            xx = r * np.cos(aa)
            zz = r * np.sin(aa)
            yy = d * np.ones_like(xx)
            vv.plot(xx,
                    yy,
                    zz,
                    lw=5,
                    lc=color,
                    ls="-",
                    axesAdjust=False,
                    axes=ax)

            translation = (xx[0], yy[0], zz[0])
            direction = (xx[0] - xx[1], yy[0] - yy[1], zz[0] - zz[1])
            cone = vv.solidCone(translation=translation,
                                scaling=(.05, .05, .1),
                                direction=direction,
                                axesAdjust=False,
                                axes=ax)
            cone.faceColor = color

        if pitch_label:
            vv.Text(ax, 'Pitch', x=0, y=.9 * d, z=.5, fontSize=28, color=color)

        if roll_arrow:
            # roll rotation (in y/z plane, i.e. around x-axis)
            yy = r * np.cos(aa[::-1])
            zz = r * np.sin(aa[::-1])
            xx = d * np.ones_like(xx)
            vv.plot(xx,
                    yy,
                    zz,
                    lw=5,
                    lc=color,
                    ls="-",
                    axesAdjust=False,
                    axes=ax)

            translation = (xx[0], yy[0], zz[0])
            direction = (xx[0] - xx[1], yy[0] - yy[1], zz[0] - zz[1])
            cone = vv.solidCone(translation=translation,
                                scaling=(.05, .05, .1),
                                direction=direction,
                                axesAdjust=False,
                                axes=ax)
            cone.faceColor = color

        if roll_label:
            vv.Text(ax,
                    'Roll',
                    x=1.25 * d,
                    y=-.8,
                    z=0,
                    fontSize=28,
                    color=color)

        if avg_direction is not None:
            # indicate direction of avg head orientation
            avg_direction = avg_direction / np.sqrt(np.sum(avg_direction**2))
            avg_direction *= length
            cyl = vv.solidCylinder(translation=(0, 0, 0),
                                   scaling=(.05, .05, length),
                                   direction=tuple(avg_direction),
                                   axesAdjust=False,
                                   axes=ax)
            cyl.faceColor = arrow_color

            cone = vv.solidCone(translation=tuple(avg_direction),
                                scaling=(.1, .1, .2),
                                direction=tuple(avg_direction),
                                axesAdjust=False,
                                axes=ax)
            cone.faceColor = arrow_color

        zoom_ = vv.view()['zoom']
        if isometric:
            vv.view(dict(azimuth=90 + ISO_AZIMUTH, elevation=ISO_ELEVATION),
                    zoom=zoom_ * zoom,
                    axes=ax)
        else:
            vv.view(dict(azimuth=90 + azimuth, elevation=elevation),
                    zoom=zoom * zoom_,
                    axes=ax)

        ax.light0.ambient = light_ambient
        ax.light0.specular = light_specular

        fig.DrawNow()
        app.ProcessEvents()

        img = vv.screenshot(None, ob=ax, sf=2, bg=None, format=None)
        ax = img

        if close_figure:
            vv.close(fig)
            fig = None

    return fig, ax
Esempio n. 5
0
    
    Parameters
    ----------
    viewparams : dict
        View parameters to set.
    axes : Axes instance
        The axes the view parameters are for.  Uses the current axes by default.
    keyword pairs
        View parameters to set.  These take precidence.
    
    If neither viewparams or any keyword pairs are given, returns the current
    view parameters (as a dict).  Otherwise, sets the view parameters given.
    """

    if axes is None:
        axes = vv.gca()

    if viewparams or kw:
        axes.SetView(viewparams, **kw)
    else:
        return axes.GetView()


if __name__ == '__main__':
    a = vv.gca()
    vv.solidTeapot()
    v = vv.view()
    # ... rotate the figure a bit and then call:
    vv.view(v)
    # Note that a.GetView() and a.SetView(v) are equivalent