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