NumPy 2.x 科研实战教程
基础要求:Python 中级,熟悉列表 / 字典 / 函数
环境:Windows + Miniconda + Python 3.14.4 + NumPy 2.4.4 + VS Code
本文所有代码均可直接保存为.py脚本运行,注释为中文。
1. 简介与定位:为什么计算化学离不开 NumPy
做 DFT(VASP / CP2K / Quantum ESPRESSO)、催化材料(CO2 还原、MOF、单原子催化)和机器学习势函数(MLIP)时,你面对的数据几乎都是数值数组:
- 结构坐标:
N × 3矩阵(N 个原子的 x/y/z) - 晶格矢量:
3 × 3矩阵(三条晶格基矢) - 受力 / 能量:
N × 3的力数组 + 标量能量 - 势能面(PES):网格点上的一组能量值
- 径向分布函数(RDF):距离直方图
这些数据的共同特征:规模大、类型统一(浮点)、需要做批量逐元素运算。Python 原生 list 存一亿个数不仅慢,而且没有"矩阵乘法""广播""向量化"这些科学计算的核心能力。
NumPy 提供:
ndarray:高效、同质(homogeneous)、多维的数组对象。- 向量化运算:
a + b直接在 C 层对一整个数组做加法,而不是 Python 循环。 - 广播(broadcasting):不同形状的数组可以自动对齐运算。
np.linalg:矩阵分解、求解、特征值,正是 DFT 里算力常数矩阵、晶格应变需要的。- 统一的随机数 API:可复现的蒙特卡洛 / 分子动力学初始速度采样。
一句话:NumPy 是 ASE、pymatgen、scikit-learn、torch 的共同底座。把它练熟,后面学这些库会非常轻松。
2. 安装与版本检查
import numpy as np
print(np.__version__) # 期望输出 2.x.x,例如 2.4.4
print(np.__config__.show()) # 查看 BLAS / LAPACK 后端(影响 linalg 性能)
Miniconda 下若需升级:
conda install -c conda-forge numpy
# 或
pip install --upgrade numpy
重要:本文基于 NumPy 2.x 编写。2.0 相对 1.x 有几处破坏性变更(后文会反复强调):
np.float_/np.int_/np.complex_/np.float等旧别名已被移除,统一用np.float64/np.int64/np.complex128。np.array(..., copy=False)语义变化:需要复制时会直接报错而不是静默复制。- 随机数推荐使用
np.random.default_rng(),旧的np.random.seed/np.random.rand属"遗留全局状态"API,应避免在新代码中使用。np.trapz已更名为np.trapezoid;np.alltrue/np.sometrue已移除,用np.all/np.any。
3. ndarray 核心概念
3.1 创建与基本属性
import numpy as np
# 从 Python list 创建
a = np.array([1.0, 2.0, 3.0])
print(a) # [1. 2. 3.]
# 三维:2 个结构,每个 3 个原子,每个原子 x/y/z
coords = np.array([
[[0.0, 0.0, 0.0], [1.2, 0.0, 0.0], [0.0, 1.2, 0.0]],
[[0.1, 0.1, 0.1], [1.3, 0.1, 0.1], [0.1, 1.3, 0.1]],
])
print("shape:", coords.shape) # (2, 3, 3)
print("ndim :", coords.ndim) # 3
print("size :", coords.size) # 18(元素总数)
print("dtype:", coords.dtype) # float64
print("itemsize:", coords.itemsize) # 8 字节/元素(float64)
3.2 dtype:为什么必须"同质"
ndarray 里所有元素类型相同,这决定了它可以被连续存储、被 C/Fortran 高效处理。常见 dtype:
| 用途 | 标量类型 | 数组 dtype 写法 |
|---|---|---|
| 双精度浮点(默认) | np.float64 | np.float64 或 "float64" |
| 单精度浮点(GPU/大规模) | np.float32 | np.float32 |
| 整数 | np.int64 / np.int32 | np.int64 |
| 复数 | np.complex128 | np.complex128 |
| 布尔 | np.bool_ | np.bool_ 或 bool |
NumPy 2.x 中,请不要再写
np.float_、np.int_、np.float——这些别名已移除,会直接AttributeError。
z = np.array([1, 2, 3], dtype=np.float64) # 强制浮点
c = np.array([1 + 2j, 3 - 4j]) # 自动推断为 complex128
b = np.array([True, False, True], dtype=np.bool_)
# 类型转换用 .astype,而不是旧式的 np.float64(x)
x = np.array([1.9, 2.7])
x_int = x.astype(np.int64) # 截断,不四舍五入 → [1 2]
3.3 内存连续性与 strides
NumPy 数组在内存里是一段连续区域,"形状"只是对这段内存的一种视图。
a = np.arange(12).reshape(3, 4)
print(a.strides) # (32, 8):沿第 0 轴跳 32 字节,沿第 1 轴跳 8 字节
print(a.flags["C_CONTIGUOUS"]) # True,行优先(C order)连续
print(a.flags["F_CONTIGUOUS"]) # False,列优先(Fortran order)不连续
- C order(行优先):最后面的轴变化最快,Python/NumPy 默认。英文资料里的
order='C'。 - F order(列优先):第一个轴变化最快,对应 Fortran/VASP 某些 I/O。
大多数时候你不需要关心顺序,但转置会改变 strides 而不移动数据(见第 8 节),理解它能帮你判断"这是视图还是副本"。
3.4 与 Python list 的区别
list | ndarray | |
|---|---|---|
| 元素类型 | 可混合 | 必须同质 |
| 内存 | 存引用,分散 | 连续一段内存 |
| 运算语义 | + 是拼接 | + 是逐元素相加 |
| 数学运算 | 需要循环 | 向量化,C 层执行 |
| 维度 | 靠嵌套 | 真正的多维,有 shape |
l = [1, 2, 3]
print(l + l) # [1, 2, 3, 1, 2, 3] 拼接
a = np.array([1, 2, 3])
print(a + a) # [2 4 6] 逐元素相加
4. 数组创建
4.1 常量数组
np.zeros((3, 4)) # 全 0
np.ones((2, 3)) # 全 1
np.full((2, 3), 7.5) # 全 7.5
np.empty((100, 3)) # 只分配内存,不初始化(值随机,慎用但快)
np.eye(3) # 3×3 单位阵(晶格变换常用)
np.identity(4) # 4×4 单位阵
科研里常见的"空力数组"、"空坐标数组":
forces = np.zeros((96, 3)) # 96 个原子的力,先占位
cell = np.eye(3) * 10.0 # 10 Å 的立方晶格
4.2 序列数组
np.arange(0, 10, 2) # [0 2 4 6 8] (stop 不包含)
np.arange(10) # [0 1 ... 9]
np.linspace(0, 1, 11) # 11 个点,含两端 [0, 0.1, ..., 1.0]
np.linspace(-1, 1, 5) # 常用于反应坐标的扫描点
np.logspace(-2, 2, 5) # 对数等间隔,10^-2 ... 10^2
坑:
np.arange用浮点步长可能有精度误差(np.arange(0, 1, 0.1)的最后一个值可能是0.9而非0.9...)。需要精确等分数时一律用np.linspace。
4.3 随机数组(新 API)
rng = np.random.default_rng(42) # 42 是种子,保证可复现
rng.random((3, 3)) # [0, 1) 均匀分布
rng.uniform(-1, 1, size=(10, 3)) # [-1, 1) 均匀,适合 MD 初始速度
rng.normal(0, 1, size=(10, 3)) # 高斯分布(均值为 0,标准差 1)
rng.integers(0, 100, size=5) # 整数,[0, 100)
rng.permutation(10) # 0~9 的随机排列,打乱数据集用
rng.choice([1, 2, 3], size=2) # 从给定集合随机抽样
旧的
np.random.seed(42)+np.random.rand(3,3)也能用,但它是全局状态,多函数 / 多线程下难以控制。新代码统一用default_rng(),每次显式传入rng对象。
4.4 从已有数组派生
a = np.array([[1, 2], [3, 4]])
np.zeros_like(a) # 与 a 同形状的 0
np.ones_like(a)
np.full_like(a, -1)
np.copy(a) # 深拷贝
5. 索引与切片
5.1 基础切片与"视图"
a = np.arange(10)
print(a[2:5]) # [2 3 4]
print(a[::-1]) # 反转
print(a[::2]) # 步长 2 → [0 2 4 6 8]
# 多维:coords[帧, 原子, 轴]
coords = np.random.default_rng(0).random((5, 8, 3)) # 5 帧 × 8 原子 × xyz
first_frame = coords[0] # (8, 3) 第一帧
x_only = coords[:, :, 0] # (5, 8) 所有帧所有原子的 x
atom_0_traj = coords[:, 0, :] # (5, 3) 0 号原子的轨迹
切片返回视图(view):不复制数据,改视图会改原数组。
b = a[2:5]
b[0] = 999
print(a) # [0 1 999 3 4 ...] —— a 被改了!
这正是"高效"的来源:切片不搬内存。若想独立,用 .copy()。
5.2 布尔掩码
energies = np.array([-5.2, -5.8, -4.9, -6.1, -5.5])
mask = energies < -5.5 # 布尔数组
print(mask) # [False True False True False]
print(energies[mask]) # 筛选出更稳定的构型 [-5.8 -6.1]
print(np.sum(mask)) # 满足条件的个数 = 2
# 组合条件用 & | ~,且每个条件必须加括号
sel = energies[(energies < -5.0) & (energies > -5.6)]
布尔掩码常用在筛选稳定构型、剔除异常值、原子类型筛选(如只取 O 原子的坐标)。
5.3 花式索引(fancy indexing)
用一个整数数组/列表去取特定位置。
a = np.arange(10)
idx = [0, 3, 7]
print(a[idx]) # [0 3 7]
# 多维度:coords[帧索引列表, 原子索引列表]
coords = np.arange(24).reshape(4, 3, 2) # 4 帧 × 3 原子 × 2 坐标
print(coords[[0, 2]]) # 取第 0、2 帧 → (2, 3, 2)
print(coords[[0, 2], [1, 1]]) # (帧0原子1, 帧2原子1) → (2, 2)
花式索引返回副本(copy),改它不影响原数组——与切片相反。布尔掩码也返回副本。
5.4 视图 vs 副本 小结
| 操作 | 返回 |
|---|---|
基础切片 a[1:3] | 视图 |
.reshape / .ravel(多数情况) | 视图 |
.T / np.transpose | 视图 |
布尔掩码 a[mask] | 副本 |
花式索引 a[[0,2]] | 副本 |
.flatten() | 副本(永远) |
.astype | 副本(通常,dtype 变了) |
判断方法:a.base is not None 说明 a 是某个数组的视图;或用 np.shares_memory(a, b)。
x = np.arange(6)
y = x[1:4]
print(np.shares_memory(x, y)) # True
z = x[[1, 2, 3]]
print(np.shares_memory(x, z)) # False
6. 广播机制(重点)
广播是 NumPy 最强大也最容易踩坑的特性。规则只有三条,记住即可:
当两个数组形状不同时,从右往左逐维比较:
- 维度数不同 → 在形状较少的数组左边补 1。
- 某维相等,或其中一个是 1 → 兼容。
- 某维都不为 1 也不相等 → 报错。
兼容后,维度为 1 的轴会被虚拟拉伸到与另一个相同(不真正复制数据)。
6.1 直观图式
形状 (4, 3) + 形状 (3,)
↓ ↓
(4, 3) + (1, 3) # 左边补 1
↓ ↓ 虚拟复制 4 行
(4, 3) + (4, 3) → 逐元素相加,得到 (4, 3)
这正是"给每个原子坐标统一加上一个平移矢量"的场景:
pos = np.random.default_rng(1).random((96, 3)) # 96 原子坐标
shift = np.array([1.0, 0.0, 0.0]) # 沿 x 平移 1 Å
pos_shifted = pos + shift # (96, 3) + (3,) 广播
print(pos_shifted.shape) # (96, 3)
6.2 维度为 1 的轴被拉伸
a = np.arange(6).reshape(2, 3) # (2, 3)
b = np.array([10, 20, 30]) # (3,)
print(a + b)
# [[10 21 32]
# [13 24 35]]
列向量场景(需要显式 reshape(-1, 1) 或 [:, None]):
col = np.array([100, 200]).reshape(2, 1) # (2, 1)
print(a + col) # (2, 1) 广播成 (2, 3)
6.3 广播在科研里的三个高频用法
# 1) 中心化:每帧坐标减去该帧质心
coords = np.random.default_rng(2).random((10, 96, 3)) # 10 帧
center = coords.mean(axis=1, keepdims=True) # (10, 1, 3)
coords_c = coords - center # (10, 96, 3) - (10, 1, 3)
# 2) 归一化:每条能量减去平均、除以标准差(每个体系一行)
E = np.array([[-5.1, -5.3, -5.2],
[-4.0, -4.1, -3.9]])
E_std = (E - E.mean(axis=1, keepdims=True)) / E.std(axis=1, keepdims=True)
# 3) 笛卡尔坐标 = 分数坐标 @ 晶格矩阵(见第 14 节)
坑:
(2, 3)与(3,)能广播;但(2, 3)与(2,)不能直接广播(因为从右往左,3 与 2 不相等也不为 1)。此时需要(2, 1)。
7. 向量化与 ufunc
7.1 为什么快
NumPy 的运算(+ - * / **,np.sqrt 等)是 ufunc(universal function):在 C 层对数组逐元素执行,并利用 CPU 向量化指令 / BLAS,几乎没有 Python 解释器开销。
import time
N = 10_000_000
a = np.random.default_rng(0).random(N)
# 纯 Python 循环(慢)
t0 = time.perf_counter()
s = 0.0
for x in a:
s += x
t1 = time.perf_counter()
print("循环耗时:", t1 - t0)
# NumPy(快)
t0 = time.perf_counter()
s = np.sum(a)
t1 = time.perf_counter()
print("NumPy 耗时:", t1 - t0)
7.2 常用数学 ufunc(逐元素)
x = np.linspace(0, np.pi, 5)
np.sqrt(x) # 平方根
np.exp(x) # e^x
np.log(x + 1e-12) # 自然对数(加小量防止 log(0))
np.sin(x), np.cos(x), np.tan(x)
np.abs(x) # 绝对值(复数取模)
np.sign(x) # 符号
np.power(x, 2) # x^2,等价 x ** 2
np.floor(x), np.ceil(x), np.round(x) # 取整
DFT 里常用:能量差取 abs、Boltzmann 因子 np.exp(-E / (k_B * T))、误差函数 np.erf、径向基 np.exp(-alpha * r**2)。
7.3 ufunc 的 out 与 axis 参数
a = np.arange(6).reshape(2, 3)
np.multiply(a, 2, out=a) # 原地写回,省一次内存分配(注意:会改 a)
8. 形状操作
8.1 reshape / ravel / flatten
a = np.arange(12) # (12,)
b = a.reshape(3, 4) # (3, 4),返回视图
c = a.reshape(3, -1) # -1 自动推断 → (3, 4)
d = a.reshape(2, 3, 2) # (2, 3, 2)
r1 = a.ravel() # 视图(可能)
r2 = a.flatten() # 一定返回副本
DFT 里的典型应用:把 N×3 坐标矩阵"拍平"成 3N 长向量喂给某些优化器。
8.2 转置与轴重排
a = np.arange(24).reshape(2, 3, 4) # (帧, 原子, xyz)
a.T # 反转所有轴 → (4, 3, 2)
np.transpose(a, (2, 1, 0)) # 指定轴顺序
np.swapaxes(a, 0, 2) # 交换第 0、2 轴
np.moveaxis(a, 0, -1) # 把第 0 轴移到最后
8.3 增维 / 减维
x = np.arange(5)
x[:, None] # (5,) → (5, 1),等价 x.reshape(-1, 1)
x[None, :] # (5,) → (1, 5)
np.expand_dims(x, axis=0) # (1, 5)
np.expand_dims(x, axis=1) # (5, 1)
np.squeeze(np.zeros((1, 5, 1))) # 去掉所有长度为 1 的轴 → (5,)
8.4 拼接与堆叠
a = np.zeros((2, 3))
b = np.ones((2, 3))
np.concatenate([a, b], axis=0) # (4, 3) 上下拼
np.concatenate([a, b], axis=1) # (2, 6) 左右拼
np.vstack([a, b]) # 垂直堆叠,等价 axis=0
np.hstack([a, b]) # 水平堆叠,等价 axis=1
np.stack([a, b], axis=0) # (2, 2, 3) 增加一个新轴
np.dstack([a, b]) # (2, 3, 2) 沿第 2 轴堆叠
科研场景:把多个 OUTCAR 里的力拼成一个 (N帧, N原子, 3) 大数组 → 用 np.stack(而不是 concatenate)。
8.5 分割
a = np.arange(12)
np.split(a, 3) # 等分成 3 份
np.split(a, [3, 7]) # 在索引 3、7 处切 → [0:3], [3:7], [7:]
m = np.arange(16).reshape(4, 4)
np.vsplit(m, 2) # 按行切
np.hsplit(m, 2) # 按列切
典型用法:把 (总样本数, 3) 的数据按 80%/20% 切成训练集 / 测试集。
8.6 axis 参数的理解
axis 指定"沿哪根轴坍缩 / 操作"。记住:axis=k 就是把第 k 根轴消掉。
a = np.arange(24).reshape(2, 3, 4) # (帧, 原子, xyz)
a.sum() # 全部求和 → 标量
a.sum(axis=0) # (3, 4) 对"帧"求和 → 每个原子坐标求和
a.sum(axis=1) # (2, 4) 对"原子"求和
a.sum(axis=2) # (2, 3) 对"xyz"求和
一个记忆技巧:axis=k 的结果,就是把原来第 k 个维度"压缩掉"。
9. 数学与统计
9.1 聚合函数
energies = np.array([-5.2, -5.8, -4.9, -6.1, -5.5])
energies.sum() # 求和
energies.mean() # 平均
energies.std() # 标准差(默认 ddof=0,即总体标准差)
energies.var() # 方差
energies.min() # 最小值(最稳定构型能量)
energies.max()
np.median(energies) # 中位数
np.percentile(energies, 25) # 第 25 百分位
# 沿轴聚合
coords = np.random.default_rng(3).random((10, 96, 3))
coords.mean(axis=0) # (96, 3) 平均结构
coords.std(axis=0) # (96, 3) 每个原子的均方位移
9.2 argmin / argmax:找最稳定构型
energies = np.array([-5.2, -5.8, -4.9, -6.1, -5.5])
i_min = np.argmin(energies) # 3
print("最稳定构型索引:", i_min, "能量:", energies[i_min])
9.3 np.where / np.clip
x = np.array([-3, -1, 0, 2, 5])
np.where(x > 0, x, 0) # 逐元素三元运算 → [0 0 0 2 5]
np.where(x > 0, "正", "非正") # 也支持字符串
np.clip(x, -1, 1) # 截断到 [-1, 1] → [-1 -1 0 1 1]
科研用法:
- 把能量 / 距离截断到合理范围(
np.clip防止exp溢出)。 - 给低于阈值的键长打标签(
np.where(dist < cutoff, 1, 0)判断成键)。
9.4 累积与差分
a = np.array([1, 2, 3, 4])
np.cumsum(a) # [1 3 6 10] 累积和 → 计算累积反应进度
np.diff(a) # [1 1 1] 一阶差分 → 近似梯度
np.gradient(a) # 中心差分,边界用单边
10. 线性代数 np.linalg
NumPy 的 np.linalg 底层调用 BLAS/LAPACK,速度快且数值稳定。
10.1 矩阵乘法
A = np.array([[1, 2], [3, 4]])
B = np.array([[5, 6], [7, 8]])
A @ B # 矩阵乘法,推荐
np.matmul(A, B) # 等价
np.dot(A, B) # 二维时也等价(一维时是点积)
# 批量矩阵乘法:把旋转矩阵作用到每个原子坐标
R = np.array([[0, -1, 0], [1, 0, 0], [0, 0, 1]]) # 绕 z 轴旋转 90°
pos = np.random.default_rng(4).random((96, 3))
rotated = pos @ R.T # (96,3)@(3,3) → (96,3)
区分
A * B(逐元素 Hadamard 积)与A @ B(矩阵乘法)。科研里算原子坐标 @ 变换矩阵用@。
10.2 逆、行列式、特征值
A = np.array([[4, 2], [2, 3]])
np.linalg.inv(A) # 逆矩阵(若奇异会报 LinAlgError)
np.linalg.det(A) # 行列式 → 晶格体积 = |det(cell)|
np.linalg.eig(A) # 返回 (特征值, 特征向量)
# 对称/厄米矩阵用 eigh,更快更稳(力常数矩阵就是对称的)
w, v = np.linalg.eigh(A) # w 是特征值,v 是列特征向量
DFT 关键应用:力常数矩阵(Hessian)对角化得到振动频率,这正是 np.linalg.eigh 的教科书场景——Hessian 是实对称矩阵,eigh 保证特征值全实数、特征向量正交。
10.3 范数
v = np.array([3.0, 4.0])
np.linalg.norm(v) # 默认 2-范数 → 5.0(原子间距!)
np.linalg.norm(v, ord=1) # 1-范数
np.linalg.norm(v, ord=np.inf) # 无穷范数
# 批量算距离:N 个原子到原点距离
pos = np.random.default_rng(5).random((96, 3))
dist = np.linalg.norm(pos, axis=1) # (96,)
10.4 求解线性方程组 / 最小二乘
# 解 A x = b
A = np.array([[3, 1], [1, 2]])
b = np.array([9, 8])
x = np.linalg.solve(A, b) # [2. 3.]
# 超定方程最小二乘(拟合势函数参数常用)
# 拟合 y = c0 + c1 * r,即 [1, r] @ [c0, c1] = y
r = np.array([1.0, 2.0, 3.0, 4.0])
y = np.array([2.1, 3.9, 6.2, 7.8])
design = np.column_stack([np.ones_like(r), r]) # (4, 2)
c, *_ = np.linalg.lstsq(design, y, rcond=None)
print(c) # [截距, 斜率] 的拟合值
10.5 分解
A = np.random.default_rng(6).random((4, 3))
U, S, Vt = np.linalg.svd(A, full_matrices=False) # 奇异值分解
Q, R = np.linalg.qr(A) # QR 分解
A_inv = np.linalg.pinv(A) # 伪逆(可处理奇异矩阵)
11. 随机数与 default_rng
11.1 可复现性
rng1 = np.random.default_rng(42)
rng2 = np.random.default_rng(42)
print(np.all(rng1.random(10) == rng2.random(10))) # True,同种子同序列
科研里"复现"是第一原则:把种子写进脚本顶部,任何涉及随机(MD 初速度、训练集打乱、Monte Carlo)的结果才能复现。
11.2 常用分布速查
rng = np.random.default_rng(2026)
rng.random((100, 3)) # [0,1) 均匀
rng.uniform(low, high, size) # [low, high) 均匀
rng.normal(mu, sigma, size) # 高斯
rng.standard_normal((100, 3)) # 标准正态
rng.integers(0, 10, size=5) # 整数
rng.permutation(100) # 排列
rng.choice(a, size=n, replace=False) # 无放回抽样
rng.shuffle(arr) # 原地打乱
MD 初始速度的经典写法(麦克斯韦-玻尔兹曼采样):
n_atoms = 96
masses = np.ones(n_atoms) * 12.0 # 简化:全部 C 原子质量
T = 300.0
k_B = 8.617333262e-5 # eV/K
sigma = np.sqrt(k_B * T / masses[:, None]) # (96, 1) 每个原子的速度标准差
rng = np.random.default_rng(42)
vel = rng.normal(0, sigma) # 广播 → (96, 3)
vel -= vel.mean(axis=0) # 去除质心平移
11.3 为什么不用 np.random.seed
np.random.seed(0) 设置的是全局 RandomState,之后 np.random.rand 都依赖这个隐藏的全局状态。缺点:
- 两个模块各调一次
np.random.rand,顺序一变结果全变。 - 多进程 / 多线程下全局状态难控,易产生相关序列。
default_rng() 返回独立的 Generator 对象,显式传递、局部可控。新代码一律这样写。
12. 文件读写
12.1 文本:loadtxt / savetxt / genfromtxt
# 写:保存能量列表(DFT 扫描结果)
energies = np.array([-5.2, -5.8, -4.9])
np.savetxt("energies.txt", energies, header="total_energy(eV)", fmt="%.6f")
# 读:简单规则文本
e = np.loadtxt("energies.txt") # 自动跳过 # 开头的注释行
# 读:带标题、逗号分隔、可能含缺失值
# genfromtxt 更健壮,能处理缺失值、命名列
data = np.genfromtxt("data.csv", delimiter=",",
skip_header=1, names=True,
missing_values="nan", filling_values=0.0)
print(data["energy"]) # 按列名取
NumPy 2.x 的
loadtxt/genfromtxt都支持encoding参数,读含中文/GBK 的文本时指定encoding="gbk"可避免 Windows 下报错。
12.2 二进制:save / savez / load
二进制 .npy 更快、无损、跨平台,是中间数据的最佳格式。
forces = np.random.default_rng(0).random((96, 3))
np.save("forces.npy", forces) # 单数组
f = np.load("forces.npy") # 读回
# 多数组存一个文件(.npz,类似字典)
np.savez("calc.npz", energy=energies, forces=forces, cell=np.eye(3)*10)
np.savez_compressed("calc.npz", energy=energies, forces=forces) # 压缩版
d = np.load("calc.npz")
print(d["energy"], d["forces"]) # 按键访问
建议:把每个结构的能量/力/坐标存成一个 .npz,比反复解析 OUTCAR 快得多。
13. 与 pandas / ASE 的互转
13.1 pandas
import pandas as pd
df = pd.DataFrame({"energy": [-5.2, -5.8, -4.9],
"force_norm": [0.1, 0.3, 0.2]})
arr = df.to_numpy() # DataFrame → ndarray(所有列)
energy = df["energy"].to_numpy() # 单列 → 1D 数组
# 反向:ndarray → DataFrame
df2 = pd.DataFrame(arr, columns=["energy", "force_norm"])
13.2 ASE
ASE 的 Atoms 对象里的核心数据本来就是 NumPy 数组。
from ase import Atoms
atoms = Atoms("H2O", positions=[[0, 0, 0], [0.96, 0, 0], [0, 0.96, 0]])
pos = atoms.positions # (3, 3) ndarray,注意:这是视图,勿原地乱改
cell = atoms.cell[:] # (3, 3) 晶格矩阵
numbers = atoms.numbers # (3,) 原子序数
# 反向:用数组构造新 Atoms
new_atoms = Atoms(numbers=numbers, positions=pos, cell=cell)
# 手动改坐标要赋新数组(ASE 推荐复制后再改)
pos_new = pos.copy()
pos_new += 0.1
atoms.positions = pos_new
计算化学流程基本是:ASE/pymatgen 读结构 → NumPy 数组做数学 → 写回 ASE/pymatgen。NumPy 是这个链条的"中间语言"。
14. 科研实战示例
14.1 批量处理 DFT 能量 / 受力
import numpy as np
rng = np.random.default_rng(7)
# 模拟:100 个构型的能量(eV)与受力(N,3)
n_struct = 100
n_atoms = 64
energies = rng.normal(-5.0, 0.5, n_struct) # 100 个总能量
forces = rng.normal(0, 0.2, (n_struct, n_atoms, 3)) # 每个构型每个原子受力
# 1) 找最稳定构型
i_min = np.argmin(energies)
print("最稳定构型:", i_min, "E =", energies[i_min], "eV")
# 2) 受力是否收敛:每帧最大受力分量
f_max = np.abs(forces).max(axis=(1, 2)) # (100,) 每帧最大力
converged = f_max < 0.03 # 力收敛判据 0.03 eV/Å
print("已收敛构型数:", converged.sum(), "/", n_struct)
# 3) 能量统计
print("平均能量:", energies.mean(), "±", energies.std(), "eV")
# 4) 保留收敛且能量较低的构型
keep = converged & (energies < energies.mean())
print("筛选保留:", keep.sum(), "个")
14.2 晶格向量矩阵运算
import numpy as np
# 晶格矩阵:三行分别是三条基矢 a, b, c(单位 Å)
cell = np.array([
[10.0, 0.0, 0.0],
[0.0, 11.0, 0.0],
[0.0, 0.0, 12.0],
])
# 1) 晶胞体积
vol = np.abs(np.linalg.det(cell))
print("晶胞体积:", vol, "Å^3")
# 2) 分数坐标 → 笛卡尔坐标
frac = np.array([[0.5, 0.5, 0.5],
[0.0, 0.0, 0.0]]) # (2, 3) 分数坐标
cart = frac @ cell # (2,3) @ (3,3) = 笛卡尔坐标
print("笛卡尔坐标:\n", cart)
# 3) 笛卡尔坐标 → 分数坐标
frac_back = cart @ np.linalg.inv(cell)
# 4) 应变:给晶格施加 +1% 拉伸
strain = np.eye(3) * 1.01
cell_strained = cell @ strain
print("应变后体积:", np.abs(np.linalg.det(cell_strained)))
14.3 径向分布函数(RDF)向量化实现思路
RDF 的核心是"计算所有原子对距离,再分箱统计"。向量化用广播一次性算出所有距离矩阵,避免双重循环。
import numpy as np
def rdf(pos, box, rmax, nbins):
"""立方盒子内径向分布函数(仅示意,未归一化到理想气体密度)。
pos : (N, 3) 原子笛卡尔坐标
box : 立方盒子边长
rmax : 截断半径
nbins: 分箱数
"""
N = pos.shape[0]
# 1) 所有原子对位移矢量 (N, N, 3),广播一次算完
diff = pos[:, None, :] - pos[None, :, :]
# 2) 最小镜像约定(立方盒)
diff -= box * np.round(diff / box)
# 3) 所有距离 (N, N)
dist = np.linalg.norm(diff, axis=2)
# 4) 取上三角(排除自身 + 不重复计数)
iu = np.triu_indices(N, k=1)
d = dist[iu]
# 5) 分箱直方图
hist, edges = np.histogram(d, bins=nbins, range=(0, rmax))
return edges[:-1], hist
rng = np.random.default_rng(8)
pos = rng.random((200, 3)) * 20.0 # 200 个原子,20 Å 盒子
r, h = rdf(pos, box=20.0, rmax=10.0, nbins=100)
print("RDF 直方图长度:", len(h))
真实 RDF 还要除以
4πr²ρΔr归一化,这里只演示向量化思想:用[:, None, :]广播替代双重 Python 循环,速度提升可达几十到上百倍。
14.4 势能面网格数据(PES)
import numpy as np
import matplotlib.pyplot as plt
# 两个反应坐标:如 CO2 还原中的 C-O 键长 r1、r2
r1 = np.linspace(1.0, 2.0, 200)
r2 = np.linspace(1.0, 2.0, 200)
# meshgrid 生成 (200, 200) 的网格
R1, R2 = np.meshgrid(r1, r2)
# 势能函数(示例:双势阱,示意两个稳定吸附构型)
def potential(r1, r2):
return (r1 - 1.3) ** 2 * (r1 - 1.7) ** 2 + \
(r2 - 1.4) ** 2 * (r2 - 1.6) ** 2 + \
0.1 * (r1 - r2) ** 2
E = potential(R1, R2) # (200, 200) 势能面,向量化一次算完
# 找势能面全局最低点
i_min = np.unravel_index(np.argmin(E), E.shape)
print("全局最低点坐标:", R1[i_min], R2[i_min], "能量:", E[i_min])
# 画等高线
plt.contourf(R1, R2, E, levels=50, cmap="viridis")
plt.colorbar(label="Energy (eV)")
plt.xlabel("r1 (Å)")
plt.ylabel("r2 (Å)")
plt.title("势能面")
plt.show()
关键点:meshgrid + 向量化函数,把"双重循环算能量"变成一次数组运算。
15. 常见坑与性能建议
15.1 避免 Python 循环
# 慢:Python 循环逐元素累加
total = 0
for x in arr:
total += x
# 快:向量化
total = arr.sum()
经验法则:凡是能写成数组运算的,绝不写 for。实在需要循环,把循环移到最外层、把内层交给 NumPy。
15.2 dtype 精度
# 坑:整数数组除法得到浮点,但整数赋值会截断
a = np.array([1, 2, 3])
a[0] = 2.9 # 存入整数 → 2,静默截断!
print(a) # [2 2 3]
# 正确:确保 dtype 是浮点
a = np.array([1.0, 2.0, 3.0])
a[0] = 2.9
print(a) # [2.9 2. 3.]
- 单精度 float32:占用内存减半、GPU 友好,但累加误差明显;高精度能量 / 积分用 float64。
- 大数 + 小数相加会丢精度,注意数值稳定性(如
np.log(x + 1e-12)避免log(0))。
15.3 NaN / inf 处理
DFT 计算失败、除零、溢出都会产生 NaN / inf,必须先清洗再做统计。
data = np.array([1.0, 2.0, np.nan, 4.0, np.inf])
np.isnan(data) # [False False True False False]
np.isinf(data) # [False False False False True]
np.isfinite(data) # 既不是 nan 也不是 inf
clean = data[np.isfinite(data)] # 剔除非法值
np.nanmean(data) # 忽略 NaN 的平均
np.nansum(data) # 忽略 NaN 的求和
np.nan_to_num(data) # NaN→0, inf→很大有限值
重要:
data.mean()只要含一个NaN就返回NaN。统计前务必先np.isfinite过滤或改用np.nanmean。
15.4 视图陷阱
a = np.arange(6)
b = a[1:4] # 视图
b *= 10 # 会改 a!
# 想独立:显式复制
b = a[1:4].copy()
15.5 内存与就地运算
大规模数组(如几百万帧轨迹)要注意内存,尽量用就地运算减少临时数组:
a = np.ones(1_000_000)
a += 1 # 就地,不产生新数组(快、省内存)
a = a + 1 # 产生新数组(慢)
a *= 2.0
np.add(a, 1, out=a) # 显式就地
15.6 其他 2.x 注意点
- 用
np.bool_而非np.bool;用np.str_而非np.unicode_。 np.alltrue/np.sometrue已移除 →np.all/np.any。- 梯形积分
np.trapz→np.trapezoid。 - 复制语义:
np.array(x, copy=False)若必须复制会抛ValueError,明确要"能省则省"用np.asarray(x)。
16. 延伸资源
- 官方文档:https://numpy.org/doc/stable/ — API 权威参考,遇到不确定的函数先查这里。
- NumPy 2.0 迁移指南:https://numpy.org/doc/stable/numpy_2_0_migration_guide.html — 列出了全部破坏性变更。
- From Python to Numpy:https://pnavaro.github.io/python-fortran/ 与 Nicolas P. Rougier 的《From Python to Numpy》——系统讲数组思维与向量化。
- Scipy Lecture Notes:https://scipy-lectures.org/ — NumPy/SciPy/Matplotlib 权威入门。
- 进阶工具:
scipy(优化、积分、插值、FFT)、pymatgen(材料)、ASE(原子模拟)、numba(把 Python 循环 JIT 编译到接近 C 的速度)。
学习建议:不要死记 API,而是把科研里每个"循环算数"的场景都改写成"广播 + 向量化",练 20 个例子,NumPy 思维就内化了。
完。本文所有示例基于 NumPy 2.x(已在 2.4.x 验证的公开稳定 API),如遇版本差异请以官方文档为准。
评论交流
欢迎留下你的想法