141 lines
6.5 KiB
Python
141 lines
6.5 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)
|
||
return
|
||
|
||
def add_particles(plotter,pp,col,itime=0,
|
||
name=None,
|
||
color='white',
|
||
scalars=None,cmap=None,clim=None,
|
||
opacity=1.0,
|
||
theta_resolution=10,phi_resolution=10):
|
||
'''
|
||
name (str, optional) – The name for the added mesh/actor so that it can be easily updated. If an actor of this name already exists in the rendering window, it will be replaced by the new actor.
|
||
|
||
color (string or 3 item list, optional, defaults to white) –
|
||
Use to make the entire mesh have a single solid color. Either a string, RGB list, or hex color string. For example: color='white', color='w', color=[1, 1, 1], or color='#FFFFFF'. Color will be overridden if scalars are specified.
|
||
|
||
scalars (str, optional) –
|
||
Scalars used to “color” the mesh. Accepts a key of the col dictionary. If both color and scalars are None, then the active scalars are used.
|
||
|
||
cmap (str, list, optional) –
|
||
Name of the Matplotlib colormap to use when mapping the scalars. See available Matplotlib colormaps. Only applicable for when displaying scalars. Requires Matplotlib to be installed. colormap is also an accepted alias for this. If colorcet or cmocean are installed, their colormaps can be specified by name.
|
||
You can also specify a list of colors to override an existing colormap with a custom one. For example, to create a three color colormap you might specify ['green', 'red', 'blue']
|
||
|
||
clim (2 item list, optional) –
|
||
Color bar range for scalars. Defaults to minimum and maximum of scalars array. Example: [-1, 2]. rng is also an accepted alias for this.
|
||
|
||
opacity (float, str, array-like) –
|
||
Opacity of the mesh. If a single float value is given, it will be the global opacity of the mesh and uniformly applied everywhere - should be between 0 and 1. A string can also be specified to map the scalars range to a predefined opacity transfer function (options include: ‘linear’, ‘linear_r’, ‘geom’, ‘geom_r’). A string could also be used to map a scalars array from the mesh to the opacity (must have same number of elements as the scalars argument). Or you can pass a custom made transfer function that is an array either n_colors in length or shorter.
|
||
'''
|
||
import pyvista
|
||
ncol,npart,ntime = pp.shape
|
||
assert(itime<ntime)
|
||
assert('x' in col)
|
||
assert('y' in col)
|
||
assert('z' in col)
|
||
assert('r' in col)
|
||
pts = pyvista.PolyData(pp[(col['x'],col['y'],col['z']),:,itime].transpose())
|
||
pts['r'] = pp[(col['r']),:,itime].squeeze()
|
||
|
||
if scalars:
|
||
assert(scalars in col)
|
||
pts[scalars] = pp[(col[scalars]),:,itime].squeeze()
|
||
color=None
|
||
|
||
geom = pyvista.Sphere(radius=1,theta_resolution=theta_resolution,phi_resolution=phi_resolution)
|
||
glyphs = pts.glyph(scale='r',factor=1,geom=geom)
|
||
plotter.add_mesh(glyphs,name=name,
|
||
color=color,
|
||
scalars=scalars,cmap=cmap,clim=clim,
|
||
opacity=opacity)
|
||
return
|
||
|
||
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 field_to_pvmesh(field):
|
||
import pyvista
|
||
mesh = pyvista.UniformGrid()
|
||
mesh.dimensions = field.dim(axis=None)
|
||
mesh.origin = field.origin # The bottom left corner of the data set
|
||
mesh.spacing = field.spacing # These are the cell sizes along each axis
|
||
mesh.point_arrays['values'] = field.data.flatten(order='F') # TBD: check order
|
||
return mesh
|
||
|
||
# def common_mesh(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
|
||
# return mesh
|
||
|
||
# def mesh_qcriterion(chunks_uvw,gridg_uvwp):
|
||
# for chunk in chunks_uvw:
|
||
# assert(chunk['ighost']==1)
|
||
# import pyvista
|
||
# dx = gridg_uvwp[0][0][1]-gridg_uvwp[0][0][1]
|
||
# dy = gridg_uvwp[0][1][1]-gridg_uvwp[0][1][1]
|
||
# dz = gridg_uvwp[0][2][1]-gridg_uvwp[0][2][1]
|
||
# xq0 = gridg_uvwp[3][0][chunks]
|
||
|
||
|
||
# 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
|
||
|
||
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 |