Commit fb597c5f authored by Lan Huang's avatar Lan Huang

add geom.view()

parent 4571189e
......@@ -27,7 +27,8 @@ setup(name='westpy',
'urllib3',
'prompt-toolkit',
'sphinx',
'sphinx_rtd_theme'
'sphinx_rtd_theme',
'py3Dmol'
],
python_requires='>=2.7, >=3.0, !=3.0.*, !=3.1.*, !=3.2.*, <4',
zip_safe=True)
......@@ -150,18 +150,23 @@ class Geometry(object) :
self.atoms.append( Atom(symbol, position, units) )
self.isSet["atoms"] = True
#
def __addAtomsFromXYZLines(self, lines ) :
def __addAtomsFromXYZLines(self, lines, decode=True ) :
"""Adds atoms from XYZ lines.
:param lines: lines read from XYZ file (only one image)
:type lines: list of string
:param decode:
:type bool:
"""
#
from westpy import Angstrom
natoms = int(lines[0])
for line in lines[2:2+natoms] :
symbol, x, y, z = line.split()[:4]
self.addAtom( symbol.decode("utf-8"), (float(x), float(y), float(z)), units=Angstrom )
if decode:
self.addAtom( symbol.decode("utf-8"), (float(x), float(y), float(z)), units=Angstrom )
else:
self.addAtom( symbol, (float(x), float(y), float(z)), units=Angstrom )
#
def addAtomsFromXYZFile(self, fname ) :
"""Adds atoms from XYZ file (only one image).
......@@ -178,7 +183,7 @@ class Geometry(object) :
#
with open(fname,'r') as file:
lines = file.readlines()
self.__addAtomsFromXYZLines(lines)
self.__addAtomsFromXYZLines(lines, decode=False)
#
def addAtomsFromOnlineXYZ(self, url ) :
"""Adds atoms from XYZ file (only one image) located at url.
......@@ -196,7 +201,7 @@ class Geometry(object) :
import urllib.request
with urllib.request.urlopen(url) as response :
lines = response.readlines()
self.__addAtomsFromXYZLines(lines)
self.__addAtomsFromXYZLines(lines, decode=True)
#
def getNumberOfAtoms(self) :
"""Returns number of atoms.
......@@ -294,3 +299,63 @@ class Geometry(object) :
from westpy import download
for key in self.species.keys() :
download( self.species[key]["url"], fname=self.species[key]["fname"] )
def view(self, style='stick',width=800,height=800, ix=1, iy=1, iz=1, debug=False):
"""Display simulation box geom in Angstrom.
ix, iy, iz is the perodic display to system
style can be line, stick, sphere.
:param style:
:param width:
:param height:
:param ix:
:param iy:
:param iz:
:param debug:
:return:
"""
import py3Dmol
import numpy as np
from westpy.units import Angstrom
BOHR2A = 1.0 / Angstrom
a1 = self.cell["a1"] * BOHR2A
a2 = self.cell["a2"] * BOHR2A
a3 = self.cell["a3"] * BOHR2A
nat = self.getNumberOfAtoms()
times = ix * iy * iz
if times > 1:
nat *= times
#generate xyz data
xyz = str(nat) + '\n\n'
for atom in self.atoms:
if times:
for i in range(ix):
for j in range(iy):
for k in range(iz):
xyz += atom.symbol + " " + " ".join(
map(str, atom.position * BOHR2A + i * a1 + j * a2 + k * a3)) + '\n'
#creat viewer
xyzview = py3Dmol.view(width=width, height=height)
xyzview.addModel(xyz, 'xyz')
if debug:
print(xyz)
xyzview.setStyle({style: {}})
#draw the box
a0 = np.array([0.0, 0.0, 0.0])
from_ = [a0, a1 + a2, a1 + a3, a2 + a3]
to_ = [[a1, a2, a3], [a1, a2, a1 + a2 + a3], [a1, a3, a1 + a2 + a3], [a2, a3, a1 + a2 + a3]]
for frm, li_to in zip(from_, to_):
x0, y0, z0 = frm
for to in li_to:
x1, y1, z1 = to
xyzview.addLine(
{'color': 'blue', 'start': {'x': x0, 'y': y0, 'z': z0}, 'end': {'x': x1, 'y': y1, 'z': z1}});
#show
xyzview.zoomTo()
xyzview.show()
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment