def tovtk(self, N, filename):
"""Dump probes to VTK file."""
is_root = comm.Get_rank() == 0
z = self.array(N=N)
if is_root:
d = self.dims
d = (d[0], d[1], d[2]) if len(d) > 2 else (d[0], d[1], 1)
grid = pyvtk.StructuredGrid(d, self.create_dense_grid())
v = pyvtk.VtkData(grid, "Probe data. Evaluations = {}".format(self.probes.number_of_evaluations()))
if self.probes.value_size() == 1:
v.point_data.append(pyvtk.Scalars(z, name="Scalar", lookup_table='default'))
elif self.probes.value_size() == 3:
v.point_data.append(pyvtk.Vectors(z, name="Vector"))
elif self.probes.value_size() == 9: # StatisticsProbes
if N == 0:
num_evals = self.probes.number_of_evaluations()
v.point_data.append(pyvtk.Vectors(z[:, :3]/num_evals, name="UMEAN"))
rs = ["uu", "vv", "ww", "uv", "uw", "vw"]
for i in range(3, 9):
v.point_data.append(pyvtk.Scalars(z[:, i]/num_evals, name=rs[i-3], lookup_table='default'))
else: # Just dump latest snapshot
v.point_data.append(pyvtk.Vectors(z[:, :3], name="U"))
else:
raise TypeError("Only vector or scalar data supported for VTK")
v.tofile(filename)
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),
points=points)
# -*- coding: utf-8 -*-
"""
Created on Wed Jun 7 14:58:44 2017
@author: Jonas Lindemann
"""
import numpy as np
import pyvtk as vtk
print("Reading from uvw.dat...")
xyzuvw = np.loadtxt('uvw.dat', skiprows=2)
print("Converting to points and vectors")
points = xyzuvw[:, 0:3].tolist()
vectors = xyzuvw[:, 3:].tolist()
pointdata = vtk.PointData(vtk.Vectors(vectors, name="vec1"),
vtk.Vectors(vectors, name="vec2"))
data = vtk.VtkData(vtk.StructuredGrid([96, 65, 48], points), pointdata)
data.tofile('uvw', 'ascii')
def _get_vtk_mesh_circ(self, dimensions, points):
self.grid = vtk.StructuredGrid(dimensions=dimensions, points=points)
def VTKgen(lat,
lon,
mask,
depth=None,
h=None,
temp=None,
salt=None,
rho=None,
dye1=None,
dye2=None,
dye3=None,
u=None,
v=None,
w=None,
seaice=None,
shelf_base=None,
shelf_thick=None,
writebottom=False,
fname='test',
dirname='VTK',
date=None,
t=0):
""" Writes ocean and ice shelf geometry and data (e.g., tracer, vel., sea-ice) into VTK files."""
import numpy as np
NY, NX = lat.shape
if not os.path.isdir(dirname):
os.system('mkdir ' + dirname)
NY, NX = lat.shape
if depth is not None:
newlat = np.resize(lat, (2, NY, NX))
newlon = np.resize(lon, (2, NY, NX))
newdepth, bottom = get_depth(h, depth, mask)
pp = f3(newlon, newlat, newdepth)
structure = pyvtk.StructuredGrid([2, NY, NX], pp)
path_to_file = str('%s/%s-bathymetry.vtk' % (dirname, fname))
if os.path.isfile(path_to_file):
print ' \n' + '==> ' + 'Bathymetry has already been written, moving on ...\n' + ''
else:
# create bottom/shape and depths
newdepth = f1(newdepth)
bottom = f1(bottom)
pointdata = pyvtk.PointData(
pyvtk.Scalars(newdepth, name='Depth'),
pyvtk.Scalars(bottom, name='Bottom9999'))
# saving the data
vtk = pyvtk.VtkData(structure, pointdata)
vtk.tofile(dirname + '/' + fname + '-bathymetry', 'binary')
if writebottom == True:
print ' \n' + '==> ' + 'Writing tracers/vel. just at the bottom layer ...\n' + ''
data = []
if temp is not None:
tmp = np.zeros((2, NY, NX))
if len(
temp.shape
) == 2: # in case the user provides 2D array with bottom data
tmp[:, :, :] = temp[:, :]
else:
tmp[:, :, :] = temp[-1, :, :]
temp = f1(tmp)
data.append("pyvtk.Scalars(temp,name='Temp')")
if salt is not None:
tmp = np.zeros((2, NY, NX))
if len(salt.shape) == 2:
tmp[:, :, :] = salt[:, :]
else:
tmp[:, :, :] = salt[-1, :, :]
salt = f1(tmp)
data.append("pyvtk.Scalars(salt,name='Salt')")
if rho is not None:
tmp = np.zeros((2, NY, NX))
if len(rho.shape) == 2:
tmp[:, :, :] = rho[:, :]
else:
tmp[:, :, :] = rho[-1, :, :]
rho = f1(tmp)
data.append("pyvtk.Scalars(rho,name='Rho')")
if dye1 is not None:
tmp = np.zeros((2, NY, NX))
if len(dye1.shape) == 2:
tmp[:, :, :] = dye1[:, :]
else:
tmp[:, :, :] = dye1[-1, :, :]
dye1 = f1(tmp)
data.append("pyvtk.Scalars(dye1,name='Dye1')")
if dye2 is not None:
tmp = np.zeros((2, NY, NX))
if len(dye2.shape) == 2:
tmp[:, :, :] = dye2[:, :]
else:
tmp[:, :, :] = dye2[-1, :, :]
dye2 = f1(tmp)
data.append("pyvtk.Scalars(dye2,name='Dye2')")
if dye3 is not None:
tmp = np.zeros((2, NY, NX))
if len(dye3.shape) == 2:
tmp[:, :, :] = dye3[:, :]
else:
tmp[:, :, :] = dye3[-1, :, :]
dye3 = f1(tmp)
data.append("pyvtk.Scalars(dye3,name='Dye3')")
if u is not None and v is not None:
w = np.zeros((2, NY, NX)) # no w vel for now
tmpu = np.zeros((2, NY, NX))
tmpv = np.zeros((2, NY, NX))
if len(u.shape) == 2:
tmpu[:, :, :] = u[:, :]
tmpv[:, :, :] = v[:, :]
else:
tmpu[:, :, :] = u[-1, :, :]
tmpv[:, :, :] = v[-1, :, :]
vel = f3(tmpu, tmpv, w)
data.append("pyvtk.Vectors(vel,name='Velocity')")
if temp is not None or salt is not None or rho is not None or u is not None:
for d in range(len(data)):
if d == 0:
tmp = data[d]
else:
tmp = tmp + ',' + data[d]
s = str("pyvtk.PointData(%s)" % (tmp))
pointdata = eval(s)
# saving the data
vtk = pyvtk.VtkData(structure, pointdata)
if date is not None:
s = str("vtk.tofile('%s/%s-%s-bottom-%05d','binary')" %
(dirname, fname, date, t))
eval(s)
else:
s = str("vtk.tofile('%s/%s-bottom-%05d','binary')" %
(dirname, fname, t))
eval(s)
if shelf_base is not None and shelf_thick is not None:
NZ_IS = 2
newlat = np.resize(lat, (NZ_IS, NY, NX))
newlon = np.resize(lon, (NZ_IS, NY, NX))
dum, z = get_iceshelf(shelf_base, shelf_thick, NZ_IS)
iceshelf = f1(dum)
pp = f3(newlon, newlat, z)
structure = pyvtk.StructuredGrid([NZ_IS, NY, NX], pp)
pointdata = pyvtk.PointData(
pyvtk.Scalars(iceshelf, name='IceShelf9999'))
vtk = pyvtk.VtkData(structure, pointdata)
if t > 0:
s = str("vtk.tofile('%s/%s-ice-shelf-%05d','binary')" %
(dirname, fname, t))
eval(s)
else:
vtk.tofile(dirname + '/' + fname + '-ice-shelf', 'binary')
if writebottom == False:
data = []
if temp is not None:
temp = f1(temp)
data.append("pyvtk.Scalars(temp,name='Temp')")
if salt is not None:
salt = f1(salt)
data.append("pyvtk.Scalars(salt,name='Salt')")
if rho is not None:
rho = f1(rho)
data.append("pyvtk.Scalars(rho,name='Rho')")
if dye1 is not None:
dye1 = f1(dye1)
data.append("pyvtk.Scalars(dye1,name='Dye1')")
if dye2 is not None:
dye2 = f1(dye2)
data.append("pyvtk.Scalars(dye2,name='Dye2')")
if dye3 is not None:
dye3 = f1(dye3)
data.append("pyvtk.Scalars(dye3,name='Dye3')")
if u is not None and v is not None:
if w is not None:
vel = f3(u, v, w)
else:
w = np.zeros(u.shape)
vel = f3(u, v, w)
data.append("pyvtk.Vectors(vel,name='Velocity')")
if seaice is not None:
NZ, NY, NX = h.shape
sice1 = np.zeros((NZ, NY, NX))
sice1[0, :, :] = seaice[:, :]
sice2 = np.zeros((NZ, NY, NX))
seaice[seaice >= 0.15] = 1.0 # all values >=15% are unit
sice2[0, :, :] = seaice[:, :]
seaice1 = f1(sice1)
seaice2 = f1(sice2)
data.append("pyvtk.Scalars(seaice1,name='Sea-ice')")
data.append("pyvtk.Scalars(seaice2,name='Sea-ice-binary')")
if temp is not None or salt is not None or rho is not None or u is not None or seaice is not None:
NZ, NY, NX = h.shape
# resize lon lat for real mesh
newlat = np.resize(lat, (NZ, NY, NX))
newlon = np.resize(lon, (NZ, NY, NX))
pp = f3(newlon, newlat, h)
structure = pyvtk.StructuredGrid([NZ, NY, NX], pp)
for d in range(len(data)):
if d == 0:
tmp = data[d]
else:
tmp = tmp + ',' + data[d]
s = str("pyvtk.PointData(%s)" % (tmp))
pointdata = eval(s)
# saving the data
vtk = pyvtk.VtkData(structure, pointdata)
if date is not None:
s = str("vtk.tofile('%s/%s-%s-%05d','binary')" %
(dirname, fname, date, t))
eval(s)
else:
s = str("vtk.tofile('%s/%s-%05d','binary')" %
(dirname, fname, t))
eval(s)