183 lines
6.6 KiB
Python
183 lines
6.6 KiB
Python
def add_domain_bounds(plotter,bounds,color='black',line_width=2.):
|
|
import pyvista
|
|
domain = pyvista.Box(bounds=bounds)
|
|
plotter.add_mesh(domain,color=color,style='wireframe',line_width=line_width,opacity=0.999,lighting=False)
|
|
return
|
|
|
|
def enable_shadows_hacked(pl):
|
|
# Reimplements pyvistas "enable_shadows()" method to also render translucent
|
|
# objects (without shadows). Can be used to add objects which do not throw
|
|
# shadows (but may look weird) by setting their opacity to a value close to 1,
|
|
# e.g. 0.999. This is useful when drawing a bounding box using "add_domain_bounds()".
|
|
import vtk
|
|
shadows = vtk.vtkShadowMapPass()
|
|
transl = vtk.vtkTranslucentPass()
|
|
seq = vtk.vtkSequencePass()
|
|
passes = vtk.vtkRenderPassCollection()
|
|
passes.AddItem(shadows.GetShadowMapBakerPass())
|
|
passes.AddItem(shadows)
|
|
passes.AddItem(transl)
|
|
seq.SetPasses(passes)
|
|
# Tell the renderer to use our render pass pipeline
|
|
cameraP = vtk.vtkCameraPass()
|
|
cameraP.SetDelegatePass(seq)
|
|
pl.renderer.SetPass(cameraP)
|
|
return
|
|
|
|
def camera_from_view(pl,bounds,viewvec,
|
|
focus=(0.5,0.5,0.5),
|
|
viewup=(0,1,0),
|
|
is_focus_relative=True):
|
|
import numpy as np
|
|
focal_point = get_focal_point(bounds,focus,is_focus_relative)
|
|
pl.view_vector(-np.array(viewvec),viewup)
|
|
pl.set_focus(focal_point)
|
|
pl.reset_camera_clipping_range()
|
|
return
|
|
|
|
def camera_from_distance(pl,bounds,viewvec,dist,
|
|
focus=(0.5,0.5,0.5),
|
|
viewup=(0,1,0),
|
|
zoom=1.0,
|
|
is_focus_relative=True):
|
|
import numpy as np
|
|
focal_point = get_focal_point(bounds,focus,is_focus_relative)
|
|
position = focal_point-viewvec/np.linalg.norm(viewvec)*dist
|
|
return camera_from_position(pl,
|
|
bounds,position,
|
|
focus=focal_point,
|
|
viewup=viewup,
|
|
zoom=zoom,
|
|
is_focus_relative=False)
|
|
|
|
def camera_from_position(pl,bounds,position,
|
|
focus=(0.5,0.5,0.5),
|
|
viewup=(0,1,0),
|
|
zoom=1.0,
|
|
is_focus_relative=True):
|
|
focal_point = get_focal_point(bounds,focus,is_focus_relative)
|
|
pl.set_focus(focal_point)
|
|
pl.set_position(position)
|
|
pl.set_viewup(viewup)
|
|
pl.camera.zoom(zoom)
|
|
pl.reset_camera_clipping_range()
|
|
return
|
|
|
|
def get_focal_point(bounds,focus,is_focus_relative=True):
|
|
import numpy as np
|
|
if is_focus_relative:
|
|
focal_point = np.array(
|
|
(focus[0]*(bounds[1]+bounds[0]),
|
|
focus[1]*(bounds[3]+bounds[2]),
|
|
focus[2]*(bounds[5]+bounds[4])))
|
|
else:
|
|
focal_point = np.array(focus)
|
|
return focal_point
|
|
|
|
def chunk_to_pvmesh(chunk,gridg):
|
|
import pyvista
|
|
mesh = pyvista.UniformGrid()
|
|
mesh.dimensions = (
|
|
chunk['nxl']+2*chunk['ighost'],
|
|
chunk['nyl']+2*chunk['ighost'],
|
|
chunk['nzl']+2*chunk['ighost']
|
|
)
|
|
xg,yg,zg = gridg
|
|
dx,dy,dz = xg[2]-xg[1],yg[2]-yg[1],zg[2]-zg[1]
|
|
x0 = xg[chunk['ibeg']]-chunk['ighost']*dx
|
|
y0 = yg[chunk['jbeg']]-chunk['ighost']*dy
|
|
z0 = zg[chunk['kbeg']]-chunk['ighost']*dz
|
|
mesh.origin = (x0,y0,z0) # The bottom left corner of the data set
|
|
mesh.spacing = (dx,dy,dz) # These are the cell sizes along each axis
|
|
mesh.point_arrays['values'] = chunk['data'].flatten(order='F') # Flatten the array!
|
|
return mesh
|
|
|
|
def clip_out_of_bounds(mesh,bounds,axis=0,inplace=True,clip_lower=True,clip_upper=True):
|
|
assert axis<3, "axis must be in [0,1,2]."
|
|
new_bounds = bounds[2*axis:2*axis+2]
|
|
old_bounds = mesh.bounds[2*axis:2*axis+2]
|
|
#print('DEBUG: New bounds: ',new_bounds)
|
|
#print('DEBUG: Old bounds: ',old_bounds)
|
|
if new_bounds[0]>old_bounds[0] and clip_lower:
|
|
#print('DEBUG: Clipping lower',new_bounds[0],old_bounds[0])
|
|
normal = [0,0,0]
|
|
normal[axis] = -1
|
|
origin = [0,0,0]
|
|
origin[axis] = new_bounds[0]
|
|
if inplace:
|
|
mesh.clip(normal=normal,origin=origin,inplace=True)
|
|
else:
|
|
mesh = mesh.clip(normal=normal,origin=origin,inplace=False)
|
|
if new_bounds[1]<old_bounds[1] and clip_upper:
|
|
#print('DEBUG: Clipping upper',new_bounds[1],old_bounds[1])
|
|
normal = [0,0,0]
|
|
normal[axis] = 1
|
|
origin = [0,0,0]
|
|
origin[axis] = new_bounds[1]
|
|
if inplace:
|
|
mesh.clip(normal=normal,origin=origin,inplace=True)
|
|
else:
|
|
mesh = mesh.clip(normal=normal,origin=origin,inplace=False)
|
|
if inplace:
|
|
return None
|
|
else:
|
|
return mesh
|
|
|
|
def translate_circular(pd,translation,bounds,axis=0):
|
|
'''Translates pyvista PolyData objects while taking into account
|
|
the bounding box'''
|
|
assert(axis<3)
|
|
import pyvista
|
|
# Map translation onto [0,L]
|
|
L = bounds[2*axis+1]
|
|
translation = translation%L
|
|
# Assemble normal direction and origin of cut, as well as translation vector
|
|
shift_forw = [0.,0.,0.]
|
|
shift_back = [0.,0.,0.]
|
|
cut_origin = [0.,0.,0.]
|
|
cut_normal = [0,0,0]
|
|
shift_forw[axis] = translation
|
|
shift_back[axis] = translation-L
|
|
cut_origin[axis] = L-translation
|
|
cut_normal[axis] = 1
|
|
# Split PolyData at new L after shift, then translate to new position
|
|
# The clipping is shallow, e.g. the call to translate affects both clipped
|
|
# and unclipped meshes. Thus create a deep copy first before translating.
|
|
pd_lo,pd_hi = pd.clip(normal=cut_normal,origin=cut_origin,return_clipped=True)
|
|
pd_lo = pyvista.PolyData(pd_lo,deep=True)
|
|
pd_hi = pyvista.PolyData(pd_hi,deep=True)
|
|
pd_lo.translate(shift_forw)
|
|
pd_hi.translate(shift_back)
|
|
# return the merged PolyData
|
|
return pd_lo+pd_hi
|
|
|
|
def warped_surface_rectilinear(x,y,z,data,axis):
|
|
import numpy as np
|
|
import pyvista as pv
|
|
x = np.array(x).flatten()
|
|
y = np.array(y).flatten()
|
|
z = np.array(z).flatten()
|
|
if axis==0: assert x.shape[0]==1
|
|
if axis==1: assert y.shape[0]==1
|
|
if axis==2: assert z.shape[0]==1
|
|
normal = [0,0,0]
|
|
normal[axis] = 1
|
|
grid = pv.RectilinearGrid(x,y,z)
|
|
grid['warp'] = data.transpose().ravel()
|
|
return grid.warp_by_scalar(normal=normal)
|
|
|
|
def warped_surface_uniform(origin,spacing,data,axis,factor=1.0):
|
|
import numpy as np
|
|
import pyvista as pv
|
|
assert data.shape[axis]==1
|
|
normal = [0,0,0]
|
|
normal[axis] = 1
|
|
grid = pv.UniformGrid()
|
|
grid.dimensions = data.shape
|
|
grid.origin = origin
|
|
grid.spacing = spacing
|
|
grid['warp'] = data.transpose().ravel()
|
|
return grid.warp_by_scalar(normal=normal,factor=factor)
|
|
|
|
#def extend_structuredgrid(sg,bounds,rep,axis):
|