root / ase / visualize / vtk / sources.py @ 1
Historique | Voir | Annoter | Télécharger (3,84 ko)
| 1 |
|
|---|---|
| 2 |
import numpy as np |
| 3 |
|
| 4 |
from vtk import vtkProperty, vtkSphereSource, vtkArrowSource, vtkConeSource |
| 5 |
|
| 6 |
from ase.data import atomic_numbers |
| 7 |
from ase.data import covalent_radii as atomic_radii |
| 8 |
from ase.data.colors import jmol_colors as atomic_colors |
| 9 |
#from ase.data.colors import cpk_colors as atomic_colors
|
| 10 |
|
| 11 |
avg_radius = np.mean(atomic_radii[np.isfinite(atomic_radii)]) |
| 12 |
|
| 13 |
# -------------------------------------------------------------------
|
| 14 |
|
| 15 |
class vtkCustomGlyphSource: |
| 16 |
def __init__(self, scale, vtk_glyph_source): |
| 17 |
self.scale = scale
|
| 18 |
self.vtk_glyph_source = vtk_glyph_source
|
| 19 |
self.vtk_property = vtkProperty()
|
| 20 |
|
| 21 |
def get_scale(self): |
| 22 |
return self.scale |
| 23 |
|
| 24 |
def get_property(self): |
| 25 |
return self.vtk_property |
| 26 |
|
| 27 |
def get_output(self): |
| 28 |
return self.vtk_glyph_source.GetOutput() |
| 29 |
|
| 30 |
class vtkAtomSource(vtkCustomGlyphSource): |
| 31 |
def __init__(self, name, scale=1, fraction=0.25): |
| 32 |
vtkCustomGlyphSource.__init__(self, scale, vtkSphereSource())
|
| 33 |
|
| 34 |
self.number = atomic_numbers[name]
|
| 35 |
self.radius = atomic_radii[self.number] |
| 36 |
self.color = atomic_colors[self.number] |
| 37 |
|
| 38 |
self.vtk_property.SetColor(self.color[0],self.color[1],self.color[2]) |
| 39 |
self.vtk_property.SetInterpolationToPhong()
|
| 40 |
self.vtk_property.SetDiffuse(0.7) |
| 41 |
self.vtk_property.SetSpecular(0.4) |
| 42 |
self.vtk_property.SetSpecularPower(20) |
| 43 |
|
| 44 |
self.vtk_glyph_source.SetPhiResolution(16) |
| 45 |
self.vtk_glyph_source.SetThetaResolution(16) |
| 46 |
self.vtk_glyph_source.SetRadius(fraction*self.radius) |
| 47 |
|
| 48 |
# -------------------------------------------------------------------
|
| 49 |
|
| 50 |
class vtkClampedGlyphSource(vtkCustomGlyphSource): |
| 51 |
def __init__(self, scale, vtk_glyph_source, range_min, range_max): |
| 52 |
vtkCustomGlyphSource.__init__(self, scale, vtk_glyph_source)
|
| 53 |
self.range = (range_min, range_max,)
|
| 54 |
|
| 55 |
def get_range(self): |
| 56 |
return self.range |
| 57 |
|
| 58 |
class vtkForceSource(vtkClampedGlyphSource): |
| 59 |
def __init__(self, maxnorm, scale=1): |
| 60 |
vtkClampedGlyphSource.__init__(self, scale, vtkArrowSource(),
|
| 61 |
range_min=0.0, range_max=maxnorm)
|
| 62 |
|
| 63 |
self.vtk_property.SetColor(1.0, 0.25, 0.25) # forces are red |
| 64 |
self.vtk_property.SetInterpolationToPhong()
|
| 65 |
self.vtk_property.SetDiffuse(0.7) |
| 66 |
self.vtk_property.SetSpecular(0.4) |
| 67 |
self.vtk_property.SetSpecularPower(20) |
| 68 |
|
| 69 |
self.vtk_glyph_source.SetShaftResolution(12) |
| 70 |
self.vtk_glyph_source.SetShaftRadius(0.03*avg_radius) #default 0.03 |
| 71 |
self.vtk_glyph_source.SetTipResolution(20) |
| 72 |
self.vtk_glyph_source.SetTipLength(0.3*avg_radius) #default 0.35 |
| 73 |
self.vtk_glyph_source.SetTipRadius(0.1*avg_radius) #default 0.1 |
| 74 |
|
| 75 |
class vtkVelocitySource(vtkClampedGlyphSource): |
| 76 |
def __init__(self, maxnorm, scale=1): |
| 77 |
vtkClampedGlyphSource.__init__(self, scale, vtkConeSource(),
|
| 78 |
range_min=0.0, range_max=maxnorm)
|
| 79 |
|
| 80 |
self.vtk_property.SetColor(0.25, 0.25, 1.0) # velocities blue |
| 81 |
self.vtk_property.SetInterpolationToPhong()
|
| 82 |
self.vtk_property.SetDiffuse(0.9) |
| 83 |
self.vtk_property.SetSpecular(1.0) |
| 84 |
self.vtk_property.SetSpecularPower(50) |
| 85 |
|
| 86 |
self.vtk_glyph_source.SetAngle(6) |
| 87 |
self.vtk_glyph_source.SetHeight(avg_radius)
|
| 88 |
self.vtk_glyph_source.SetResolution(16) |
| 89 |
self.vtk_glyph_source.SetCenter(0.05*avg_radius, 0.0, 0.0) |
| 90 |
|
| 91 |
# -------------------------------------------------------------------
|
| 92 |
|
| 93 |
class vtkBondSource(vtkCustomGlyphSource): |
| 94 |
def __init__(self, width, scale=1): |
| 95 |
vtkCustomGlyphSource.__init__(self, scale, vtkCylinderSource())
|
| 96 |
|
| 97 |
self.width = width
|
| 98 |
|
| 99 |
self.vtk_property.SetColor(0.25, 1.0, 0.25) # bonds are green |
| 100 |
# (and so are you)
|
| 101 |
|
| 102 |
self.vtk_glyph_source.SetRadius(self.scale*self.width) |
| 103 |
self.vtk_glyph_source.SetResolution(16) |
| 104 |
|