LAMMPS Python 接口 Scatter/Gather 原子数据操作完全指南：从 gather_atoms 到 scatter

科研
科学计算
高性能计算

【免费下载链接】lammps

Public development project of the LAMMPS MD software package

项目地址：https://gitcode.com/gh_mirrors/la/lammps

导读

本文围绕 LAMMPS 官方文档中 "Scatter/gather operations" 一节展开，系统讲解通过 Python 接口在模拟过程中批量读取（gather）与写回（scatter）原子级属性（坐标、电荷、类型、力等）的完整方法。读完本文，你将掌握gather_atoms、gather_atoms_concat、gather_atoms_subset、scatter_atoms、scatter_atoms_subset五个核心方法的参数语义与调用规范，理解它们与extract_atom的本质区别，学会用 ctypes 向量正确构造与传递数据，并能结合 src/library.cpp 的源码理解其在 MPI 并行环境下的底层实现原理，从而在分子动力学工作流中实现"读数据、改数据、写回 LAMMPS、继续跑"的闭环控制。

一、五兄弟 API 总览

gather/scatter 系列是 LAMMPS Python 模块（python/lammps/core.py）提供的五对一的操作，全部是 C 库接口的 Python 封装。它们解决同一个核心问题：并行模拟中每个处理器只持有部分原子（nlocal），如何让 Python 侧拿到全局一致的原子数据，又如何把 Python 侧计算好的数据分发回正确的原子。

方法	作用	数据顺序	适用场景
`lmp.gather_atoms(name, dtype, count)`	按原子 ID 排序收集全部原子的属性	有序（按 atom ID）	需要全局有序视图，如坐标 x
`lmp.gather_atoms_concat(name, dtype, count)`	按各处理器拼接收集全部原子属性	无序（处理器分块拼接）	只需全量数据、不关心顺序
`lmp.gather_atoms_subset(name, dtype, count, ndata, ids)`	收集指定 ID 子集原子的属性	与 ids 列表同序	只关心少数原子的状态
`lmp.scatter_atoms(name, dtype, count, data)`	将有序数据分发回全部原子	要求按 atom ID 连续有序	改完坐标/电荷后写回
`lmp.scatter_atoms_subset(name, dtype, count, ndata, ids, data)`	将数据分发回指定 ID 子集	与 ids 列表同序	只修改少数原子

这些方法在 python/lammps/core.py 中的封装非常直白：将name编码为字节串后，根据dtype分配对应长度的 ctypes 向量（c_int或c_double），再调用底层self.lib.lammps_gather_atoms(...)等 C 函数。这意味着最终的数据类型、内存布局都由 ctypes 决定，理解这一点是正确使用的前提。

二、gather 三兄弟：把分散在各处理器上的原子数据收集到 Python

2.1 gather_atoms：按原子 ID 排序的全局向量

data = lmp.gather_atoms("x", 1, 3) # 收集所有原子的坐标 data = lmp.gather_atoms("q", 1, 1) # 收集所有原子的电荷 data = lmp.gather_atoms("type", 0, 1) # 收集所有原子的类型

name是 LAMMPS 内部通过atom->extract()可识别的属性名，常见有"x"（坐标）、"v"（速度）、"f"（力）、"q"（电荷）、"type"（类型）、"id"（ID）、"image"（镜像标志）等。三个参数的含义：

dtype：0 表示返回int值，1 表示返回double值；
count：每个原子对应的取值个数。标量属性（如type、q）传 1；数组属性（如x、v、f）传 3；
name="image"时传count=3有特殊含义：LAMMPS 内部把原子的 x/y/z 镜像标志打包在单个整数里，此时gather_atoms会把它解包成三个分量返回（见下文源码解析）。

返回的data是一个长度为count × natoms的一维 ctypes 向量，按原子 ID 从小到大排列：x[0..2]是 ID=1 的原子坐标，x[3..5]是 ID=2 的原子坐标，依此类推。

2.2 gather_atoms_concat：无排序要求的全量收集

data = lmp.gather_atoms_concat("x", 1, 3) # 每个处理器原子块的简单拼接

与gather_atoms不同，gather_atoms_concat不做按 ID 排序，直接按各处理器nlocal原子的本地顺序拼接（底层用MPI_Allgatherv）。它的价值在于：当你不关心顺序、只需要拿到全部数据做统计聚合（如计算全局均值）时，它不需要原子 ID 连续，适用面更广。如果你还需要每个原子的 ID，可以再调用一次gather_atoms_concat("id", 0, 1)自行对齐。

2.3 gather_atoms_subset：只收集你关心的原子

ids = [10, 20, 30, 40] # 感兴趣的原子 ID data = lmp.gather_atoms_subset("x", 1, 3, len(ids), ids)

ndata是要查询的原子个数，ids是这些原子的 ID 列表（长度必须等于ndata）。返回的向量长度为count × ndata，顺序与 ids 列表一一对应：data[0..2]是 ID=10 的坐标，data[3..5]是 ID=20 的坐标。适合做单原子追踪、局部探针等场景，通信量也远小于全量收集。

三、scatter 两兄弟：把 Python 数据写回 LAMMPS

