使用VisPy在Jupyter中构建分子可视化工具

引言：为什么选择VisPy进行分子可视化？

在科学计算和生物信息学领域，分子结构的可视化是理解复杂生物分子相互作用的关键。传统的可视化工具往往需要复杂的桌面应用程序，而VisPy（Visualization in Python）提供了一个革命性的解决方案——直接在Jupyter笔记本中实现高性能的3D分子可视化。

VisPy是一个基于OpenGL的高性能交互式2D/3D数据可视化库，它利用现代GPU的计算能力来处理大规模数据集。对于分子可视化而言，这意味着：

• 🚀 实时渲染：即使处理成千上万个原子，也能保持流畅的交互体验
• 🎯 精确控制：通过GLSL着色器精确控制每个原子的外观和光照效果
• 📊 无缝集成：与Python科学计算生态完美融合（NumPy、Pandas等）
• 🌐 Web友好：在Jupyter中通过WebGL实现跨平台可视化

环境准备与安装

系统要求

• Python 3.7+
• Jupyter Notebook 或 JupyterLab
• 支持WebGL的现代浏览器

安装VisPy及相关依赖

# 使用pip安装VisPy
pip install vispy

# 安装Jupyter渲染后端
pip install jupyter-rfb

# 可选：安装科学计算常用库
pip install numpy matplotlib

验证安装

import vispy
print(f"VisPy版本: {vispy.__version__}")

# 检查后端支持
from vispy.app.backends import BACKEND_NAMES
print("可用后端:", BACKEND_NAMES)

核心概念：VisPy的架构设计

VisPy采用分层架构，为不同需求的用户提供适当的抽象级别：

graph TB
    A[VisPy架构] --> B[高级API]
    A --> C[中级API Scene系统]
    A --> D[低级API gloo]
    
    B --> E[Plotting模块]
    B --> F[简单可视化]
    
    C --> G[Visuals视觉元素]
    C --> H[Transforms变换系统]
    C --> I[Scene Graph场景图]
    
    D --> J[OpenGL ES 2.0封装]
    D --> K[GLSL着色器管理]
    D --> L[缓冲区操作]

对于分子可视化，我们主要使用gloo（低级API）来实现精确的渲染控制。

构建分子可视化器：分步指南

步骤1：定义GLSL着色器

分子可视化的核心在于顶点着色器和片段着色器的编写：

// 顶点着色器 - 处理原子位置和大小
#version 120

uniform mat4 u_model;
uniform mat4 u_view;
uniform mat4 u_projection;
uniform vec3 u_light_position;

attribute vec3  a_position;
attribute vec3  a_color;
attribute float a_radius;

varying vec3  v_color;
varying vec4  v_eye_position;
varying float v_radius;
varying vec3  v_light_direction;

void main(void) {
    v_radius = a_radius;
    v_color = a_color;
    
    v_eye_position = u_view * u_model * vec4(a_position, 1.0);
    v_light_direction = normalize(u_light_position);
    
    gl_Position = u_projection * v_eye_position;
    
    // 基于距离调整点精灵大小
    vec4 proj_corner = u_projection * vec4(a_radius, a_radius, 
                                         v_eye_position.z, v_eye_position.w);
    gl_PointSize = 512.0 * proj_corner.x / proj_corner.w;
}

// 片段着色器 - 处理光照和材质
#version 120

uniform mat4 u_model;
uniform mat4 u_view;
uniform mat4 u_projection;
uniform vec3 u_light_position;
uniform vec3 u_light_spec_position;

varying vec3  v_color;
varying vec4  v_eye_position;
varying float v_radius;
varying vec3  v_light_direction;

void main() {
    // 计算球体表面的点
    vec2 texcoord = gl_PointCoord * 2.0 - vec2(1.0);
    float x = texcoord.x;
    float y = texcoord.y;
    float d = 1.0 - x*x - y*y;
    
    if (d <= 0.0)
        discard;

    float z = sqrt(d);
    vec4 pos = v_eye_position;
    pos.z += v_radius * z;
    vec3 pos2 = pos.xyz;
    pos = u_projection * pos;
    
    vec3 normal = vec3(x, y, z);
    float diffuse = clamp(dot(normal, v_light_direction), 0.0, 1.0);

    // 高光反射计算
    vec3 M = pos2.xyz;
    vec3 O = v_eye_position.xyz;
    vec3 L = u_light_spec_position;
    vec3 K = normalize(normalize(L - M) + normalize(O - M));
    float specular = clamp(pow(abs(dot(normal, K)), 40.), 0.0, 1.0);
    
    vec3 v_light = vec3(1., 1., 1.);
    gl_FragColor.rgba = vec4(.15*v_color + .55*diffuse * v_color
                            + .35*specular * v_light, 1.0);
}

