Files
PlantAI-dev/tools/pymol-env/lib/python3.12/site-packages/pymol/cgo.py
T
2026-08-03 18:42:47 +08:00

618 lines
18 KiB
Python

#A* -------------------------------------------------------------------
#B* This file contains source code for the PyMOL computer program
#C* Copyright (c) Schrodinger, LLC.
#D* -------------------------------------------------------------------
#E* It is unlawful to modify or remove this copyright notice.
#F* -------------------------------------------------------------------
#G* Please see the accompanying LICENSE file for further information.
#H* -------------------------------------------------------------------
#I* Additional authors of this source file include:
#-*
#-*
#-*
#Z* -------------------------------------------------------------------
from chempy import cpv
#import popen2
import os
from pymol import cmd
from .cmd import DEFAULT_ERROR, DEFAULT_SUCCESS, _raising
POINTS = 0.0
LINES = 1.0
LINE_LOOP = 2.0
LINE_STRIP = 3.0
TRIANGLES = 4.0
TRIANGLE_STRIP = 5.0
TRIANGLE_FAN = 6.0
#QUADS = 7.0
#QUAD_STRIP = 8.0
#POLYGON = 9.0
STOP = 0.0
NULL = 1.0
BEGIN = 2.0
END = 3.0
VERTEX = 4.0
NORMAL = 5.0
COLOR = 6.0
SPHERE = 7.0
TRIANGLE = 8.0
CYLINDER = 9.0
LINEWIDTH = 10.0
WIDTHSCALE = 11.0
ENABLE = 12.0
DISABLE = 13.0
SAUSAGE = 14.0
CUSTOM_CYLINDER = 15.0
DOTWIDTH = 16.0
ALPHA_TRIANGLE = 17.0
ELLIPSOID = 18.0
#SHAPE_VERTEX = 16.0
#SHAPE_COLOR = 17.0
#SHAPE_NORMAL = 18.0
FONT = 19.0
FONT_SCALE = 20.0
FONT_VERTEX = 21.0
FONT_AXES = 22.0
CHAR = 23.0
ALPHA = 25.0
QUADRIC = 26.0 # NOTE: Only works with ellipsoids and disks
CONE = 27.0
PICK_COLOR = 31.0 # 0x1F [PICK_COLOR, index, bond/cPickable_t]
BEZIER = 66.0 # 0x42
LIGHTING = float(0x0B50)
# enum cPickable_t:
cPickableAtom = -1.0
cPickableLabel = -2.0
cPickableGadget = -3.0
cPickableNoPick = -4.0
cPickableThrough = -5.0
def molauto(name="mols", sele="(all)", marg="-nice", _self=cmd):
_self.save("molauto.pdb",sele)
print("molauto %s -nocentre molauto.pdb | molscript -r > molauto.r3d"%marg)
os.system("molauto %s -nocentre molauto.pdb | molscript -r > molauto.r3d"%marg)
f = open("molauto.r3d")
rr = RenderReader(f)
f.close()
_self.load_cgo(rr.obj,name)
def measure_text(font,text,
axes=[[1.0,0.0,0.0],[0.0,1.0,0.0],[0.0,0.0,1.0]]):
w = 0
x = axes[0]
for char in text:
if char in font:
w = w + font[char][0]*x[0]
return w
def wire_text(cgo,font,pos,text,
axes = [[1.0,0.0,0.0],[0.0,1.0,0.0],[0.0,0.0,1.0]]): # modifies pos
x = axes[0]
y = axes[1]
for char in text:
if char in font:
fc = font[char]
stroke = 0
w = fc[0]
f = fc[1]
c = 0
l = len(f)-2
while c<l:
if not f[c]:
if stroke: cgo.append(END)
cgo.append(BEGIN)
cgo.append(LINE_STRIP)
stroke = 1
ax = f[c+1]
ay = f[c+2]
cgo.append(VERTEX)
cgo.append(pos[0]+x[0]*ax+y[0]*ay)
cgo.append(pos[1]+x[1]*ax+y[1]*ay)
cgo.append(pos[2]+x[2]*ax+y[2]*ay)
c = c + 3
pos[0] = pos[0] + w*x[0]
pos[1] = pos[1] + w*x[1]
pos[2] = pos[2] + w*x[2]
if stroke: cgo.append(END)
def cyl_text(cgo,font,pos,text,radius=0.1,color=[1.0,1.0,1.0],
axes = [[1.0,0.0,0.0],[0.0,1.0,0.0],[0.0,0.0,1.0]]): # modifies pos
x = axes[0]
y = axes[1]
for char in text:
if char in font:
fc = font[char]
stroke = 0
w = fc[0]
f = fc[1]
c = 0
l = len(f)-2
while c<l:
ax = f[c+1]
ay = f[c+2]
next = [(pos[0]+x[0]*ax+y[0]*ay),
(pos[1]+x[1]*ax+y[1]*ay),
(pos[2]+x[2]*ax+y[2]*ay)]
if f[c]:
if stroke:
cgo.append(SAUSAGE)
cgo.extend(last)
cgo.extend(next)
cgo.append(radius)
cgo.extend(color)
cgo.extend(color)
else:
stroke = 1
last = next
c = c + 3
pos[0] = pos[0] + w*x[0]
pos[1] = pos[1] + w*x[1]
pos[2] = pos[2] + w*x[2]
def from_r3d(fname):
result = DEFAULT_ERROR
input = None
if '://' in fname:
from urllib.request import urlopen
input = urlopen(fname)
elif os.path.exists(fname):
input = open(fname)
if input:
rr = RenderReader(input)
result = rr.obj
return result
class RenderReader:
def append_last(self):
if self.app_fn:
self.app_fn()
self.app_fn=None
def append_tri(self):
if self.l_vert:
d0 = cpv.sub(self.l_vert[0],self.l_vert[1])
d1 = cpv.sub(self.l_vert[0],self.l_vert[2])
n0 = cpv.cross_product(d0,d1)
n0 = cpv.normalize_failsafe(n0)
if not self.tri_flag:
self.obj.append(BEGIN)
self.obj.append(TRIANGLES)
self.tri_flag = 1
indices = [0, 1, 2]
if not self.l_norm:
# TODO could simplify this if ray tracing would support
# object-level two_sided_lighting. Duplicating the
# face with an offset is a hack and produces visible
# lines on edges.
n1 = [-n0[0],-n0[1],-n0[2]]
ns = cpv.scale(n0,0.002)
indices = [0, 1, 2, 4, 3, 5]
l_vert_offsetted = [cpv.add(v, ns) for v in self.l_vert]
l_vert_offsetted.extend(cpv.sub(v, ns) for v in self.l_vert)
self.l_vert = l_vert_offsetted
self.l_norm = [n0, n0, n0, n1, n1, n1]
elif cpv.dot_product(self.l_norm[0], n0) < 0:
indices = [0, 2, 1]
for i in indices:
self.obj.append(COLOR) # assuming unicolor
self.obj.extend(self.t_colr[i % 3])
self.obj.append(NORMAL)
self.obj.extend(self.l_norm[i])
self.obj.append(VERTEX)
self.obj.extend(self.l_vert[i])
self.l_vert=None
self.t_colr=None
self.l_norm=None
def append_cyl(self):
if self.l_vert and self.c_colr and self.l_radi:
if self.tri_flag:
self.tri_flag=0
self.obj.append(END)
self.obj.append(SAUSAGE)
d = cpv.sub(self.l_vert[1],self.l_vert[0])
d = cpv.normalize_failsafe(d)
d0 = cpv.scale(d,self.l_radi/4.0)
self.obj.extend(cpv.add(self.l_vert[0],d0))
self.obj.extend(cpv.sub(self.l_vert[1],d0))
self.obj.append(self.l_radi)
self.obj.extend(self.c_colr[0])
self.obj.extend(self.c_colr[1])
self.l_vert=None
self.c_colr=None
self.l_radi=None
def tri(self,f):
self.append_last()
l = f.readline()
if l:
self.app_fn=self.append_tri
s = l.split()
self.l_vert = [[float(s[0]),float(s[1]),float(s[2])],
[float(s[3]),float(s[4]),float(s[5])],
[float(s[6]),float(s[7]),float(s[8])]]
self.t_colr_t = [float(s[9]),float(s[10]),float(s[11])]
self.t_colr = [self.t_colr_t,self.t_colr_t,self.t_colr_t]
def tri_normal(self,f):
l = f.readline()
if l:
s = l.split()
self.l_norm = [[float(s[0]),float(s[1]),float(s[2])],
[float(s[3]),float(s[4]),float(s[5])],
[float(s[6]),float(s[7]),float(s[8])]]
def cyl(self,f):
self.append_last()
l = f.readline()
if l:
self.app_fn = self.append_cyl
s = l.split()
self.l_vert = [[float(s[0]),float(s[1]),float(s[2])],
[float(s[4]),float(s[5]),float(s[6])]]
self.l_radi = float(s[3])
self.c_colr_t = [float(s[8]),float(s[9]),float(s[10])]
self.c_colr = [self.c_colr_t,self.c_colr_t]
def sphere(self,f):
self.append_last()
l = f.readline()
if l:
s = l.split()
self.obj.append(COLOR)
self.obj.extend([float(s[4]),float(s[5]),float(s[6])])
self.obj.append(SPHERE)
self.obj.extend([float(s[0]),float(s[1]),float(s[2]),float(s[3])])
def quadric(self,f):
self.append_last()
l = f.readline()
if l:
s = l.split()
self.obj.append(COLOR)
self.obj.extend([float(s[4]),float(s[5]),float(s[6])])
self.obj.append(QUADRIC)
self.obj.extend([float(s[0]),float(s[1]),float(s[2]),float(s[3])])
l = f.readline()
if l:
s = l.split()
self.obj.extend([float(s[0]),float(s[1]),float(s[2]),
float(s[3]),float(s[4]),float(s[5]),
float(s[6]),float(s[7]),float(s[8]),float(s[9])])
def mat_prop(self,f):
self.append_last()
l = f.readline()
print("mat_prop"+l)
if l:
s = l.split()
(mphong, mspec, sr, sg, sb, clrity) = map(float,s[0:6])
if clrity>0.999:
clrity=0.999
self.obj.extend([ALPHA, 1.0-clrity])
opts1 = int(s[6])
opts4 = int(s[9])
for x in range(opts4):
f.readline();
def mat_reset(self,f):
self.append_last()
self.obj.extend([ALPHA, 1.0])
def reset(self,f):
pass
def __init__(self,input):
# Author: Warren DeLano
# Modifications: Robert Campbell
self.app_fn = None
self.l_vert = None
self.t_colr = None
self.c_colr = None
self.l_radi = None
self.l_norm = None
self.o_vert = None
self.tri_flag = 0
self.cc = 0
self.obj = []
for a in range(20):
input.readline()
dispatch = [
None,
self.tri,
self.sphere,
self.cyl,
None,
self.cyl,
None,
self.tri_normal,
self.mat_prop,
self.reset,
None,
None,
None,
None,
self.quadric,
]
ld = len(dispatch)
while 1:
l = input.readline()
if not l:
break
if l[0] != '#':
v = l.split()
n=int(v[0])
if(n<ld):
dd = dispatch[n]
if dd:
dd(input)
else:
# skip over lines that don't match desired object type
input.readline()
elif ( n != 9 ):
# don't read another line if render object type 0
input.readline()
self.append_last()
if self.tri_flag:
self.obj.append(END)
input.close()
def torus(center=(0., 0., 0.), normal=(0., 0., 1.), radius=1., color='',
cradius=.25, samples=20, csamples=20, _self=cmd):
'''
Generate and return a torus CGO with given center, normal
and ring radius.
'''
from math import cos, sin, pi
if color and isinstance(color, str):
color = list(_self.get_color_tuple(color))
obj = []
axis = cpv.cross_product(normal, (0., 0., 1.))
angle = -cpv.get_angle(normal, (0., 0., 1.))
matrix = cpv.rotation_matrix(angle, cpv.normalize(axis))
obj_vertex = lambda x, y, z: obj.extend([VERTEX] + cpv.add(center,
cpv.transform(matrix, [x, y, z])))
obj_normal = lambda x, y, z: obj.extend([NORMAL] +
cpv.transform(matrix, [x, y, z]))
r = radius
cr = cradius
rr = 1.5 * cr
dv = 2 * pi / csamples
dw = 2 * pi / samples
v = 0.0
w = 0.0
while w < 2 * pi:
v = 0.0
c_w = cos(w)
s_w = sin(w)
c_wdw = cos(w + dw)
s_wdw = sin(w + dw)
obj.append(BEGIN)
obj.append(TRIANGLE_STRIP)
if color:
obj.append(COLOR)
obj.extend(color)
while v < 2 * pi + dv:
c_v = cos(v)
s_v = sin(v)
c_vdv = cos(v + dv)
s_vdv = sin(v + dv)
obj_normal(
(r + rr * c_v) * c_w - (r + cr * c_v) * c_w,
(r + rr * c_v) * s_w - (r + cr * c_v) * s_w,
(rr * s_v - cr * s_v))
obj_vertex(
(r + cr * c_v) * c_w,
(r + cr * c_v) * s_w,
cr * s_v)
obj_normal(
(r + rr * c_vdv) * c_wdw - (r + cr * c_vdv) * c_wdw,
(r + rr * c_vdv) * s_wdw - (r + cr * c_vdv) * s_wdw,
rr * s_vdv - cr * s_vdv)
obj_vertex(
(r + cr * c_vdv) * c_wdw,
(r + cr * c_vdv) * s_wdw,
cr * s_vdv)
v += dv
obj.append(END)
w += dw
return obj
def from_plystr(contents, surfacenormals=True, alphaunit=1.):
'''
PLY - Polygon File Format
'''
import sys
if isinstance(contents, bytes):
contents = contents.decode(errors='ignore')
lines_iter = iter(contents.splitlines())
if next(lines_iter) != 'ply':
raise ValueError('not a ply file')
types = {'char': int, 'uchar': int, 'short': int, 'ushort': int,
'int': int, 'uint': int, 'float': float, 'double': float,
'int8': int, 'uint8': int, 'int16': int, 'uint16': int,
'int32': int, 'uint32': int, 'float32': float, 'float64': float}
elements = []
for line in lines_iter:
a = line.split()
if not a:
continue
command = a[0]
if command == 'end_header':
break
if command == 'format' or command == 'comment':
continue
if command == 'element':
element = {'type': a[1], 'count': int(a[2]), 'properties': []}
elements.append(element)
continue
if command == 'property':
element['properties'].append(a[1:])
continue
if command == 'obj_info':
# ignore
continue
print('unknown instruction: ' + command)
table = {}
for element in elements:
table[element['type']] = records = []
properties = element['properties']
for i in range(element['count']):
line = next(lines_iter)
a_iter = iter(line.split())
rec = {}
for prop in properties:
if prop[0] == 'list':
rec[prop[-1]] = [types[prop[2]](next(a_iter))
for j in range(types[prop[2]](next(a_iter)))]
else:
rec[prop[-1]] = types[prop[0]](next(a_iter))
# auto-detect alpha divider
if alphaunit == 1. and prop[-1] == 'alpha' and rec['alpha'] > 1:
alphaunit = 255.
# detect if normals need to be calculated
if surfacenormals and prop[-1] == 'nx':
surfacenormals = False
records.append(rec)
vertices = table['vertex']
obj = []
enum_quad = [0, 1, 2, 2, 3, 0]
def colorfromelem(elem):
if 'red' in elem:
obj.append(COLOR)
obj.append(elem['red'] / 255.0)
obj.append(elem['green'] / 255.0)
obj.append(elem['blue'] / 255.0)
if 'alpha' in elem:
obj.append(ALPHA)
obj.append(elem['alpha'] / alphaunit)
def compute_surface_normals():
'''
Compute average normals from all adjacent triangles
on each vertex
'''
from functools import reduce
# don't use cpv.normalize which has an RSMALL4 limit
normalize = lambda v: cpv.scale(v, 1. / cpv.length(v))
for face in table['face']:
if 'vertex_index' in face:
indices = face['vertex_index']
elif 'vertex_indices' in face:
indices = face['vertex_indices']
else:
return
f_vert = [vertices[i] for i in indices]
f_xyz = [(v['x'], v['y'], v['z']) for v in f_vert]
try:
normal = normalize(cpv.cross_product(
cpv.sub(f_xyz[1], f_xyz[0]),
cpv.sub(f_xyz[2], f_xyz[1])))
except ZeroDivisionError:
continue
for v in f_vert:
v.setdefault('normals', []).append(normal)
for v in vertices:
try:
v['nx'], v['ny'], v['nz'] = normalize(
reduce(cpv.add, v.pop('normals')))
except (KeyError, ZeroDivisionError):
continue
if surfacenormals and 'face' in table:
compute_surface_normals()
for ptype, op in [
('face', TRIANGLES),
('edge', LINES),
('range_grid', POINTS),
]:
if ptype in table:
obj.append(BEGIN)
obj.append(op)
for face in table[ptype]:
colorfromelem(face)
if 'vertex_index' in face:
indices = face['vertex_index']
elif 'vertex_indices' in face:
indices = face['vertex_indices']
else:
indices = [face['vertex1'], face['vertex2']]
N = len(indices)
for i in (enum_quad if N == 4 else range(N)):
vertex = vertices[indices[i]]
colorfromelem(vertex)
if 'nx' in vertex:
obj.append(NORMAL)
obj.append(vertex['nx'])
obj.append(vertex['ny'])
obj.append(vertex['nz'])
obj.append(VERTEX)
obj.append(vertex['x'])
obj.append(vertex['y'])
obj.append(vertex['z'])
obj.append(END)
obj.append(STOP)
return obj