3.1 scatter_atoms：写回全部原子

lmp.scatter_atoms("x", 1, 3, data) # data 必须按原子 ID 连续有序

这是gather_atoms的逆操作：把长度为count × natoms、按原子 ID 排序的一维向量分发到各个原子（原子可能由不同处理器拥有）。参数dtype、count的语义与 gather 完全一致，要求传入的data是 ctypes 的c_int或c_double向量。

3.2 scatter_atoms_subset：只写回部分原子

lmp.scatter_atoms_subset("x", 1, 3, ndata, ids, data)

ndata与ids的语义与gather_atoms_subset相同，data的长度为count × ndata，顺序与ids一致。典型用法是先 gather 一个子集、修改、再 scatter 回去，实现"定点修改"。

四、gather 与 extract_atom 的四个关键区别

文档明确强调，gather_atoms返回的数据结构与extract_atom("x")完全不同，两者在四个维度上存在本质差异：

数据结构不同：gather_atoms("x")返回一维向量，索引为x[i]；extract_atom("x")返回二维数组，索引为x[i][j]。
原子顺序不同：gather_atoms()按原子 ID 排序；extract_atom()返回的是各处理器本地原子顺序，不排序。
原子范围不同：gather_atoms()返回全模拟的所有原子；extract_atom()只返回当前处理器本地的原子（nlocal），在并行运行时不同处理器上看到的内容不一样。
内存语义不同（最重要）：gather_atoms()返回的是内部数据的拷贝；extract_atom()返回的数组直接指向 LAMMPS 内部内存。这意味着直接对extract_atom()数组赋值就能就地修改 LAMMPS 内部状态；而对gather_atoms()向量赋值后，必须再调用scatter_atoms("x")才能把修改生效。

这一区别决定了编程模式：需要修改模拟状态时，要么用extract_atom直接写，要么走"gather → 改 → scatter"三步流程。后者在需要全局有序视图或跨处理器统一处理时更干净。

五、参数细节：dtype、count 与 image 的特殊打包

参数	取值	含义
`dtype`	0	数据为`int`类型（如`type`、`id`）
`dtype`	1	数据为`double`类型（如`x`、`v`、`f`、`q`）
`count`	1	每个原子 1 个值（标量属性）
`count`	3	每个原子 3 个值（数组属性）

name="image"搭配count=3是唯一的特例：LAMMPS 内部将原子的三个镜像标志打包存储为一个 32 位整数，传count=3时 gather 会自动解包为三个分量（x/y/z 方向各自的镜像序号），scatter 时则反向打包回单个整数。如果对image传count=1，得到的则是原始打包整数。

六、ctypes 向量构造：scatter 的必修课

对于 scatter 方法，传入的data必须是 ctypes 的c_int或c_double向量（与dtype对应）。文档给出了三种标准构造方式：

6.1 分配后逐项赋值

from ctypes import c_double natoms = lmp.get_natoms() # 原子总数 n3 = 3 * natoms x = (n3 * c_double)() # 分配 3*natoms 个 double 的向量 x[0] = ... # ID=1 的 x 坐标 x[1] = ... # ID=1 的 y 坐标 x[2] = ... # ID=1 的 z 坐标 x[3] = ... # ID=2 的 x 坐标 ... x[n3-1] = ... # ID=natoms 的 z 坐标 lmp.scatter_atoms("x", 1, 3, x)

6.2 构造函数初始化

from ctypes import c_double natoms = 2 n3 = 3 * natoms x = (n3 * c_double)(0.0, 0.0, 0.0, 1.0, 1.0, 1.0) # 构造函数中直接给初值 lmp.scatter_atoms("x", 1, 3, x)

6.3 从 Python list 转换

coords = [1.0, 2.0, 3.0, -3.0, -2.0, -1.0] x = (c_double * len(coords))(*coords)

捷径：gather 方法返回的本身就是 ctypes 向量，所以最常见的做法是直接修改 gather 的返回值再 scatter 回去，无需重新构造。

七、完整实战：修改坐标并继续运行

仓库自带的示例 python/examples/simple.py 展示了最典型的应用——把第 1 个原子的 x 坐标扰动后写回并继续跑：

# 从 LAMMPS 收集坐标与速度 x = lmp.gather_atoms("x", 1, 3) v = lmp.gather_atoms("v", 1, 3) # 扰动第 1 个原子的 x 坐标 epsilon = 0.1 x[0] += epsilon # 写回 LAMMPS lmp.scatter_atoms("x", 1, 3, x) # 用修改后的坐标继续运行 lmp.command("run 1")

python/examples/demo.py 也演示了同一模式：xc = lmp.gather_atoms("x",1,3)后xc[0] = xc[0] + 1.0，再lmp.scatter_atoms("x",1,3,xc)，随后通过extract_atom("x")验证修改已经生效。注意其中extract_atom索引为x[0][0]（二维），而gather_atoms索引为xc[0]（一维），正是第四节所述区别的直观体现。

这套"gather → 改 → scatter → run"模式可用于实现任意粒子级的外部控制：添加随机扰动、强制设定某原子位置、根据外部势函数改写速度等，是 LAMMPS 与外部算法（如强化学习、实时可视化、采样增强）对接的常用桥梁。