步骤2：创建分子可视化画布类

import numpy as np
from vispy import gloo, app
from vispy.util.transforms import perspective, translate, rotate
from vispy.io import load_data_file

class MolecularViewer(app.Canvas):
    def __init__(self):
        # 初始化画布
        app.Canvas.__init__(self, title='分子可视化器', 
                           keys='interactive', size=(800, 600))
        
        self.ps = self.pixel_scale
        self.translate = 40
        
        # 创建着色器程序
        self.program = gloo.Program(vertex, fragment)
        
        # 初始化变换矩阵
        self.view = translate((0, 0, -self.translate))
        self.model = np.eye(4, dtype=np.float32)
        self.projection = np.eye(4, dtype=np.float32)
        self.apply_zoom()
        
        # 加载分子数据
        self.load_molecule_data()
        self.setup_lighting()
        
        # 设置旋转参数
        self.theta = 0
        self.phi = 0
        
        # 配置OpenGL状态
        gloo.set_state(depth_test=True, clear_color='black')
        
        # 启动自动旋转计时器
        self.timer = app.Timer('auto', connect=self.on_timer, start=True)
        self.show()

    def load_molecule_data(self):
        """加载分子数据并配置缓冲区"""
        try:
            # 从VisPy数据仓库加载示例数据
            fname = load_data_file('molecular_viewer/micelle.npz')
            molecule = np.load(fname)['molecule']
            self._nAtoms = molecule.shape[0]
            
            # 提取原子坐标、颜色和半径
            self.coords = molecule[:, :3]          # x, y, z坐标
            self.atomsColours = molecule[:, 3:6]   # RGB颜色
            self.atomsScales = molecule[:, 6]      # 原子半径比例
            
            # 创建结构化数据数组
            data = np.zeros(self._nAtoms, [
                ('a_position', np.float32, 3),
                ('a_color', np.float32, 3),
                ('a_radius', np.float32)
            ])
            
            data['a_position'] = self.coords
            data['a_color'] = self.atomsColours
            data['a_radius'] = self.atomsScales * self.ps
            
            # 绑定顶点缓冲区
            self.program.bind(gloo.VertexBuffer(data))
            
        except Exception as e:
            print(f"数据加载错误: {e}")
            # 创建示例数据作为备选
            self.create_sample_data()

    def create_sample_data(self):
        """创建示例分子数据"""
        self._nAtoms = 100
        # 随机生成原子位置
        self.coords = np.random.uniform(-10, 10, (self._nAtoms, 3))
        # 随机颜色
        self.atomsColours = np.random.uniform(0, 1, (self._nAtoms, 3))
        # 固定半径比例
        self.atomsScales = np.ones(self._nAtoms) * 0.5
        
        data = np.zeros(self._nAtoms, [
            ('a_position', np.float32, 3),
            ('a_color', np.float32, 3),
            ('a_radius', np.float32)
        ])
        
        data['a_position'] = self.coords
        data['a_color'] = self.atomsColours
        data['a_radius'] = self.atomsScales * self.ps
        
        self.program.bind(gloo.VertexBuffer(data))

    def setup_lighting(self):
        """配置光照参数"""
        self.program['u_model'] = self.model
        self.program['u_view'] = self.view
        self.program['u_light_position'] = [0., 0., 2.]  # 平行光方向
        self.program['u_light_spec_position'] = [-5., 5., -5.]  # 高光位置

    def on_timer(self, event):
        """定时器回调 - 实现自动旋转"""
        self.theta += 0.5
        self.phi += 0.3
        self.model = np.dot(rotate(self.theta, (0, 0, 1)),
                          rotate(self.phi, (0, 1, 0)))
        self.program['u_model'] = self.model
        self.update()

    def on_resize(self, event):
        """处理画布大小变化"""
        width, height = event.physical_size
        gloo.set_viewport(0, 0, width, height)
        self.projection = perspective(25.0, width / float(height), 2.0, 100.0)
        self.program['u_projection'] = self.projection

    def apply_zoom(self):
        """应用缩放变换"""
        width, height = self.physical_size
        gloo.set_viewport(0, 0, width, height)
        self.projection = perspective(25.0, width / float(height), 2.0, 100.0)
        self.program['u_projection'] = self.projection

    def on_mouse_wheel(self, event):
        """鼠标滚轮缩放控制"""
        self.translate -= event.delta[1]
        self.translate = max(-1, self.translate)
        self.view = translate((0, 0, -self.translate))
        self.program['u_view'] = self.view
        self.update()

    def on_key_press(self, event):
        """键盘控制"""
        if event.text == ' ':  # 空格键暂停/继续旋转
            if self.timer.running:
                self.timer.stop()
            else:
                self.timer.start()
        elif event.text == 'r':  # R键重置视角
            self.translate = 40
            self.view = translate((0, 0, -self.translate))
            self.program['u_view'] = self.view
            self.update()

    def on_draw(self, event):
        """渲染循环"""
        gloo.clear()
        self.program.draw('points')

