ProdForceSeA

产品支持情况

产品	是否支持
Ascend 950PR/Ascend 950DT	√
Atlas A3 训练系列产品/Atlas A3 推理系列产品	×
Atlas A2 训练系列产品/Atlas A2 推理系列产品	×
Atlas 200I/500 A2 推理产品	×
Atlas 推理系列产品	×
Atlas 训练系列产品	×

功能说明

• 算子功能：ProdForceSeA（Product Force for Smooth Edition Angular）是DeepMD-kit中DeepPot-SE模型的反向力计算算子。将神经网络对描述符的导数（net_deriv）与描述符对坐标的导数（in_deriv）逐元素相乘后归约累加，得到每个原子的三维力向量。

• 计算公式：

中心原子力（取负号）：

$$force[k][i][d]=-\sum _{a=0}^{ndescrpt-1}net\_deriv[k][i][a]\times in\_deriv[k][i][a][d]$$

邻居原子力（取正号，按nlist索引累加）：

$$force[k][j][d]+=\sum _{a=aa\_start}^{aa\_end-1}net\_deriv[k][i][a]\times in\_deriv[k][i][a][d]$$

其中 $ndescrpt=4\times nnei$，$nnei=n\_a\_sel+n\_r\_sel$，$j=nlist[k][i][j]$，$[aa\_start,aa\_end)=[j\times 4,j\times 4+4)$。

参数说明

参数名	输入/输出/属性	描述	数据类型	数据格式
net_deriv	输入	网络对描述符的导数，shape=(nframes, nloc*nnei*4)。	FLOAT	ND
in_deriv	输入	描述符对坐标的导数，shape=(nframes, nloc*nnei*4*3)。	FLOAT	ND
nlist	输入	邻居列表，shape=(nframes, nloc*nnei)，-1表示虚拟邻居。	INT32	ND
natoms	输入	原子数信息，shape=(ntypes+2,)，natoms[0]=nloc, natoms[1]=nall。	INT32	ND
n_a_sel	属性	角向邻居选择数，取值范围[0, +∞)，required。	INT	-
n_r_sel	属性	径向邻居选择数，取值范围[0, +∞)，required。n_a_sel + n_r_sel > 0。	INT	-
atom_force	输出	原子受力，shape=(nframes, nall, 3)。	FLOAT	ND

约束说明

• 仅支持float32数据类型。
• net_deriv与in_deriv的dtype必须一致。
• nlist 中元素值为-1表示虚拟邻居，跳过不处理。
• n_a_sel>=0且n_r_sel>=0，且n_a_sel+n_r_sel>0。
• natoms元素数>=3，natoms[0]=nloc>=0，natoms[1]=nall>=nloc。
• 邻居力累加使用原子加操作，输出为非确定性结果。

调用说明

调用方式	样例代码	说明
图模式调用	test_geir_prod_force_se_a	通过算子IR构图方式调用ProdForceSeA算子。