mirror of
https://github.com/schrodinger/pymol-open-source.git
synced 2026-06-04 20:04:21 +08:00
350 lines
9.4 KiB
Python
350 lines
9.4 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* -------------------------------------------------------------------
|
|
|
|
import dataclasses
|
|
import math
|
|
from typing import List, Dict
|
|
|
|
from chempy import cpv
|
|
|
|
from pymol import cgo
|
|
|
|
from . import cmd
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class Point:
|
|
"""
|
|
Represents a 3D point in space
|
|
|
|
Attributes:
|
|
x (float): x-coordinate
|
|
y (float): y-coordinate
|
|
z (float): z-coordinate
|
|
"""
|
|
x: float = 0.0
|
|
y: float = 0.0
|
|
z: float = 0.0
|
|
|
|
def array(self) -> List[float]:
|
|
"""
|
|
Converts point into 3-component array
|
|
:returns : array of xyz coordinates
|
|
"""
|
|
return [self.x, self.y, self.z]
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class Color:
|
|
"""
|
|
Represents a Color in RGBA space
|
|
|
|
Attributes:
|
|
r (float): red channel value
|
|
g (float): green channel value
|
|
b (float): blue channel value
|
|
a (float): alpha channel value
|
|
"""
|
|
r: float = 0.0
|
|
g: float = 0.0
|
|
b: float = 0.0
|
|
a: float = 1.0
|
|
|
|
def array(self) -> List[float]:
|
|
"""
|
|
Converts color into 3-component array
|
|
"""
|
|
return [self.r, self.g, self.b]
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class CGO:
|
|
"""
|
|
Base class for all CGOs
|
|
|
|
For all CGO types, whenever attributes are changed, `rebuild` should be
|
|
called to update the data cache.
|
|
"""
|
|
_data: List[float] = dataclasses.field(default_factory=list)
|
|
|
|
def __post_init__(self):
|
|
"""
|
|
Rebuilds CGO. All child classes should implement this.
|
|
"""
|
|
self.rebuild()
|
|
|
|
@property
|
|
def data(self):
|
|
return self._data
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class Sphere(CGO):
|
|
"""
|
|
Sphere CGO
|
|
|
|
Attributes:
|
|
center (Point): center position
|
|
radius (float): sphere radius
|
|
color (Color): sphere color
|
|
"""
|
|
center: Point = Point(0, 0, 0)
|
|
radius: float = 0.0
|
|
color: Color = Point(0, 0, 0)
|
|
|
|
def rebuild(self) -> None:
|
|
"""
|
|
Rebuilds sphere
|
|
"""
|
|
self._data.clear()
|
|
self._data.append(cgo.COLOR)
|
|
self._data.extend(self.color.array())
|
|
self._data.append(cgo.SPHERE)
|
|
self._data.extend(self.center.array())
|
|
self._data.append(self.radius)
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class Cylinder(CGO):
|
|
"""
|
|
Cylinder CGO
|
|
|
|
Attributes:
|
|
point1 (Point): center position of first side
|
|
point2 (Point): center position of second side
|
|
radius (float): sphere radius
|
|
color1 (Color): color of first cylinder half
|
|
color2 (Color): color of second cylinder half
|
|
"""
|
|
point1: Point = Point(0, 0, 0)
|
|
point2: Point = Point(1, 1, 1)
|
|
radius: float = 1.0
|
|
color1: Color = Point(0, 0, 0)
|
|
color2: Color = Color(0, 0, 0)
|
|
|
|
def rebuild(self) -> None:
|
|
"""
|
|
Rebuilds cylinder
|
|
"""
|
|
self._data.clear()
|
|
self._data.append(cgo.CYLINDER)
|
|
self._data.extend(self.point1.array())
|
|
self._data.extend(self.point2.array())
|
|
self._data.append(self.radius)
|
|
self._data.extend(self.color1.array())
|
|
self._data.extend(self.color2.array())
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class Torus(CGO):
|
|
"""
|
|
Torus CGO
|
|
|
|
Attributes:
|
|
center (Point): torus center position
|
|
normal (Point): torus normal direction
|
|
radius (float): ring radius
|
|
color (Color): torus color
|
|
cradius (float): torus ring radius
|
|
samples (int): number of samples to generate torus segment
|
|
csamples (int): number of samples to generate torus ring
|
|
"""
|
|
center: Point = Point(0.0, 0.0, 0.0)
|
|
normal: Point = Point(0.0, 0.0, 1.0)
|
|
radius: float = 1.0
|
|
color: Color = Color(0.0, 1.0, 1.0, 1.0)
|
|
cradius: float = 0.25
|
|
samples: int = 20
|
|
csamples: int = 20
|
|
|
|
def rebuild(self) -> None:
|
|
"""
|
|
Rebuilds torus
|
|
"""
|
|
obj = []
|
|
|
|
axis = cpv.cross_product(self.normal.array(), (0., 0., 1.))
|
|
angle = -cpv.get_angle(self.normal.array(), (0., 0., 1.))
|
|
matrix = cpv.rotation_matrix(angle, cpv.normalize(axis))
|
|
|
|
def obj_vertex(x, y, z):
|
|
return [cgo.VERTEX] + cpv.add(self.center.array(),
|
|
cpv.transform(matrix, [x, y, z]))
|
|
|
|
def obj_normal(x, y, z):
|
|
return [cgo.NORMAL] + cpv.transform(matrix, [x, y, z])
|
|
|
|
r = self.radius
|
|
cr = self.cradius
|
|
rr = 1.5 * cr
|
|
dv = 2 * math.pi / self.csamples
|
|
dw = 2 * math.pi / self.samples
|
|
v = 0.0
|
|
w = 0.0
|
|
|
|
while w < 2 * math.pi:
|
|
v = 0.0
|
|
c_w = math.cos(w)
|
|
s_w = math.sin(w)
|
|
c_wdw = math.cos(w + dw)
|
|
s_wdw = math.sin(w + dw)
|
|
|
|
obj.append(cgo.BEGIN)
|
|
obj.append(cgo.TRIANGLE_STRIP)
|
|
|
|
obj.append(cgo.COLOR)
|
|
obj.extend(self.color.array())
|
|
|
|
while v < 2 * math.pi + dv:
|
|
c_v = math.cos(v)
|
|
s_v = math.sin(v)
|
|
c_vdv = math.cos(v + dv)
|
|
s_vdv = math.sin(v + dv)
|
|
obj.extend(
|
|
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.extend(
|
|
obj_vertex((r + cr * c_v) * c_w, (r + cr * c_v) * s_w,
|
|
cr * s_v))
|
|
obj.extend(
|
|
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.extend(
|
|
obj_vertex((r + cr * c_vdv) * c_wdw,
|
|
(r + cr * c_vdv) * s_wdw, cr * s_vdv))
|
|
v += dv
|
|
|
|
obj.append(cgo.END)
|
|
w += dw
|
|
|
|
self._data = obj
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class Cone(CGO):
|
|
"""
|
|
Cone CGO
|
|
|
|
Attributes:
|
|
tip (Point): cone tip position
|
|
base_center (Point): center position of cone base
|
|
radius (float): cone base radius
|
|
color (Color): cone color
|
|
"""
|
|
tip: Point = Point(1.0, 0.0, 0.0)
|
|
base_center: Point = Point(0.0, 0.0, 0.0)
|
|
radius: float = 1.0
|
|
color: Color = Color(0.0, 0.0, 0.0)
|
|
|
|
def rebuild(self) -> None:
|
|
"""
|
|
Rebuilds cone
|
|
"""
|
|
self._data = []
|
|
self._data.append(cgo.CONE)
|
|
self._data.extend(self.tip.array())
|
|
self._data.extend(self.base_center.array())
|
|
self._data.extend([self.radius, 0.0])
|
|
self._data.extend(self.color.array())
|
|
self._data.extend(self.color.array())
|
|
self._data.extend([1.0, 0.0])
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class Bezier(CGO):
|
|
"""
|
|
Cubic Bezier CGO
|
|
|
|
Attributes:
|
|
control_pt_A (Point): on-curve bezier spline position (from)
|
|
A_right_handle (Point): influences curve from A to B
|
|
control_pt_B (Point): on-curve bezier spline position (to)
|
|
B_left_handle (Point): influences curve from A to B
|
|
A_color (Color): color at point A
|
|
B_color (Color): color at point B
|
|
"""
|
|
control_pt_A = Point(-5.0, 0.0, 0.0)
|
|
A_right_handle = Point(0.0, 10.0, 0.0)
|
|
B_left_handle = Point(1.0, -10.0, 0.0)
|
|
control_pt_B = Point(5.0, 0.0, 0.0)
|
|
|
|
def rebuild(self) -> None:
|
|
"""
|
|
Rebuilds bezier spline
|
|
"""
|
|
self._data = []
|
|
self._data.append(cgo.BEZIER)
|
|
self._data.extend(self.control_pt_A.array())
|
|
self._data.extend(self.A_right_handle.array())
|
|
self._data.extend(self.B_left_handle.array())
|
|
self._data.extend(self.control_pt_B.array())
|
|
|
|
|
|
class CGOBuilder:
|
|
"""
|
|
PyMOL CGO Builder
|
|
|
|
Holds a storage of CGOs to load into PyMOL as a single object.
|
|
|
|
Attributes:
|
|
cgos (Dict[CGO]): Collection of CGOs
|
|
"""
|
|
cgos: Dict[int, CGO] = {}
|
|
_token_count = 0
|
|
|
|
def add(self, cgo: CGO) -> int:
|
|
"""
|
|
Adds CGO to builder
|
|
|
|
:param cgo (CGO): cgo to add
|
|
|
|
:return: token ID for added CGO
|
|
"""
|
|
token = self._token_count
|
|
self.cgos[token] = cgo
|
|
self._token_count += 1
|
|
return token
|
|
|
|
def remove(self, token: int) -> None:
|
|
"""
|
|
Removes CGO from builder
|
|
|
|
:param token (int): token of CGO to remove
|
|
"""
|
|
self.cgos.pop(token, None)
|
|
|
|
def clear(self) -> None:
|
|
"""
|
|
Clears the CGO Collection
|
|
"""
|
|
self.cgos.clear()
|
|
|
|
def load(self, name: str, flush: bool = True, _self=cmd) -> None:
|
|
"""
|
|
Loads the CGO as an object into pymol
|
|
|
|
:param name (str): Name of object
|
|
:param flush (bool): Should reset the object data cache after loading
|
|
"""
|
|
all_serialized = []
|
|
for _, cgo_obj in self.cgos.items():
|
|
all_serialized.extend(cgo_obj.data)
|
|
|
|
_self.load_cgo(all_serialized, name)
|
|
if flush:
|
|
self.cgos.clear()
|