步骤3：在Jupyter中集成和使用

# 在Jupyter单元格中创建和显示可视化器
viewer = MolecularViewer()
viewer  # 这将自动显示在输出区域

# 控制方法示例
viewer.timer.stop()    # 暂停旋转
viewer.timer.start()   # 继续旋转

# 交互控制说明
print("""
交互控制指南：
- 鼠标拖动：旋转视角
- 鼠标滚轮：缩放
- 空格键：暂停/继续自动旋转
- R键：重置视角
""")

高级功能扩展

1. 自定义分子数据加载

def load_pdb_file(self, pdb_content):
    """从PDB文件内容加载分子数据"""
    atoms = []
    colors = []
    radii = []
    
    for line in pdb_content.split('\n'):
        if line.startswith('ATOM') or line.startswith('HETATM'):
            # 解析原子信息
            x = float(line[30:38])
            y = float(line[38:46])
            z = float(line[46:54])
            element = line[76:78].strip()
            
            atoms.append([x, y, z])
            colors.append(self.get_atom_color(element))
            radii.append(self.get_atom_radius(element))
    
    return np.array(atoms), np.array(colors), np.array(radii)

def get_atom_color(self, element):
    """根据元素类型返回颜色"""
    color_map = {
        'C': [0.4, 0.4, 0.4],   # 碳 - 灰色
        'O': [1.0, 0.0, 0.0],   # 氧 - 红色
        'N': [0.0, 0.0, 1.0],   # 氮 - 蓝色
        'H': [1.0, 1.0, 1.0],   # 氢 - 白色
        'S': [1.0, 1.0, 0.0],   # 硫 - 黄色
        'P': [1.0, 0.5, 0.0],   # 磷 - 橙色
    }
    return color_map.get(element, [0.5, 0.5, 0.5])  # 默认灰色

def get_atom_radius(self, element):
    """根据元素类型返回半径"""
    radius_map = {
        'C': 0.7, 'O': 0.66, 'N': 0.65, 
        'H': 0.25, 'S': 1.0, 'P': 1.0
    }
    return radius_map.get(element, 0.5)

2. 添加化学键可视化

def add_bond_visualization(self):
    """添加化学键的线段可视化"""
    # 简化的化学键检测（实际应用中需要更复杂的算法）
    bond_vertices = []
    bond_colors = []
    
    for i in range(len(self.coords)):
        for j in range(i+1, len(self.coords)):
            dist = np.linalg.norm(self.coords[i] - self.coords[j])
            if dist < 2.0:  # 简单的距离阈值
                # 添加键的起点和终点
                bond_vertices.extend([self.coords[i], self.coords[j]])
                # 使用相邻原子的平均颜色
                avg_color = (self.atomsColours[i] + self.atomsColours[j]) / 2
                bond_colors.extend([avg_color, avg_color])
    
    if bond_vertices:
        bond_data = np.zeros(len(bond_vertices), [
            ('a_position', np.float32, 3),
            ('a_color', np.float32, 3)
        ])
        bond_data['a_position'] = bond_vertices
        bond_data['a_color'] = bond_colors
        
        # 创建化学键的着色器程序
        self.bond_program = gloo.Program(bond_vertex_shader, bond_fragment_shader)
        self.bond_program.bind(gloo.VertexBuffer(bond_data))

3. 添加选择和高亮功能