用途:nvalchemi 用图结构表示原子体系,核心是 AtomicData(单个分子、晶体等)与 Batch(多个 AtomicData 拼接为连续张量,便于 GPU 批量计算)。两者均为 Pydantic BaseModel 子类,并通过 DataMixin 支持设备与 dtype 操作。
主要能力:AtomicData 必填 positions [n_nodes,3] 与 atomic_numbers [n_nodes],可选节点、边、稠密邻居及体系级字段(cell、pbc、energy、stress、virial、dipole、charge、graph_embeddings 等),自定义数据可放入 info 字典。支持 from_atoms 从 ASE Atoms 构建,提供 num_nodes、num_edges、device、dtype、chemical_hash 等属性,以及字典式读写、add_node_property/add_edge_property/add_system_property、to/cpu/cuda/clone、model_dump 序列化;相等性按 chemical_hash 判断。Batch 通过 from_data_list 构建,也可用 Batch.empty 预分配缓冲区;提供 num_graphs、batch_idx、batch_ptr、num_nodes_list 等尺寸属性,支持索引取单个图或子批、to_data_list 还原、add_key 增改键、append/append_data 拼接、set_group_layout 分组,以及 put/defrag/zero 等预分配缓冲操作,并支持 to/cpu/cuda/clone/contiguous/pin_memory 与 model_dump 序列化。
使用方式:从 nvalchemi.data 导入 AtomicData 与 Batch;构建时传入张量字段,批量时用 Batch.from_data_list(data_list),高性能场景用 Batch.empty(num_systems, num_nodes, num_edges, template) 并配合 put/defrag。分布式场景可用 send/recv 及 isend/irecv,接收端需提供 template 以确定键、dtype 与分组结构;0 图哨兵批仅传元数据头。
适用场景:构建原子体系数据集、邻居列表与周期性体系、批量前向计算、流式动力学的高吞吐缓冲,以及 torch.distributed 点对点通信。
限制:除 positions 与 atomic_numbers 外字段均可选;分组批不支持 put 与 append_data,append 要求两侧同为分组或同为未分组;group_layout 缓存在原地修改 group_idx 后不会自动失效,需重新调用 set_group_layout 重建。其余细节资料未提供。