added particle visualization methods

This commit is contained in:
Michael Krayer 2021-05-27 23:38:53 +02:00
parent 1756a7a916
commit 0619930cfa
1 changed files with 17 additions and 1 deletions

View File

@ -76,7 +76,7 @@ class Particles:
attr = {} attr = {}
for key in self.attr: for key in self.attr:
attr[key] = self.attr[key].copy() attr[key] = self.attr[key].copy()
return Particles(self.num,self.time,attr,period) return Particles(self.num,self.time,attr,self.period)
def add_attribute(self,key,val): def add_attribute(self,key,val):
import numpy import numpy
if isinstance(val,(tuple,list,numpy.ndarray)): if isinstance(val,(tuple,list,numpy.ndarray)):
@ -125,6 +125,22 @@ class Particles:
key = ('x','y','z')[axis] key = ('x','y','z')[axis]
self.attr[key] %= self.period[axis] self.attr[key] %= self.period[axis]
return return
def to_vtk(self,deep=False):
import pyvista as pv
import numpy as np
position = np.vstack([self.attr[key] for key in ('x','y','z')]).transpose()
mesh = pv.PolyData(position,deep=deep)
for key in self.attr:
mesh[key] = self.attr[key]
return mesh
def glyph(self,theta_resolution=10,phi_resolution=10,deep=False):
import pyvista as pv
assert self.has_attribute('r'), "Attribute 'r' required."
geom = pv.Sphere(radius=1,theta_resolution=theta_resolution,
phi_resolution=phi_resolution)
return self.to_vtk(deep=deep).glyph(scale='r',factor=1,geom=geom)
class Trajectories: class Trajectories:
def __init__(self,part,unraveled=False,copy_particles=False): def __init__(self,part,unraveled=False,copy_particles=False):
import numpy import numpy