八、源码级原理：gather 的 MPI 合并技巧与 scatter 的 atom map 依赖

深入 src/library.cpp 可以看到每个方法背后的并行实现逻辑。

8.1 gather_atoms：MPI_Allreduce 求和实现"按 ID 排序"

lammps_gather_atoms（src/library.cpp）的实现技巧十分巧妙：每个处理器申请一个长度为count*natoms的copy向量并清零，然后遍历自己拥有的 nlocal 个原子，把每个原子的值填入copy[count*(tag[i]-1)]对应的槽位（即用原子 ID 定位数组下标），最后执行一次MPI_Allreduce(copy, data, ..., MPI_SUM, ...)。因为每个槽位只被唯一一个处理器写入，其余处理器在该槽位贡献 0，求和后data就是完整的、按原子 ID 排序的全局向量。这也解释了文档的隐含前提——atom ID 必须存在且连续（代码中通过tag_consecutive()检查），否则tag[i]-1无法安全用作数组下标。

8.2 gather_atoms_concat：MPI_Allgatherv 直接拼接

lammps_gather_atoms_concat（src/library.cpp）走的是另一条路径：先通过MPI_Allgather收集各处理器的nlocal，构造recvcounts/displs后调用MPI_Allgatherv把各处理器本地块直接拼进data。所以返回顺序是各处理器原子块的拼接顺序，而非 ID 顺序，且只要求 atom ID 存在，不要求连续。

8.3 image 标志的解包/打包位运算

无论是 gather 还是 scatter，源码都对name="image"且count=3的情形做了特殊位运算处理。解包（gather）时：

copy[offset++] = (image & IMGMASK) - IMGMAX; copy[offset++] = ((image >> IMGBITS) & IMGMASK) - IMGMAX; copy[offset++] = ((image >> IMG2BITS) & IMGMASK) - IMGMAX;

打包（scatter）时反向执行：

int image = dptr[offset++] + IMGMAX; image += (dptr[offset++] + IMGMAX) << IMGBITS; image += (dptr[offset++] + IMGMAX) << IMG2BITS;

这印证了第六节所述：count=3与image的组合会在 C 层完成单整数到三分量的自动转换。

8.4 subset 与 scatter 对 atom map 的要求

从源码限制条件看，不同方法的约束各不相同：

gather_atoms：要求 atom ID 存在且连续；
gather_atoms_concat：只要求 atom ID 存在；
gather_atoms_subset与scatter_atoms_subset：要求 atom ID 存在且启用了 atom map（代码检查map_style == Atom::MAP_NONE时报错），因为需要通过atom->map(id)根据任意 ID 反查本地索引；
scatter_atoms：要求 ID 存在、连续、且启用 atom map（scatter 遍历所有 ID 调用map(i+1)定位原子）。

在 Python 侧构造 LAMMPS 实例时，应确保使用lmp = lammps(..., args=["-log","none"])或通过输入脚本atom_modify map yes开启 atom map，以满足 subset/scatter 系列的前置条件。LAMMPS 官方 Python 模块文档见 doc/src/Python_run.rst，其中列出了完整的 Python 接口模块索引。

九、限制与注意事项

综合 src/library.cpp 中各函数的 Doxygen 注释（如 lammps_gather_atoms 的说明），使用 gather/scatter 系列需要注意以下限制：

LAMMPS_BIGBIG 不兼容：所有五个函数在编译选项-DLAMMPS_BIGBIG（64 位 atom ID + 64 位原子数）下都会直接报错，要求原子总数不超过 2147483647（32 位有符号整数上限）。
atom ID 要求：tag_enable必须为 1；gather_atoms/scatter_atoms还要求 ID 连续；subset 系列必须启用 atom map。
数据长度必须精确：data数组长度必须是count × natoms（subset 为count × ndata），少了会越界、多了会浪费，务必用lmp.get_natoms()查询原子总数。
dtype 与 ctypes 类型必须匹配：dtype=0配c_int向量，dtype=1配c_double向量，混用会导致数据解析错误。
gather 返回的是拷贝：修改后必须显式 scatter 回 LAMMPS 才生效（区别于extract_atom的零拷贝语义）。

仓库的测试套件也对这套机制做了覆盖验证，可参考 unittest/fortran/test_fortran_gather_scatter.f90 与 unittest/fortran/wrap_gather_scatter.cpp，其中用具体原子数、属性名与期望值逐一校验 gather/scatter 的往返一致性，是学习正确调用方式的绝佳对照材料。

十、小结

gather/scatter 系列是 LAMMPS Python 接口中最实用的"全局数据通道"：gather_atoms提供按 ID 排序的全量快照，gather_atoms_concat提供无序全量拼接，gather_atoms_subset支持定点查询，scatter_atoms/scatter_atoms_subset完成反向写回。理解其参数语义（dtype/count）、与extract_atom的差异、ctypes 向量的构造方式，以及背后 MPI 合并与 atom map 的依赖，就能在各类自定义工作流中安全高效地操纵任意原子属性。