fix a bug when init_wfc=nao in pw basis for nspin=4

[mohanchen]

fix a bug when init_wfc=nao in pw basis for nspin=4

[mohanchen]

一个潜伏已久的 SOC 崩溃 Bug：当 NAO 遇上 Psi 类的「双重语义」

一、崩溃现场

GaAs 自旋轨道耦合（SOC）算例， init_wfc=nao ，程序一启动就崩。日志停在 "INIT BASIS" 之后、"INIT SCF" 之前，没有报错信息，只有一句冰冷的 segfault。

换 init_wfc=atomic 或 random ？不崩。换 nspin=1 ？也不崩。 只有 SOC + NAO 这个组合会触发 。

二、抽丝剥茧

第一层：NAO 算多了 band？

NAO 初始化器算出的 nbands_start_=72 ，而用户要的 nbands=30 。72 > 30，于是走进了 diag_subspace_init ——子空间对角化，把 72 个初始轨道投影到 30 个本征态。

这条路只有 NAO 会走（atomic/random 的 nbands_start_ 等于 nbands ，直接跳过）。嫌疑初步锁定。

第二层：内存爆炸？不，是越界

加调试输出，发现程序崩在 diag_subspace_init 的 hPsi 调用。 hPsi 内部调 Nonlocal::act ，后者有个 ZGEMM：

this->npw 来自 ngk_ik ，而 ngk_ik 来自 psi_temp.get_current_nbas() 。打印出来： 928 。

但 vkb （非局域势投影子）只有 464 行（= npwx ，单个旋量分量的平面波数）。K=928 > 464，gemm 越界读取，直接 segfault。

那 928 是哪来的？ 2 × 464 = 928 ，正好是 nbasis （含旋量分量 npol=2 ）。

第三层：get_current_ngk() 的双重人格

追到 psi.cpp ，看到这个函数：

同一个函数， npol=1 返回 npw （449）， npol=2 返回 nbasis （928） 。语义完全不同。函数名叫 "current_ngk"（当前 k 点的平面波数），按名字应该是 npw ，但 SOC 分支返回了含旋量分量的 nbasis 。

diag_subspace_init 用 dmin = evc.get_current_ngk() 构造 psi_temp ， psi_temp.current_nbasis = dmin = 928 ，传给 act 的 gemm 当 K → 越界。

第四层：dmin 的双重职责

本以为改 get_current_ngk() 就行。改了，结果 SCF 里其他 gemm 全报 "illegal value"。

原来 dmin 同时干两件事：

用途 需要的值 原值 hcc/scc gemm 的 K（<psi H psi> 内积维度） psi_temp.current_nbasis → act gemm 的 K npw =449（不含 npol，vkb 只有 npw 行） ✗

一个变量，两个语义，互相矛盾 。改谁都炸。

最终修复

不动 get_current_ngk() ，只把 psi_temp 构造的第4个参数从 dmin 换成 evc.get_current_nbas() （始终返回 current_nbasis = npw ）：

dmin （928）继续给 hcc/scc gemm 用， current_nbasis （449）给 act 用，各取所需。SCF 正常迭代。

三、Psi 类和 diag 的设计混乱

这个 bug 暴露了几个深层问题：

1. get_current_ngk() 的语义分裂 一个 getter 在不同 npol 下返回不同物理量（ npw vs nbasis ），违背最小惊讶原则。调用者根本不知道自己拿到的是含旋量还是不含旋量的值。全代码库有十几处调用，每处的期望语义可能都不一样——能跑全靠「碰巧对上」。

2. nbands 的含 npol 约定不明 NAO 初始化器算 nbands_start_ 时乘了 npol （得到 72），塞进 Psi::nbands 。但 Psi::nbands 在其他初始化器里不含 npol （atomic/random 是 30）。同一个成员变量，不同初始化器塞不同语义的值，下游 to_range 又无脑乘 npol ——NAO 就多乘了一次。

3. dmin 一个变量扛两个职责 diag_subspace_init 里 dmin 既是 gemm 的 K（需要含 npol），又是 psi_temp.current_nbasis （需要不含 npol）。这种变量复用是 bug 温床——改一个用途，另一个就炸。

4. current_nbasis vs nbasis 的关系隐晦 current_nbasis 理论上是「当前 k 点的活跃基函数数」，应该 = npw （不含 npol）。但 get_current_ngk() 在 SOC 下返回 nbasis （含 npol），绕过了 current_nbasis 。说明开发者自己也没理清这两个字段的关系。

四、感想

这个 bug 能潜伏这么久，是因为触发条件苛刻： SOC（npol=2）+ NAO 初始化（nbands_start > nbands）+ PW 基 ，三个条件缺一不可。日常测试覆盖不到这个组合。

但更深层的教训是： 当一个类的成员变量语义随模式切换而变化时，每个调用点都是潜在的 bug 。 Psi 类的 nbands 、 nbasis 、 current_nbasis 在 npol=1/2 下的语义关系，值得用一张表彻底理清，而不是让每个调用者自己猜。

修改仅两行（CPU + GPU 各一行），定位却花了整整一天。这就是遗留代码的代价——不是改代码难，是搞清楚该改哪里难。

[mohanchen]

If anyone interested in this story written in Chinese, you can let LLM helps you translating to English.

[kirk0830]

I think the key problem is the psi plays dual role: the multidimensional array and the wavefunction which is actually a physical quantity. The npol should only appear "between" the wavefunction and the array, say the interface between the container (programming concepts) and the physical object, otherwise, the npol would definitely bring about confusion when developers want to implement functionalities based on it (including me).

I remember when I was trying to implement the init_wfc = nao functionality, the four times memory consumption and sometimes the two times confused me quite much, since physically for a spinor wavefunction the twice memory space is enough, as the generalized Kohn-Sham theory tells, while the four times memory usage was discovered as a technical trick, while there was no information, annotations, docstrings, documentation explaining this!

A data structure should have no knowledge about the physics, the physical quantity should not know the detail of how to access the memory, but psi is still a hybrid one and requires the developers to know both the physics and memory operations, which spoils the idea of "physical programming" as I firstly heard in 2023.

Cheers.