FlyCrysNet

FlyCrysNet

FlyCrysNet 是一个基于 PyTorch 和 PyTorch Geometric 的独立晶体带隙预测模型。该模型将周期性晶体图编码器与源自完整果蝇雄性中枢神经系统连接组（MaleCNS）的稀疏潜在消息传递相结合。

晶体编码器包括原子嵌入、径向基距离编码、四层边门控图卷积和全局平均池化。池化后的晶体表示被投影到 256 个潜在节点，每个节点包含 8 个通道。随后，模型根据 MaleCNS 有向加权拓扑执行自适应门控的稀疏传播。拓扑门和突触权重幂均可学习，最终通过注意力池化和全局跳跃路径预测晶体带隙。

脚本还包含四项必要的拓扑对照：仅自连接、随机拓扑、度保持随机拓扑和无权重 MaleCNS 拓扑。所有变体采用相同的固定数据划分和训练设置，并分别根据验证集 RMSE 独立选择最佳检查点。

环境依赖

1

pip install torch torch-geometric pymatgen numpy pandas openpyxl scikit-learn matplotlib tqdm pyarrow

仅当 USE_OPTUNA = True 时需要安装 Optuna：

1

pip install optuna

数据准备

在脚本所在目录中准备以下文件和文件夹：

1
2
3
4
5
6
7
8

FlyCrysNet_v3.py
data.xlsx
cif/
|-- 1.cif
|-- 2.cif
|-- 3.cif
fly_connectome/
|-- connectome-weights-male-cns-v1.0-minconf-0.5.feather

connectome-weights-male-cns-v1.0-minconf-0.5.feather 文件约为 1.1 GB。请保持文件名不变，并将其放入 ./fly_connectome/。

脚本读取 data.xlsx 的前两列：

脚本会把第一列的值转换为整数文件名。例如，第一列中的 1 对应 ./cif/1.cif。

MaleCNS 拓扑

第一次运行时，脚本会对 Feather 文件执行两遍扫描。第一遍根据加权度对神经元排序并选取前 256 个神经元，第二遍提取这些神经元构成的有向加权诱导子图。脚本对正突触权重执行 log1p 变换，加入自连接，对稀疏拓扑进行归一化并保存缓存。

拓扑缓存保存为：

1

FlyCrysNet_v3/cache/MaleCNS_top256.npz

当源 Feather 文件、文件修改时间和 CONNECTOME_NODES 均未改变时，后续运行会直接复用该缓存。因此，第一次运行所需时间会明显长于后续运行。

使用方法

1

python FlyCrysNet_v3.py

脚本会自动完成 CIF 检查、周期性晶体图构建、固定的 80/10/10 训练集/验证集/测试集划分、目标标准化、模型训练、早停、评估、绘图和拓扑消融实验。当 CUDA 可用时自动使用 GPU 和 bfloat16 自动混合精度，否则使用 CPU。

将 RUN_ABLATIONS = False 可设置为仅运行完整加权 MaleCNS 模型。将 USE_OPTUNA = True 可启用可选的超参数搜索。

所有结果保存在由 MODEL_NAME 和 RUN_VERSION 确定的目录中。按照默认设置，输出目录为：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11

FlyCrysNet_v3/
|-- FlyCrysNet_best.pt
|-- figure/
|-- dat/
|-- table/
|-- log/
|-- split/
|-- cache/
|   |-- MaleCNS_top256.npz
|   `-- crystal_graphs/
`-- ablation_checkpoints/

输出内容包括训练集、验证集和测试集的 MAE、RMSE 与 R2，奇偶图及其数据，RMSE 训练曲线及其数据，固定数据划分，完整模型的最佳权重，消融模型权重，拓扑对照指标、预测结果、训练曲线与图像，晶体图与 MaleCNS 拓扑缓存，以及完整训练日志。

引用

MaleCNS 连接组文献：

• Berg S, Beckett I R, Costa M, et al. Sexual dimorphism in the complete Drosophila male central nervous system connectome. Cell, 2026, 189(18): 5504-5526. e15

本工作：

论文正式发表后补充。