def get_vtk(data, time):
# temporary function generating vtk data structure for experiment
from tvtk.api import tvtk
tt = to_unixtime(time)
ds = data.sel(time=tt, method='nearest')
fv = ds.dist.values[0, ...]
vr = ds.vr.values[0, ...]
vt = ds.vt.values[0, ...]
vp = ds.vp.values[0, ...]
f, r, t, p = _extend_mesh_interp(fv, vr, vt, vp)
fmax = f.max()
fmin = fmax * 1.0e-15
f = np.clip(f[:-1, ...], fmin, fmax) # eliminate zeros
p = p[:-1, ...]
rr = r[None, None, :]
tt = t[None, :, None]
pp = p[:, None, None]
dims = f.shape
mesh = np.zeros((np.prod(dims), 3), dtype=np.float64)
mesh[:, 0] = (rr * np.sin(tt) * np.cos(pp)).ravel()
mesh[:, 1] = (rr * np.sin(tt) * np.sin(pp)).ravel()
mesh[:, 2] = (rr * np.cos(tt) * np.ones_like(pp)).ravel()
sgrid = tvtk.StructuredGrid(dimensions=dims[::-1])
sgrid.points = np.zeros((np.prod(dims), 3), dtype=np.float64)
sgrid.points = mesh
sgrid.point_data.scalars = np.log10(f.ravel())
sgrid.point_data.scalars.name = 'VDF'
return tvtk.to_vtk(sgrid)
def convert_to_vts(self, outdir='.', Radius=1., name='basin'):
from tvtk.api import tvtk, write_data
from sympy.ntheory import primefactors
lonArr = np.array([])
latArr = np.array([])
for geopolygon in self.geopolygons:
lonArr = np.append(lonArr, geopolygon.lonArr)
latArr = np.append(latArr, geopolygon.latArr)
theta = (90. - latArr) / 180. * np.pi
phi = lonArr / 180. * np.pi
x = Radius * np.sin(theta) * np.cos(phi)
y = Radius * np.sin(theta) * np.sin(phi)
z = Radius * np.cos(theta)
pts = np.empty(z.shape + (3, ), dtype=float)
pts[..., 0] = x
pts[..., 1] = y
pts[..., 2] = z
least_prime = primefactors(x.size)[0]
dims = (x.size / least_prime, least_prime, 1)
sgrid = tvtk.StructuredGrid(dimensions=dims, points=pts)
sgrid.point_data.scalars = (np.ones(x.size)).ravel(order='F')
sgrid.point_data.scalars.name = name
outfname = outdir + '/' + name + '.vts'
write_data(sgrid, outfname)
return
def write_to_vtk(conf, ar, filename):
"""
This function writes a numpy array into vtk format file.
It reads configuration into DispalyParams object.
The given array is cropped to the configured size if the crop parameter is defined in configuration file.
Based on configuration parameters the grid needed for the tvtk StructuredGrid is calculated.
Other StructuredGrid attributes are set.
The initialized StructuredGrid object is then written into vtk type file. If configuration parameter
save_two_files is set to True, one file will be saved for amplitudes with "_Amp.vtk" ending, and one for
phases with "_Phase.vtk" ending. Otherwise, a single file will be saved that contains double array
with amplitudes a nd phases.
Parameters
----------
conf : str
configuration file name
ar : numpy array
a complex array thatwill be saved
filename : str
a prefix to an output filename
Returns
-------
none
"""
params = DispalyParams(conf)
if params.crop is None:
dims = ar.shape
arr_cropped = ar.ravel()
else:
dims = params.crop
arr_cropped = crop_array_center(ar, params.crop).ravel()
amps = np.abs(arr_cropped)
phases = np.arctan2(arr_cropped.imag, arr_cropped.real)
geometry = params.get_geometry()
coordinates = get_coords(dims, geometry)
sg = tvtk.StructuredGrid()
sg.points = coordinates
sg.dimensions = (dims[2], dims[1], dims[0])
sg.extent = 0, dims[2] - 1, 0, dims[1] - 1, 0, dims[0] - 1
if params.save_two_files:
sg.point_data.scalars = amps
sg.point_data.scalars.name = "Amp"
write_array(sg, filename + "_Amp.vtk")
sg.point_data.scalars = phases
sg.point_data.scalars.name = "Phase"
write_array(sg, filename + "_Phase.vtk")
else:
sg.point_data.scalars = amps
sg.point_data.scalars.name = "Amp"
ph = tvtk.FloatArray()
ph.from_array(phases)
ph.name = "Phase"
sg.point_data.add_array(ph)
write_array(sg, filename + ".vtk")
def create_grid(grid, data, period=1):
s = np.shape(data)
nx = grid["nx"]#[0]
ny = grid["ny"]#[0]
nz = s[2]
print("data: %d,%d,%d grid: %d,%d\n" % (s[0],s[1],s[2], nx,ny))
dims = (nx, nz, ny)
sgrid = tvtk.StructuredGrid(dimensions=dims)
pts = aligned_points(grid, nz, period)
print(np.shape(pts))
sgrid.points = pts
scalar = np.zeros([nx*ny*nz])
start = 0
for y in range(ny):
end = start + nx*nz
#scalar[start:end] = (data[:,y,:]).transpose().ravel()
scalar[start:end] = (data[:,y,:]).ravel()
print(y, " = " , np.max(scalar[start:end]))
start = end
sgrid.point_data.scalars = np.ravel(scalar.copy())
sgrid.point_data.scalars.name = "data"
return sgrid
def create_grid(grid, data, period=1):
"""
Create a structured grid which can be plotted with Mayavi
or saved to a VTK file.
Example
-------
from boutdata.collect import collect
from boututils.file_import import file_import
from boututile import boutgrid
# Load grid file
g = file_import("bout.grd.nc")
# Load 3D data (x,y,z)
data = collect("P", tind=-1)[0,:,:,:]
# Create a structured grid
sgrid = boutgrid.create_grid(g, data, 1)
# Write structured grid to file
w = tvtk.XMLStructuredGridWriter(input=sgrid, file_name='sgrid.vts')
w.write()
# View the structured grid
boutgrid.view3d(sgrid)
"""
s = np.shape(data)
nx = grid["nx"] # [0]
ny = grid["ny"] # [0]
nz = s[2]
print("data: %d,%d,%d grid: %d,%d\n" % (s[0], s[1], s[2], nx, ny))
dims = (nx, nz, ny)
sgrid = tvtk.StructuredGrid(dimensions=dims)
pts = aligned_points(grid, nz, period)
print(np.shape(pts))
sgrid.points = pts
scalar = np.zeros([nx * ny * nz])
start = 0
for y in range(ny):
end = start + nx * nz
# scalar[start:end] = (data[:,y,:]).transpose().ravel()
scalar[start:end] = (data[:, y, :]).ravel()
print(y, " = ", np.max(scalar[start:end]))
start = end
sgrid.point_data.scalars = np.ravel(scalar.copy())
sgrid.point_data.scalars.name = "data"
return sgrid
def _mk_structured_grid(self):
""" Creates a StructuredGrid VTK data set using the factory's
attributes.
"""
# FIXME: We need to figure out the dimensions of the data
# here, if any.
sg = tvtk.StructuredGrid(dimensions=self.scalar_data.shape)
sg.points = c_[self.position_x.ravel(),
self.position_y.ravel(),
self.position_z.ravel(), ]
self._vtk_source = sg
self._mayavi_source = VTKDataSource(data=self._vtk_source)
def setUp(self):
datasets = [tvtk.ImageData(),
tvtk.StructuredPoints(),
tvtk.RectilinearGrid(),
tvtk.StructuredGrid(),
tvtk.PolyData(),
tvtk.UnstructuredGrid(),
]
exts = ['.vti', '.vti', '.vtr', '.vts', '.vtp', '.vtu']
self.datasets = datasets
self.exts = exts
def structured_grid():
# Make the data.
dims = (3, 4, 3)
r = linspace(5, 15, dims[0])
theta = linspace(0, 0.5 * pi, dims[1])
z = linspace(0, 10, dims[2])
pts = generate_annulus(r, theta, z)
sgrid = tvtk.StructuredGrid(dimensions=(dims[1], dims[0], dims[2]))
sgrid.points = pts
s = random.random((dims[0] * dims[1] * dims[2]))
sgrid.point_data.scalars = ravel(s.copy())
sgrid.point_data.scalars.name = 'scalars'
return sgrid
def plot(atoms, data, contours):
"""Plot atoms, unit-cell and iso-surfaces using Mayavi.
Parameters:
atoms: Atoms object
Positions, atomiz numbers and unit-cell.
data: 3-d ndarray of float
Data for iso-surfaces.
countours: list of float
Contour values.
"""
# Delay slow imports:
from mayavi import mlab
from tvtk.api import tvtk
import mayavi.tools.pipeline
mlab.figure(1, bgcolor=(1, 1, 1)) # make a white figure
# Plot the atoms as spheres:
for pos, Z in zip(atoms.positions, atoms.numbers):
mlab.points3d(*pos,
scale_factor=covalent_radii[Z],
resolution=20,
color=tuple(cpk_colors[Z]))
# Draw the unit cell:
A = atoms.cell
for i1, a in enumerate(A):
i2 = (i1 + 1) % 3
i3 = (i1 + 2) % 3
for b in [np.zeros(3), A[i2]]:
for c in [np.zeros(3), A[i3]]:
p1 = b + c
p2 = p1 + a
mlab.plot3d([p1[0], p2[0]], [p1[1], p2[1]], [p1[2], p2[2]],
tube_radius=0.1)
pts = np.dot(np.indices(data.shape, float).T / data.shape,
atoms.cell).T.reshape((3, -1)).T
sgrid = tvtk.StructuredGrid(dimensions=data.shape)
sgrid.points = pts
sgrid.point_data.scalars = np.ravel(data.copy())
sgrid.point_data.scalars.name = 'scalars'
mayavi.tools.pipeline.iso_surface(sgrid,
contours=contours,
transparent=True,
opacity=0.5,
colormap='hot')
mlab.show()
def generateStructuredGrid():
"""Generates Structured Grid"""
dims = (32, 32, 12)
sgrid = tvtk.StructuredGrid(dimensions=(dims[1], dims[0], dims[2]))
r = linspace(1, 10, dims[0])
theta = linspace(0, 2 * numpy.pi, dims[1])
z = linspace(0, 5, dims[2])
pts = generate_annulus(r, theta, z)
sgrid.points = pts
s = sqrt(pts[:, 0]**2 + pts[:, 1]**2 + pts[:, 2]**2)
sgrid.point_data.scalars = numpy.ravel(s.copy())
sgrid.point_data.scalars.name = 'scalars'
return sgrid
def get_rs_structured_grid(self, **args):
sg = tvtk.StructuredGrid()
arr0 = self.recip_arrs[list(self.recip_arrs.keys())[0]]
dims = list(arr0.shape)
sg.points = self.recip_coords
for a in self.recip_arrs.keys():
arr = tvtk.DoubleArray()
arr.from_array(self.recip_arrs[a].ravel())
arr.name = a
sg.point_data.add_array(arr)
sg.dimensions = (dims[2], dims[1], dims[0])
sg.extent = 0, dims[2] - 1, 0, dims[1] - 1, 0, dims[0] - 1
return sg
def convert_to_vts(self, outdir, component, iter0=None, iterf=None, diter=None, Radius=0.98, verbose=True):
"""
Plot depth slices of field component at given depth ranging between "valmin" and "valmax"
================================================================================================
Input parameters:
outdir - output directory
component - component for plotting
The currently available "components" are:
Material parameters: A, B, C, mu, lambda, rhoinv, vp, vsh, vsv, rho
Velocity field snapshots: vx, vy, vz
Sensitivity kernels: Q_mu, Q_kappa, alpha_mu, alpha_kappa
depth - depth for plot (km)
iter0, iterf - start/end iteration index
diter - iteration interval
Radius - radius for output sphere
=================================================================================================
"""
if not os.path.isdir(outdir): os.makedirs(outdir)
group=self[component]
from tvtk.api import tvtk, write_data
try: iterArr=np.arange(iter0 ,iterf+diter, diter, dtype=int)
except: iterArr = group.keys()
least_prime=None
for iteration in iterArr:
subgroup=group[str(iteration)]
if len(subgroup.keys())==0: continue
if verbose: print 'Output vts file for iteration =',iteration
theta=np.array([]); phi=np.array([]); r=np.array([]); field = np.array([])
for key in subgroup.keys():
subdset = subgroup[key]
field = np.append(field, (subdset[...]))
theta1 = subdset.attrs['theta']
phi1 = subdset.attrs['phi']
theta1, phi1 = np.meshgrid(theta1, phi1, indexing='ij')
theta = np.append(theta, theta1)
phi = np.append(phi, phi1)
x = Radius * np.sin(theta) * np.cos(phi)
y = Radius * np.sin(theta) * np.sin(phi)
z = Radius * np.cos(theta)
if least_prime==None: least_prime=primefactors(field.size)[0]
dims = (field.size/least_prime, least_prime, 1)
pts = np.empty(z.shape + (3,), dtype=float)
pts[..., 0] = x; pts[..., 1] = y; pts[..., 2] = z
sgrid = tvtk.StructuredGrid(dimensions=dims, points=pts)
sgrid.point_data.scalars = (field).ravel(order='F')
sgrid.point_data.scalars.name = component
outfname=outdir+'/'+component+'_'+str(iteration)+'.vts'
write_data(sgrid, outfname)
return
def __generate_grid(self, mins, maxs, cells, datas, names):
''' Generates a grid from given data
:param mins: An array of minimum coordinates for the grid for ex. [-100, 0, 0]
:param maxs: An array of maximum coordinates for the grid for ex. [-100, 0, 0]
:param cells: An array of number of cells in x, y, z direction
:param datas: Scalar data for the grid e.g. array([ cell1Rho, cell2Rho, cell3Rho, cell4Rho, .., cellNRho ])
:param names: Name for the scalar data
'''
# Create nodes
x, y, z = mgrid[mins[0]:maxs[0]:(cells[0] + 1) * complex(0, 1),
mins[1]:maxs[1]:(cells[1] + 1) * complex(0, 1),
mins[2]:maxs[2]:(cells[2] + 1) * complex(0, 1)]
# Create points for the nodes:
pts = empty(z.shape + (3, ), dtype=float)
pts[..., 0] = x
pts[..., 1] = y
pts[..., 2] = z
# Input scalars
scalars = np.array(datas)
# We reorder the points, scalars and vectors so this is as per VTK's
# requirement of x first, y next and z last.
pts = pts.transpose(2, 1, 0, 3).copy()
pts.shape = pts.size / 3, 3
scalars = scalars.T.copy()
# Create the dataset.
sg = tvtk.StructuredGrid(dimensions=x.shape, points=pts)
sg.cell_data.scalars = ravel(scalars.copy())
if isinstance(names, str) == False:
names = "custom"
sg.cell_data.scalars.name = names
# Visualize the data
d = self.scene.mlab.pipeline.add_dataset(sg)
iso = self.scene.mlab.pipeline.surface(d)
# Add labels:
# from mayavi.modules.labels import Labels
# testlabels = self.scene.mlab.pipeline.labels(d)
self.dataset = d
print scalars[0]
# Configure traits
self.configure_traits()
self.__grid_figure = mayavi.mlab.gcf(engine=self.__engine)
def exportToVTK( self, spect_filter=None, phot_num=False,\
lambda0_um = None, smooth_filter=None, \
filename='spectrum', project=False):
if not tvtk_installed:
print('TVTK API is not found')
return
omega, theta, phi = self.Args['omega'], self.Args['theta'], \
self.Args['phi']
phi = np.r_[phi, 2*np.pi]
if project is False:
val = self.get_full_spectrum(spect_filter=spect_filter, \
phot_num=phot_num, lambda0_um=lambda0_um)
scalar_name = 'spectrum'
else:
val = self.get_spot( phot_num=phot_num, lambda0_um=lambda0_um, \
spect_filter=spect_filter)
val = val[None, :, :]
omega = omega[[-1]]
filename += '_proj'
scalar_name = 'spectrum_proj'
val = np.concatenate( (val, val[:, :, [0]]), axis=-1 )\
.astype(self.dtype)
if smooth_filter is not None:
val = gaussian_filter(val, smooth_filter)
Nom = omega.size
Nth = theta.size
Nph = phi.size
omega = omega[:,None,None].astype(self.dtype)
theta = theta[None,:, None].astype(self.dtype)
phi = phi[None,None,:].astype(self.dtype)
x, y, z = (omega*np.sin(theta)*np.sin(phi)).ravel(), \
(omega*np.sin(theta)*np.cos(phi)).ravel(), \
(omega*np.cos(theta)*np.ones_like(phi)).ravel()
spc_vtk = tvtk.StructuredGrid(dimensions=(Nph, Nth, Nom),
points=np.vstack((x.ravel(),y.ravel(),z.ravel())).T)
spc_vtk.point_data.scalars = val.flatten()
spc_vtk.point_data.scalars.name = scalar_name
write_data(spc_vtk, filename)
def mesh_to_vtkStructuredGrid(X_mesh, Y_mesh, Z_mesh, vals_mesh):
# Points in the format required by VTK.
pts = np.empty(Z_mesh.shape + (3,), dtype=float)
pts[..., 0] = X_mesh
pts[..., 1] = Y_mesh
pts[..., 2] = Z_mesh
# We reorder the points, scalars and vectors so this is as per VTK's
# requirement of x first, y next and z last.
pts = pts.transpose(2, 1, 0, 3).copy()
pts = pts.reshape(int(pts.size / 3), 3)
scalars = vals_mesh.transpose(2, 0, 1).copy()
sg = tvtk.StructuredGrid(dimensions=X_mesh.shape, points=pts)
sg.point_data.scalars = scalars.ravel()
sg.point_data.scalars.name = 'values'
return sg
def __init__(self, mesh, directory=".", filename="unnamed"):
self.mesh = mesh
self.directory = directory
self.filename = filename
if isinstance(mesh, HexagonalMesh):
self.grid = tvtk.PolyData(points=mesh.vertices,
polys=mesh.hexagons)
elif isinstance(mesh, CuboidMesh):
# if the mesh is made up of nx * ny * nz cells, it has
# (nx + 1) * (ny + 1) * (nz + 1) vertices.
self.grid = tvtk.StructuredGrid(
dimensions=[ni + 1 for ni in mesh.size])
self.grid.points = mesh.grid
else:
raise NotImplementedError(
"Mesh should be CuboidMesh or HexagonalMesh.")
def test_tvtk_dataset_name(self):
"Can tvtk datasets can be converted to names correctly."
datasets = [
tvtk.ImageData(),
tvtk.StructuredPoints(),
tvtk.RectilinearGrid(),
tvtk.StructuredGrid(),
tvtk.PolyData(),
tvtk.UnstructuredGrid(),
tvtk.Property(), # Not a dataset!
'foo', # Not a TVTK object.
]
expect = [
'image_data', 'image_data', 'rectilinear_grid', 'structured_grid',
'poly_data', 'unstructured_grid', 'none', 'none'
]
result = [pipeline_info.get_tvtk_dataset_name(d) for d in datasets]
self.assertEqual(result, expect)
def save_points_vtk(result, outfile, **kwds):
'''
Save a points result to a VTK file.
'''
from tvtk.api import tvtk, write_data
if "indices" in result.array_data_by_type():
return _save_meshlike_points_vtk(result, outfile, **kwds)
for typ, nam, arr in result.triplets():
print typ, nam, arr.shape
dim = arr.shape[:-1]
points = arr.reshape(numpy.product(dim), arr.shape[-1])
pd = tvtk.StructuredGrid(dimensions=(1, dim[0], dim[1]), points=points)
n, e = osp.splitext(outfile)
fname = n + "-" + nam + e
print "writing %s" % fname
write_data(pd, fname)
def py2vtk(dat, t, name, ext):
nx, ny, nz = dat.shape[1:4]
# Generate the grid,points
xx, yy, zz = np.mgrid[0:nx, 0:ny, 0:nz]
pts = np.empty((nx, ny, nz, 3), dtype=int)
pts[..., 0] = xx
pts[..., 1] = yy
pts[..., 2] = zz
# We reorder the points and vectors so this is as per VTK's
# requirement of x first, y next and z last.
pts = pts.transpose(2, 1, 0, 3).copy()
pts.shape = pts.size // 3, 3
vectors = dat[..., t].transpose(3, 2, 1, 0).copy()
vectors.shape = vectors.size // 3, 3
sg = tvtk.StructuredGrid(dimensions=xx.shape, points=pts)
sg.point_data.vectors = vectors
sg.point_data.vectors.name = name
write_data(sg, name + ext)
def save_grid_to_disk(grid, file_path):
"""
save a VTK structured grid to a .vtk file
:param grid: grid to save
:param file_path: path to use
"""
grid = grid.cpu().numpy()
x, y, z = grid[0, ...], grid[1, ...], grid[2, ...]
pts = np.empty(x.shape + (3, ), dtype=float)
pts[..., 0], pts[..., 1], pts[..., 2] = x, y, z
pts = pts.transpose(2, 1, 0, 3).copy()
pts.shape = pts.size // 3, 3
sg = tvtk.StructuredGrid(dimensions=x.shape, points=pts)
write_data(sg, file_path)
def test6(self):
def generate_annulus(r, theta, z):
""" Generate points for structured grid for a cylindrical annular
volume. This method is useful for generating a unstructured
cylindrical mesh for VTK.
"""
# Find the x values and y values for each plane.
x_plane = (np.cos(theta) * r[:, None]).ravel()
y_plane = (np.sin(theta) * r[:, None]).ravel()
# Allocate an array for all the points. We'll have len(x_plane)
# points on each plane, and we have a plane for each z value, so
# we need len(x_plane)*len(z) points.
points = np.empty([len(x_plane) * len(z), 3])
# Loop through the points for each plane and fill them with the
# correct x,y,z values.
start = 0
for z_plane in z:
end = start + len(x_plane)
# slice out a plane of the output points and fill it
# with the x,y, and z values for this plane. The x,y
# values are the same for every plane. The z value
# is set to the current z
plane_points = points[start:end]
plane_points[:, 0] = x_plane
plane_points[:, 1] = y_plane
plane_points[:, 2] = z_plane
start = end
return points
dims = (3, 4, 3)
r = np.linspace(5, 15, dims[0])
theta = np.linspace(0, 0.5 * np.pi, dims[1])
z = np.linspace(0, 10, dims[2])
pts = generate_annulus(r, theta, z)
sgrid = tvtk.StructuredGrid(dimensions=(dims[1], dims[0], dims[2]))
sgrid.points = pts
s = np.random.random((dims[0] * dims[1] * dims[2]))
sgrid.point_data.scalars = np.ravel(s.copy())
sgrid.point_data.scalars.name = 'scalars'
def scalarDataOnRegularGridActor(
points, data, dimensions,
colorRange=None,transformation='real'):
"""
Return a TVTK actor representing the plot of a function interpolated on
a regular grid.
"""
if points.shape[0] != 3 or points.ndim != 2:
raise ValueError("incorrect shape")
data = data.squeeze()
if data.ndim != 1:
raise ValueError("incorrect shape")
if not hasattr(transformation, '__call__'):
if transformation=='real':
data_transform = lambda point,val:np.real(val)
elif transformation=='imag':
data_transform = lambda point,val:np.imag(val)
elif transformation=='abs':
data_transform = lambda point,val:np.abs(val)
else:
raise ValueError("Unknown value for 'transformation'. It needs to be 'real', 'imag', 'abs' or a Python Callable!")
else:
data_transform = transformation
data = data_transform(points,data)
dims = dimensions
if colorRange is None:
minVal = np.min(data)
maxVal = np.max(data)
colorRange = (minVal, maxVal)
g = tvtk.StructuredGrid(dimensions=(dims[1], dims[0], 1), points=points.T)
g.point_data.scalars = data
# Create actors
mapper = tvtk.DataSetMapper(input=g)
mapper.scalar_range = colorRange
return tvtk.Actor(mapper=mapper)
def _make_vtk_mesh_circ(self):
"""
Create a cylindric mesh using vtk.StructuredGrid method.
Note:
this function operates only once, and all following calls
"""
# Exit if the grid already exists and CommonMesh is True
if self.grid is not None and self.CommonMesh:
return
# register the z-origin of the grid
self.zmin_orig = self.info.zmin
# Get the Z and R axes from the original data
z, r = self.info.z, self.info.r
r = r[r.size // 2:]
Nr, Nz = r.size, z.size
# Make Theta axis (half)
theta = np.r_[0:2 * np.pi:(self.Nth + 1) * 1j]
# shift the visualization domain origin if needed
if self.zmin_fixed is not None:
z -= z.min() - self.zmin_fixed
# Make the grid-point (copied from TVTK tutorial)
points = np.empty([len(theta) * len(r) * len(z), 3])
x_plane = (np.cos(theta) * r[:, None]).ravel()
y_plane = (np.sin(theta) * r[:, None]).ravel()
start = 0
for iz in range(len(z)):
end = start + len(x_plane)
z_plane = z[iz]
points[start:end, 0] = x_plane
points[start:end, 1] = y_plane
points[start:end, 2] = z_plane
start = end
# register the grid VTK container
self.grid = vtk.StructuredGrid(dimensions=(self.Nth + 1, Nr, Nz))
self.grid.points = points
''' points = np.column_stack((x,y,z)) #triangles = array([[0,1,3], [0,3,2], [1,2,3], [0,2,1]]) #temperature = array([10., 20., 30., 40.]) # The TVTK dataset. mesh = tvtk.PolyData(points=points)#, polys=triangles) mesh.point_data.scalars = elev mesh.point_data.scalars.name = 'Elevation' # Uncomment the next two lines to save the dataset to a VTK XML file. w = tvtk.XMLPolyDataWriter(input=mesh, file_name='polydata.vtp') w.write() ''' #r = rect_grid() data = np.atleast_3d(elev[10,...]) r = tvtk.StructuredGrid() r.dimensions = (data.shape[0]+1, data.shape[1]+1, data.shape[2]) r.points = np.column_stack((x,y,z)) r.cell_data.scalars = data.ravel() r.cell_data.scalars.name = 'scalars' w2 = tvtk.XMLStructuredGridWriter(input=r, file_name='structgrid.vts') w2.write() f2.close() f.close()
def plot_slice(path,file_name,f,var_num,num_eqn,minval,maxval):
# plot one slice
patch_specs_list,patch_array_list = slice_io.read_patch_list(path,file_name)
normal,m_slice,num_t_tick,translate = slice_io.q_output_name_read(file_name,all_data_dict)
orient_dict = {'x':[1,2,0],'y':[0,2,1], 'z':[0,1,2]}
orient_ord = {'x':0, 'y':1, 'z':2}
orient_plane = {'x':'yz', 'y':'xz', 'z':'xy'}
orient_real = orient_dict[normal]
src_list = []
pts_list = []
vprint('\t(plot_slice) working on: ' + file_name)
npatches = len(patch_specs_list)
for k,patch_specs in enumerate(patch_specs_list):
# load patch specs
m1 = int(patch_specs[1])
m2 = int(patch_specs[2])
ulo = float(patch_specs[3])
vlo = float(patch_specs[4])
d1 = float(patch_specs[5])
d2 = float(patch_specs[6])
# debugging outputs
vprint('\t(plot_slice) patch number ' + str(k))
vprint('\t(plot_slice) ' + str(ulo) + ', ' + str(vlo) + '---' \
+ str(ulo + d1*m1) + ', ' + str(vlo + d2*m2))
vprint('\t(plot_slice) ' + str(d1) + 'x ' + str(d2))
# center cells
#xi1lo_c = xi1lo #+ d1/2 # _c for "center"
#xi1hi_c = xi1lo_c + d1*m1
#xi2lo_c = xi2lo #+ d2/2
#xi2hi_c = xi2lo_c + d2*m2
translate_lo = translate - 5e-5
translate_hi = translate + 5e-5
translate_lo2 = translate - 8e-5
translate_hi2 = translate + 8e-5
u = np.linspace(ulo,ulo + m1*d1,m1)
v = np.linspace(vlo,vlo + m2*d2,m2)
tr = np.linspace(translate_lo,translate_hi,2)
grids = permute_orientation(orient_real,[u,v,tr])
m_real = permute_orientation(orient_real,[m1,m2,1])
x,y,z = np.meshgrid(grids[0],grids[1],grids[2],indexing='ij')
q_sol = patch_array_list[k][var_num::num_eqn].reshape(m_real[0],m_real[1],m_real[2],order='F')
q_sol = q_sol.repeat(2,axis=orient_ord[normal])
objname = orient_plane[normal] + '_' + str(translate) + '_' + str(k)
src_list.append(mlab.pipeline.scalar_field(x,y,z,q_sol,name = objname))
#sg = tvtk.StructuredGrid(dimensions=x.shape, points=pts_list[k])
#tr2 = np.linspace(translate_lo2,translate_hi2,2)
for tr2val in [translate_lo2, translate_hi2]:
tr2 = np.array(tr2val)
u1 = np.linspace(ulo,ulo + m1*d1,m1 + 1)
v1 = np.linspace(vlo,vlo + m2*d2,m2 + 1)
grids = permute_orientation(orient_real,[u1,v1,tr2])
m_real = permute_orientation(orient_real,[m1,m2,1])
x,y,z = np.meshgrid(grids[0],grids[1],grids[2],indexing='ij')
pts = np.empty(z.shape + (3,), dtype=float)
pts[..., 0] = x
pts[..., 1] = y
pts[..., 2] = z
pts = pts.transpose(2, 1, 0, 3).copy()
pts.shape = pts.size / 3, 3
sg = tvtk.StructuredGrid(dimensions=x.shape, points=pts)
d = mlab.pipeline.add_dataset(sg)
g1 = mlab.pipeline.grid_plane(d, line_width=0.25,color=(0,0,0))
g1.grid_plane.axis = normal
color_choice = 'Spectral'
axis_str = normal + '_axes'
for k,src in enumerate(src_list):
# plot two slices according to adaptive mesh level
# patch one
yp = mlab.pipeline.scalar_cut_plane(src, plane_orientation=axis_str,opacity=1.0,figure=f,vmin=minval,vmax=maxval,colormap=color_choice)
#yp = mlab.pipeline.image_plane_widget(src, plane_orientation=axis_str,opacity=0.5,figure=f,vmin=minval,vmax=maxval,colormap=color_choice)
tr_vec = np.zeros(3)
tr_vec[orient_ord[normal]] = translate - 5e-5*k/npatches
yp.implicit_plane.origin = (tr_vec[0],tr_vec[1],tr_vec[2])
yp.implicit_plane.normal = np.array([.5,.5,.0])
yp.implicit_plane.visible = False
#mlab.outline(line_width=0.025)
# patch two
yp = mlab.pipeline.scalar_cut_plane(src, plane_orientation=axis_str,opacity=1.0,figure=f,vmin=minval,vmax=maxval,colormap=color_choice)
#yp = mlab.pipeline.image_plane_widget(src, plane_orientation=axis_str,opacity=1.0,figure=f,vmin=minval,vmax=maxval,colormap=color_choice)
tr_vec[orient_ord[normal]] = translate + 5e-5*k/npatches
yp.implicit_plane.origin = (tr_vec[0],tr_vec[1],tr_vec[2])
yp.implicit_plane.visible = False
#mlab.outline(line_width=0.025)
import pdb
pdb.set_trace()
return
def _get_vtk_mesh_circ(self, dimensions, points):
self.grid = vtk.StructuredGrid(dimensions=dimensions)
self.grid.points = points
# Loop through the points for each plane and fill them with the
# correct x,y,z values.
start = 0
for z_plane in z:
end = start + len(x_plane)
# slice out a plane of the output points and fill it
# with the x,y, and z values for this plane. The x,y
# values are the same for every plane. The z value
# is set to the current z
plane_points = points[start:end]
plane_points[:, 0] = x_plane
plane_points[:, 1] = y_plane
plane_points[:, 2] = z_plane
start = end
return points
dims = (3, 4, 3)
r = linspace(5, 15, dims[0])
theta = linspace(0, 0.5 * pi, dims[1])
z = linspace(0, 10, dims[2])
pts = generate_annulus(r, theta, z)
sgrid = tvtk.StructuredGrid(dimensions=(dims[1], dims[0], dims[2]))
sgrid.points = pts
s = random.random((dims[0] * dims[1] * dims[2]))
sgrid.point_data.scalars = ravel(s.copy())
sgrid.point_data.scalars.name = 'scalars'
outfname = './test_circle.vtk'
write_data(sgrid, outfname)
def _pd_default(self):
return tvtk.StructuredGrid()
plane_points[:,2] = z_plane
start = end
return points
# Make the data.
dims = (51, 25, 25)
# Note here that the 'x' axis corresponds to 'theta'
theta = np.linspace(0, 2*np.pi, dims[0])
# 'y' corresponds to varying 'r'
r = np.linspace(1, 10, dims[1])
z = np.linspace(0, 5, dims[2])
pts = generate_annulus(r, theta, z)
# Uncomment the following if you want to add some noise to the data.
#pts += np.random.randn(dims[0]*dims[1]*dims[2], 3)*0.04
sgrid = tvtk.StructuredGrid(dimensions=dims)
sgrid.points = pts
s = np.sqrt(pts[:,0]**2 + pts[:,1]**2 + pts[:,2]**2)
sgrid.point_data.scalars = np.ravel(s.copy())
sgrid.point_data.scalars.name = 'scalars'
# Uncomment the next two lines to save the dataset to a VTK XML file.
#w = tvtk.XMLStructuredGridWriter(input=sgrid, file_name='sgrid.vts')
#w.write()
# View the data.
@mayavi2.standalone
def view():
from mayavi.sources.vtk_data_source import VTKDataSource
from mayavi.modules.api import Outline, GridPlane
# Get the data
datax = np.array(f[Bx])
datay = np.array(f[By])
dataz = np.array(f[Bz])
# Generate some points.
x, y, z = mgrid[0:320, 0:320, 0:320]
# The actual points.
pts = empty(z.shape + (3, ), dtype=np.int16)
pts[..., 0] = x
pts[..., 1] = y
pts[..., 2] = z
# Some vectors
vectors = empty(z.shape + (3, ), dtype=np.float32)
vectors[..., 0] = datax
vectors[..., 1] = datay
vectors[..., 2] = dataz
#Transposing because the order is z,y and x
pts = pts.transpose(2, 1, 0, 3).copy()
pts.shape = int(pts.size / 3), 3 #Flattening
vectors = vectors.transpose(2, 1, 0, 3).copy()
vectors.shape = int(vectors.size / 3), 3 #Flattening
sg = tvtk.StructuredGrid(dimensions=x.shape, points=pts)
sg.point_data.vectors = vectors
sg.point_data.vectors.name = 'vector_field'
write_data(sg, 'magField_vectors.vtk')
