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 提供:

  1. ndarray:高效、同质(homogeneous)、多维的数组对象。
  2. 向量化运算:a + b 直接在 C 层对一整个数组做加法,而不是 Python 循环。
  3. 广播(broadcasting):不同形状的数组可以自动对齐运算。
  4. np.linalg:矩阵分解、求解、特征值,正是 DFT 里算力常数矩阵、晶格应变需要的。
  5. 统一的随机数 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.float64np.float64 或 "float64"
单精度浮点(GPU/大规模)np.float32np.float32
整数np.int64 / np.int32np.int64
复数np.complex128np.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 的区别

listndarray
元素类型可混合必须同质
内存存引用,分散连续一段内存
运算语义+ 是拼接+ 是逐元素相加
数学运算需要循环向量化,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。
  2. 某维相等,或其中一个是 1 → 兼容。
  3. 某维都不为 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. 延伸资源

学习建议:不要死记 API,而是把科研里每个"循环算数"的场景都改写成"广播 + 向量化",练 20 个例子,NumPy 思维就内化了。


完。本文所有示例基于 NumPy 2.x(已在 2.4.x 验证的公开稳定 API),如遇版本差异请以官方文档为准。