← Back to branch summary

~ubuntu-branches/ubuntu/trusty/mayavi2/trusty

« back to all changes in this revision

Viewing changes to mayavi/tests/datasets.py

  • Committer: Bazaar Package Importer
  • Author(s): Varun Hiremath
  • Date: 2011-07-09 01:18:36 UTC
  • mfrom: (1.1.10 upstream) (2.2.3 sid)
  • Revision ID: james.westby@ubuntu.com-20110709011836-fha21zirlgkqh92s
Tags: 4.0.0-1
http://bugs.debian.org/625148
* New upstream release
* debian/control:
  - Bump Standards-Version to 3.9.2
  - Set X-Python-Version: 2.6, fixes FTBFS (Closes: #625148)
  - Update Depends
* Update debian/watch file
* Cleanup debian/rules file

Show diffs side-by-side

added added

removed removed

Lines of Context:
mayavi/tests/datasets.py
 
1
# Author: Suyog Dutt Jain <suyog.jain@aero.iitb.ac.in>
 
2
#         Prabhu Ramachandran <prabhu_r@users.sf.net>
 
3
# Copyright (c) 2008,  Enthought, Inc.
 
4
# License: BSD Style.
 
5
 
 
6
# Standard library imports.
 
7
 
 
8
import numpy
 
9
from numpy import linspace, cos, sin, pi, empty, sqrt,array,arange,random
 
10
 
 
11
# Enthought library imports
 
12
from tvtk.api import tvtk
 
13
 
 
14
 
 
15
def generate_annulus(r=None, theta=None, z=None):
 
16
    """ Generate points for structured grid for a cylindrical annular
 
17
        volume.  This method is useful for generating a unstructured
 
18
        cylindrical mesh for VTK (and perhaps other tools).
 
19
 
 
20
        Parameters
 
21
        ----------
 
22
        r : array : The radial values of the grid points.
 
23
                    It defaults to linspace(1.0, 2.0, 11).
 
24
 
 
25
        theta : array : The angular values of the x axis for the grid
 
26
                        points. It defaults to linspace(0,2*pi,11).
 
27
 
 
28
        z: array : The values along the z axis of the grid points.
 
29
                   It defaults to linspace(0,0,1.0, 11).
 
30
 
 
31
        Return
 
32
        ------
 
33
        points : array
 
34
            Nx3 array of points that make up the volume of the annulus.
 
35
            They are organized in planes starting with the first value
 
36
            of z and with the inside "ring" of the plane as the first
 
37
            set of points.  The default point array will be 1331x3.
 
38
    """
 
39
    # Default values for the annular grid.
 
40
    if r is None: r = linspace(1.0,2.0, 11)
 
41
    if theta is None: theta = linspace(0,2*pi,11)
 
42
    if z is None: z = linspace(0.0,1.0, 11)
 
43
 
 
44
    # Find the x values and y values for each plane.
 
45
    x_plane = (cos(theta)*r[:,None]).ravel()
 
46
    y_plane = (sin(theta)*r[:,None]).ravel()
 
47
 
 
48
    # Allocate an array for all the points.  We'll have len(x_plane)
 
49
    # points on each plane, and we have a plane for each z value, so
 
50
    # we need len(x_plane)*len(z) points.
 
51
    points = empty([len(x_plane)*len(z),3])
 
52
 
 
53
    # Loop through the points for each plane and fill them with the
 
54
    # correct x,y,z values.
 
55
    start = 0
 
56
    for z_plane in z:
 
57
        end = start+len(x_plane)
 
58
        # slice out a plane of the output points and fill it
 
59
        # with the x,y, and z values for this plane.  The x,y
 
60
        # values are the same for every plane.  The z value
 
61
        # is set to the current z
 
62
        plane_points = points[start:end]
 
63
        plane_points[:,0] = x_plane
 
64
        plane_points[:,1] = y_plane
 
65
        plane_points[:,2] = z_plane
 
66
        start = end
 
67
 
 
68
    return points
 
69
 
 
70
 
 
71
def single_type_ug():
 
72
    """Simple example showing how to create an unstructured grid
 
73
    consisting of cells of a single type.
 
74
    """
 
75
    points = array([[0,0,0], [1,0,0], [0,1,0], [0,0,1], # tets
 
76
                    [1,0,0], [2,0,0], [1,1,0], [1,0,1],
 
77
                    [2,0,0], [3,0,0], [2,1,0], [2,0,1],
 
78
                    ], 'f')
 
79
    tets = array([[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10, 11]])
 
80
    tet_type = tvtk.Tetra().cell_type
 
81
    ug = tvtk.UnstructuredGrid(points=points)
 
82
    ug.set_cells(tet_type, tets)
 
83
    return ug
 
84
 
 
85
def mixed_type_ug():
 
86
    """A slightly more complex example of how to generate an
 
87
    unstructured grid with different cell types.  Returns a created
 
88
    unstructured grid.
 
89
    """
 
90
    points = array([[0,0,0], [1,0,0], [0,1,0], [0,0,1], # tetra
 
91
                    [2,0,0], [3,0,0], [3,1,0], [2,1,0],
 
92
                    [2,0,1], [3,0,1], [3,1,1], [2,1,1], # Hex
 
93
                    ], 'f')
 
94
    # shift the points so we can show both.
 
95
    points[:,1] += 2.0
 
96
    # The cells
 
97
    cells = array([4, 0, 1, 2, 3, # tetra
 
98
                   8, 4, 5, 6, 7, 8, 9, 10, 11 # hex
 
99
                   ])
 
100
    # The offsets for the cells, i.e. the indices where the cells
 
101
    # start.
 
102
    offset = array([0, 5])
 
103
    tetra_type = tvtk.Tetra().cell_type # VTK_TETRA == 10
 
104
    hex_type = tvtk.Hexahedron().cell_type # VTK_HEXAHEDRON == 12
 
105
    cell_types = array([tetra_type, hex_type])
 
106
    # Create the array of cells unambiguously.
 
107
    cell_array = tvtk.CellArray()
 
108
    cell_array.set_cells(2, cells)
 
109
    # Now create the UG.
 
110
    ug = tvtk.UnstructuredGrid(points=points)
 
111
    # Now just set the cell types and reuse the ug locations and cells.
 
112
    ug.set_cells(cell_types, offset, cell_array)
 
113
    return ug
 
114
 
 
115
def generateStructuredGrid():
 
116
    """Generates Structured Grid"""
 
117
    dims = (32, 32, 12)
 
118
    sgrid = tvtk.StructuredGrid(dimensions=(dims[1], dims[0], dims[2]))
 
119
    r = linspace(1, 10, dims[0])
 
120
    theta = linspace(0, 2*numpy.pi, dims[1])
 
121
    z = linspace(0, 5, dims[2])
 
122
    pts = generate_annulus(r, theta, z)
 
123
    sgrid.points = pts
 
124
    s = sqrt(pts[:,0]**2 + pts[:,1]**2 + pts[:,2]**2)
 
125
    sgrid.point_data.scalars = numpy.ravel(s.copy())
 
126
    sgrid.point_data.scalars.name = 'scalars'
 
127
 
 
128
    return sgrid
 
129
 
 
130
def generateUnstructuredGrid_single():
 
131
    """Generates Untructured Grid"""
 
132
    ug = single_type_ug()
 
133
    temperature = arange(0, 120, 10, 'd')
 
134
    velocity = random.randn(12, 3)
 
135
    ug.point_data.scalars = temperature
 
136
    ug.point_data.scalars.name = 'temperature'
 
137
    # Some vectors.
 
138
    ug.point_data.vectors = velocity
 
139
    ug.point_data.vectors.name = 'velocity'
 
140
    return ug
 
141
 
 
142
def generateUnstructuredGrid_mixed():
 
143
    """Generates Untructured Grid"""
 
144
    ug = mixed_type_ug()
 
145
    temperature = arange(0, 120, 10, 'd')
 
146
    velocity = random.randn(12, 3)
 
147
    ug.point_data.scalars = temperature
 
148
    ug.point_data.scalars.name = 'temperature'
 
149
    # Some vectors.
 
150
    ug.point_data.vectors = velocity
 
151
    ug.point_data.vectors.name = 'velocity'
 
152
    return ug
 
153