ニューラルネットワークポテンシャルの実装　DeepPot-SE (Atom Type Embedding) ①

# ニューラルネットワークとは
　物質中の原子は相互に影響を及ぼし合っています、この相互作用をすべて足し合わせて求まる系のエネルギーをポテンシャルエネルギーと言います。ポテンシャルエネルギーと原子の座標の関係がわかっていれば、ポテンシャルエネルギーを原子の座標で微分することで、それぞれの原子にかかる力を求めることができ、構造最適化計算や分子動力学計算をすることができます。原子間の相互作用を高精度に知るには、第一原理計算という、量子力学に基づいた計算をする必要がありますが、この第一原理計算は計算コストが大きいです。そこで、原子の座標の関数であるポテンシャルエネルギーをニューラルネットワークによって近似することで、計算コストの大きい第一原理計算をバイパスし、分子動力学計算などができます。これがニューラルネットワークポテンシャルです。
　本記事では、ニューラルネットワークポテンシャルの一つである、DeepPot-SEを実際に実装します。


# 今回実装するニューラルネットワークポテンシャル
本記事では、"DeePMD-kit: A deep learning package for many-body potential energy representation and molecular dynamics"(Computer Physics Communications
Volume 228, July 2018, Pages 178-184)を実装します。アーキテクチャなどは詳しく説明しないので、論文を読んで下さい。論文と違う実装をしたところが２箇所あります。１つ目は、論文ではAtomTypeEmbeddingを使用していませんが、本記事では使用したところです。２つ目は、スイッチ関数です。これらは後で触れます。

## 前提
$E$ : 全エネルギ
$E_i$ : 原子$i$の持つエネルギー
$N$ : 全原子数
$N_i$ : 原子$i$のカットオフ半径内の原子数
$i$  : $i$番目の原子、$i = 0, 1, 2, ...., N-1$
$j$  : $i$番目の原子の隣接原子$j$、$j = 0, 1, 2, ...., N_i-1$
$R$ : 原子の座標の行列、shapeは[$N$, 3]
$R_i$ : 原子$i$から見た隣接原子の行列、shapeは[$N_i$, 3]
$\hat{R_i}$ : 変換後の原子$i$から見た隣接原子の行列、shapeは[$N_i$, 4]
$m_1$, $m_2$ : 定数、ただし$m_1 > m_2$
$g_{i1}$ : $i$原子の周りの環境用のフィルター行列、shapeは[$N_i$, $m_1$]
$g_{i2}$ : $i$原子の周りの環境用のフィルター行列、shapeは[$N_i$, $m_2$]
$D_{i}$ : 原子$i$の周りの環境を表す特徴量行列、shapeは[$m_1$, $m_2$]
$T_i$ : 原子$i$のタイプ

$\text{AtomTypeEmbedNet()}$ : 原子のタイプ埋め込み用ニューラルネットワーク
$\text{EmbedNet()}$ : 埋め込み用ニューラルネットワーク
$\text{FittingNet()}$ : フィッティング用ニューラルネットワーク


## Atom type embedding について
Embeddingというのは、人間が識別できるものから、そのモノを的確に表すような特徴量ベクトルに変換することです。今回は、Atom type embeddingなので、原子のタイプ（原子の種類）から、その原子のタイプを表す特徴量ベクトルに変換します。

まず、原子$i$のタイプ$T_i$をワンホットベクトルにします。ワンホットベクトルとは、一つだけ１になっていて、ほかの部分はすべて0になっているようなベクトルです。

例として、C H O N Si の５元素ある場合を考えます。原子のタイプはそれぞれCは1, Hは2, Oは3, Nは4, Siは5です。タイプが0の原子はないものとします。
この場合、
C原子のワンホットベクトルは、[0, 1, 0, 0, 0, 0]
H原子のワンホットベクトルは、[0, 0, 1, 0, 0, 0]
O原子のワンホットベクトルは、[0 ,0 , 0, 1, 0, 0]
N原子のワンホットベクトルは、[0, 0, 0, 0, 1, 0]
Si原子のワンホットベクトルは、[0, 0, 0, 0, 0, 1]
のようになります。
次に、これらの原子のタイプのワンホットベクトルを、AtomTypeEmbedNetというニューラルネットワークによって、特定の長さの特徴量ベクトルに変換します。
例えば、O原子の特徴量ベクトルは
$$\text{AtomTypeEmbedNet}([0, 0, 1, 0, 0])=[0.235, 0.131, 0.263, 4.242, ..., 1.633]$$
などのような、O原子を的確に表すような特徴量ベクトルに変換されます。

## Atom type embeddingを用いる理由
もともとの論文では、原子$i$の持つポテンシャルエネルギー$E_i$を予測するのに、(原子のタイプの種類数)^2個のEmbeddingNetと(原子のタイプの種類数)個のFittingNetを用意しています。これは、原子$i$のタイプと原子$j$のタイプのペア用のEmbeddingNetに原子$i$と原子$j$の位置関係を入れることによって、フィルター行列$g_i$を作成します。また、原子$i$の周りの特徴量行列$D_i$を原子$i$のタイプ用のFittingNetに入れることで原子$i$のポテンシャルエネルギー$E_i$を予測します。この方法では、高速に動作させるプログラムを作ることが難しいです。

Atom type embeddingを用いると、1個のAtomTypeEmbeddingNetと1個のEmbeddingNetと1個のFittingNetですみます。これは、原子$i$のタイプの特徴量ベクトルと原子$j$のタイプの特徴量ベクトルと原子$i$と原子$j$の位置関係をEmbeddingNetに入れることで、フィルター行列$g_i$を作成します。また、原子$i$の周りの特徴量行列$D_i$と原子$i$のタイプの特徴量ベクトルをFittingNetに入れることで原子$i$のポテンシャルエネルギー$E_i$を予測します。このように、Atom type embeddingを用いることで、より効率的に計算できます。


## 実装
先に実装したコードを貼ります。

```python
from pathlib import Path
import torch
import torch.nn as nn
import torch.nn.functional as F
from ase.neighborlist import neighbor_list as make_neighbor_list
import numpy as np
from ase import Atoms
from tqdm import trange, tqdm
from pprint import pprint
torch.set_default_tensor_type('torch.cuda.FloatTensor')

def smooth_cut_s_function(
    r_ij_norm:torch.tensor, 
    r_cutoff:float, 
)->torch.tensor:
    """DeepPot-SE function s(r_ij) of tensor version
    As r_ij becomes larger than r_cutoff_smth, the output of the function begins to decrease, 
    and as more r_ij becomes larger, the output of the function smoothly becomes zero.
    Parameters
    ----------
        r_ij_norm : torch.tensor  
            shape : (Coordination Num of atom i,)
        r_cutoff : float
            cutoff
    Returns
    -------
        s_vec : torch.tensor
            shape : (Coordination Num of atom i,)
    """
    r_cutoff_smth = 0.1

    s_vec = torch.zeros_like(r_ij_norm)
    
    flag = (r_ij_norm <= r_cutoff)

    u = (r_ij_norm[flag] - r_cutoff_smth) / (r_cutoff - r_cutoff_smth)
    s_vec[flag] =  (u*u*u*(-6*u*u + 15*u - 10) + 1) / (r_ij_norm[flag])

    return s_vec




class ThreeLayerPerceptron(nn.Module):
    def __init__(self, input_size, hidden_size, output_size, activation_function):
        super(ThreeLayerPerceptron, self).__init__()
        self.fc1 = nn.Linear(input_size, hidden_size)
        self.fc2 = nn.Linear(hidden_size, hidden_size)
        self.fc3 = nn.Linear(hidden_size, output_size)
        self.activation_function = activation_function
    
    def forward(self, x):
        x = self.fc1(x)
        x = self.activation_function(x)
        x = self.fc2(x)
        x = self.activation_function(x)
        x = self.fc3(x)
        return x


class DeepPotSEModel(nn.Module):
    """atom type embedding version of DeepPotSE
    """
    def __init__(self,
                max_atom_type:int, 
                atom_type_embed_nchanl:int,
                m1:int, 
                m2:int, 
                fitting_hidden_size:int,
                r_cutoff:float, 
                activation_function, 
                sel:int,
                ):
        """
        Parameters
        ----------
        max_atom_type : int
            Maximum type of atoms
            Example:
                For atoms of C H O N Si, 
                each atom has atom type 1, 2, 3, 4, 5, 
                so max_atom_type is 5.    
        atom_type_embed_nchanl : int
            Number of features after conversion 
            when converting the atom type to a 
            feature vector representing the atom type.
        m1 : int
            Transform S(r) into a vector with m1 features.
            See the DeepPotSE paper.
        m2 : int
            Transform S(r) into a vector with m1 features.
            Take the first half of the m2 features
            See the DeepPotSE paper.
        fitting_hidden_size : int
            Number of nodes in the hidden layer of the last fitting net
        r_cutoff : float
            Cutoff radius
        activation_function : function
            activation function
        
        sel : int
            The maximum number of atoms within the cutoff radius 
            of an atom when focusing on a particular atom.
            Set to a sufficiently large value. 
            If it is set too large, it may slow down the calculation.
        """
        super(DeepPotSEModel, self).__init__()

        self.r_cutoff = r_cutoff
        self.max_atom_type = max_atom_type
        self.m1 = m1
        self.m2 = m2
        self.activation_function = activation_function
        self.sel = sel
        self.atom_type_embed_nchanl = atom_type_embed_nchanl

        self.atom_type_embed_net = ThreeLayerPerceptron(
            input_size=max_atom_type+1, # 1-indexed
            hidden_size=atom_type_embed_nchanl, 
            output_size=atom_type_embed_nchanl, 
            activation_function=activation_function
        )

        self.embed_net = ThreeLayerPerceptron(
            input_size=1 + 2*atom_type_embed_nchanl,
            hidden_size=m1, 
            output_size=m1, 
            activation_function=activation_function
        )
        
        self.fitting_net = ThreeLayerPerceptron(
            input_size=m1*m2 + atom_type_embed_nchanl,
            hidden_size=fitting_hidden_size, 
            output_size=1, 
            activation_function=activation_function
        )
    
    def forward(
        self,
        coords:torch.tensor,
        atom_types:torch.tensor,
        atom_i_types:torch.tensor,
        atom_j_types:torch.tensor,
        atom_i_idxs:torch.tensor,
        atom_j_idxs:torch.tensor,
        shift:torch.tensor,
    ):
        """predict potential energy.
        Parameters
            coords : torch.tensor, shape : [Natoms, 3]
                The coordinates of each atom
            atom_types : torch.tensor, shape : [Natoms]
                Types of each atom
            atom_i_types : torch.tensor, shape : [Natoms*sel]
                If 0 <= atom_i_idx < Natoms, 
                then all atom_i_types[atom_i_idx:(atom_i_idx + 1)*sel] are 
                filled with atom_types[atom_i_idx].
            atom_j_types : torch.tensor, shape : [Natoms*sel]
                If 0 <= atom_i_idx < Natoms and 0 <= j < sel,
                then the value of atom_j_types[atom_i_idx*sel + j] is 
                the type of the j-th neighbor atom of atom i
            atom_i_idxs : torch.tensor, shape : [Natoms*sel]
                If 0 <= atom_i_idx < Natoms, 
                then all atom_i_idxs[atom_i_idx:(atom_i_idx + 1)*sel] are 
                filled with atom_i_idx.
            atom_j_idxs : torch.tensor, shape : [Natoms*sel]
                If 0 <= atom_i_idx < Natoms and 0 <= j < sel,
                then the value of atom_j_idxs[atom_i_idx*sel + j] is 
                the atom_idx of the j-th neighbor atom of atom i
            shift : torch.tensor, shape : [Natoms*sel, 3]
                How much the atom j is shifted by the periodic boundary condition
        """
        coords.requires_grad_(True)
        relative_coords = self.coords_to_relative_coords(
            coords=coords,
            atom_i_idxs=atom_i_idxs,
            atom_j_idxs=atom_j_idxs,
            shift=shift,
        )

        generalized_coords = self.relative_coods_to_generalized_coords(
            relative_coords=relative_coords, 
            r_cutoff=self.r_cutoff
        )

        atom_type_one_hot = torch.eye(self.max_atom_type + 1)
        atom_type_embedded_matrix = self.atom_type_embed_net(atom_type_one_hot)
        total_potential_energy = torch.tensor(0.0)

        s_rij = generalized_coords[:,0].reshape(-1, 1)
        atom_i_embedded_matrix = atom_type_embedded_matrix[atom_i_types]
        atom_j_embedded_matrix = atom_type_embedded_matrix[atom_j_types]

        g_i1s_before_embed = torch.concat((
            s_rij,
            atom_i_embedded_matrix,
            atom_j_embedded_matrix,
        ), dim=1)

        g_i1s = self.embed_net(g_i1s_before_embed).reshape(coords.shape[0], self.sel, self.m1)
        generalized_coords = generalized_coords.reshape(coords.shape[0], self.sel, 4)
        # g_i1^T @ generalized_coords
        left = torch.bmm(
            torch.transpose(g_i1s, 1, 2),
            generalized_coords
        )

        # generalized_coords^T @ g_i1[:, :m2]
        right = torch.bmm(
            torch.transpose(generalized_coords, 1, 2),
            g_i1s[:, :, :self.m2]
        )

        D_i_s_reshaped = torch.bmm(left, right).reshape(coords.shape[0], self.m1 * self.m2)

        feature_vectors = torch.concat((
            D_i_s_reshaped,
            atom_type_embedded_matrix[atom_types]
        ), dim=1)

        total_potential_energy = torch.sum(
            self.fitting_net(feature_vectors)
        )

        return total_potential_energy



    def coords_to_relative_coords(
        self,
        coords:torch.tensor, 
        atom_i_idxs:torch.tensor,
        atom_j_idxs:torch.tensor,
        shift:torch.tensor,
    ):
        coords_concat = torch.concat((
            coords,
            torch.full((1, 3), 1e5)
        ), dim=0)
        relative_coords = coords_concat[atom_j_idxs] - coords_concat[atom_i_idxs] + shift

        return relative_coords
        
```




## ThreeLayerPerceptronについて
ただの三層のニューラルネットワークです。

## DeepPotSEModelの__init__について
max_atom_typeはintで、原子のタイプを１インデックスで始めたときの最大の原子のタイプです。C, H, O, N, Siの５元素あった場合は、それぞれの原子のタイプは1, 2, 3, 4, 5となるので、max_atom_typeは5です。

atom_type_embed_nchanlはintで、原子のタイプのワンホットベクトルをAtom Type Embeddingによって原子の特徴量ベクトルに変換するのですが、変換後の特徴量数です。

m1, m2は論文と同じです。

fitting_hidden_sizeはintで、FittingNetの隠れ層のノードの数です。

r_cutoffはカットオフ半径、activation_functionは活性化関数、selはカットオフ半径内に入りうる原子の数の最大値です。

```python
        self.atom_type_embed_net = ThreeLayerPerceptron(
            input_size=max_atom_type+1, # 1-indexed
            hidden_size=atom_type_embed_nchanl, 
            output_size=atom_type_embed_nchanl, 
            activation_function=activation_function
        )
```
ここでは、Atom Type Embedding用のニューラルネットワークを定義しています。
input_sizeはmax_atom_type + 1です。タイプが0の原子は今回存在せず、データがないときはこれで埋めます。

```python
        self.embed_net = ThreeLayerPerceptron(
            input_size=1 + 2*atom_type_embed_nchanl,
            hidden_size=m1, 
            output_size=m1, 
            activation_function=activation_function
        )
```
ここでは、フィルター行列を作成するためのニューラルネットワークを定義しています。input_sizeは1 + 2*atom_type_embed_nchanlです。1というのは論文で言う$s(r_{ij})$です。これに追加して、原子$i$のタイプの情報と原子$j$のタイプの情報を一緒に入れたいので、$s(r_{ij})$と原子$i$のタイプのEmbedding後の特徴量ベクトルと原子$j$のタイプのEmbedding後の特徴量ベクトルを入れます。

```python
        self.fitting_net = ThreeLayerPerceptron(
            input_size=m1*m2 + atom_type_embed_nchanl,
            hidden_size=fitting_hidden_size, 
            output_size=1, 
            activation_function=activation_function
        )
```
ここでは、原子$i$の周りの環境を表す特徴量ベクトルを入力して、原子$i$のもつポテンシャルエネルギー$E_i$を予測するためのニューラルネットワークを定義しています。
input_sizeはm1 * m2 + atom_type_embed_nchanlで、m1*m2は原子$i$の周りの環境を表す特徴量ベクトルの大きさです。これに追加して、原子$i$のタイプの情報を一緒に入れたいので、$i$のタイプのEmbedding後の特徴量ベクトルを一緒に入れます。

## DeepPotSEModelのforwardについて

はじめにforwardの引数を説明します。
coordsは全原子の座標$R$を表し、shapeは[N, 3]です。

atom_typesは全原子のタイプを表し、shapeは[N]です。

atom_i_idxsは原子$i$のindexが入ったベクトルで、shapeは[N*sel]です。
jを0 <= j < selとして、atom_i_idxs[atom_i_idx + j]にはatom_i_idxが入っています。
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/405525/617956bb-fbe7-9712-401b-8dc9dca8cd08.png)


atom_j_idxsは原子$i$からみた原子$j$のindexが入ったベクトルで、shapeは[N*sel]です。
jを0 <= j < selとして、atom_j_idxs[atom_i_idx + j]にはatom_j_idxが入っています。下の図の例だと、atom_i_idxが1の原子の周りにはatom_idxが3, 2, 8, 5の原子があります。また、iの周りにsel個未満の原子しかない場合は、残りの部分はNで埋められます。
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/405525/e8b770af-7e87-05e2-0dd0-ea14ef937bc6.png)


atom_i_typesは原子$i$のタイプの情報が入っているベクトルで、shapeは[N*sel]です。
jを0 <= j < selとして、atom_i_types[atom_i_idx + j]には原子iのタイプが入っています。下の図の例だと、atom_i_idxが1の原子のタイプである3で埋まっています。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/405525/9fc40f22-e6c3-3f60-7d6d-5d7e753f4115.png)

atom_j_typesは原子$j$のタイプの情報が入っているベクトルで、shapeは[N*sel]です。
jを0 <= j < selとして、atom_j_types[atom_i_idx + j]には原子jのタイプが入っています。下の図の例だと、atom_i_idxが1の原子の周りには、原子のタイプが1,2,1,3の原子があります。原子がない場合は0で埋めます。原子のタイプは、1から始めており、原子のタイプが0の原子は存在しないものとしています。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/405525/3416d560-c271-f530-f50e-dfc04cc358bc.png)

shiftは原子jがどれくらい周期境界によってシフトしているかを表す行列で、shapeは[N*sel, 3]です。なぜshiftが必要かというと、原子iからみた原子jの周期境界を考えたときの方向ベクトルは単純にcoords[atom_j_idx] - coords[atom_i_idx]にはならないからです。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/405525/feaf7358-1628-b4a1-e51f-a736825faaff.png)

原子jをセルの長さ分ずらすようなベクトルをshiftとすると、
原子iからみた原子jの方向ベクトル$r_{ij}$は
coords[atom_j_idx] + shift - coords[atom_i_idx]
として求めることができます。
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/405525/d55cd16a-6bfe-8961-d977-7a9ca48d20e8.png)

引数の説明が終わったので、forwardの中を説明します。

```python
        coords.requires_grad_(True)
```
ここでは、最終的に、ポテンシャルエネルギーを位置(coords)で微分することで力を求めるため、requires_gradをTrueにしておきます。
```python
        relative_coords = self.coords_to_relative_coords(
            coords=coords,
            atom_i_idxs=atom_i_idxs,
            atom_j_idxs=atom_j_idxs,
            shift=shift,
        )
```
ここでは、$R$(coords)を$R_i$(relative_coords)に変換しています。

```python
    def coords_to_relative_coords(
        self,
        coords:torch.tensor, 
        atom_i_idxs:torch.tensor,
        atom_j_idxs:torch.tensor,
        shift:torch.tensor,
    ):
        coords_concat = torch.concat((
            coords,
            torch.full((1, 3), 1e5)
        ), dim=0)
        relative_coords = coords_concat[atom_j_idxs] - coords_concat[atom_i_idxs] + shift

        return relative_coords
```
coords_concatでは、coordsの一番下に十分大きな値を入れておきます。
相対座標は先程説明したように、coords_concat[atom_j_idxs] - coords_concat[atom_i_idxs] + shiftで求めることが出来ます。

```python
        generalized_coords = self.relative_coods_to_generalized_coords(
            relative_coords=relative_coords, 
            r_cutoff=self.r_cutoff
        )
```
```python
    def relative_coods_to_generalized_coords(
        self,
        relative_coords:torch.Tensor, 
        r_cutoff:float
    )->torch.Tensor:
        generalized_coords = torch.full((relative_coords.shape[0], 4), 1e2)
        r_ij_norm = torch.linalg.norm(relative_coords, dim=1)
        s_vec = smooth_cut_s_function(
            r_ij_norm=r_ij_norm,
            r_cutoff=r_cutoff
        )
        generalized_coords[:,0] = s_vec
        generalized_coords[:,1:] = relative_coords * s_vec.view(-1, 1) / r_ij_norm.view(-1, 1)
        return generalized_coords

def smooth_cut_s_function(
    r_ij_norm:torch.tensor, 
    r_cutoff:float, 
)->torch.tensor:
    """DeepPot-SE function s(r_ij) of tensor version
    As r_ij becomes larger than r_cutoff_smth, the output of the function begins to decrease, 
    and as more r_ij becomes larger, the output of the function smoothly becomes zero.
    Parameters
    ----------
        r_ij_norm : torch.tensor  
            shape : (Coordination Num of atom i,)
        r_cutoff : float
            cutoff
    Returns
    -------
        s_vec : torch.tensor
            shape : (Coordination Num of atom i,)
    """
    r_cutoff_smth = 0.1

    s_vec = torch.zeros_like(r_ij_norm)
    
    flag = (r_ij_norm <= r_cutoff)

    u = (r_ij_norm[flag] - r_cutoff_smth) / (r_cutoff - r_cutoff_smth)
    s_vec[flag] =  (u*u*u*(-6*u*u + 15*u - 10) + 1) / (r_ij_norm[flag])

    return s_vec


```
ここでは、$R_i$(relative_coords)を$\hat{R_i}$(generalized_coords)に変換しています。スイッチ関数は論文とは違うものを用いました。多項式の方が手で微分しやすいので、多項式のスイッチ関数を用いました。
"DPA-1: Pretraining of Attention-based Deep Potential Model for Molecular Simulation"という論文と同じスイッチ関数を用いています。

```python
        atom_type_one_hot = torch.eye(self.max_atom_type + 1)
        atom_type_embedded_matrix = self.atom_type_embed_net(atom_type_one_hot)
```
ここでは、原子のタイプのワンホットベクトルをバッチ処理によって、embeddingしています。atom_type_embedded_matrix[atom_type]には、atom_typeをembeddingした特徴量ベクトルがはいっています。

```python
        s_rij = generalized_coords[:,0].reshape(-1, 1)
        atom_i_embedded_matrix = atom_type_embedded_matrix[atom_i_types]
        atom_j_embedded_matrix = atom_type_embedded_matrix[atom_j_types]

        g_i1s_before_embed = torch.concat((
            s_rij,
            atom_i_embedded_matrix,
            atom_j_embedded_matrix,
        ), dim=1)
```
ここでは、$s(r_{ij})$をembeddingする前処理をしています。
$s(r_{ij})$、原子iのatom type embeddingと原子jのatom type embeddingを
g_i1s_before_embedにしています。

```python
        g_i1s = self.embed_net(g_i1s_before_embed).reshape(coords.shape[0], self.sel, self.m1)
        generalized_coords = generalized_coords.reshape(coords.shape[0], self.sel, 4)
```
ここでは、g_i1s_before_embedをバッチ処理によってembeddingし、shapeを[N, sel, m1]に変えています。
また、generalized_coordsのshapeを[N, sel, 4]に変えています。

```python
        # g_i1^T @ generalized_coords
        left = torch.bmm(
            torch.transpose(g_i1s, 1, 2),
            generalized_coords
        )

        # generalized_coords^T @ g_i1[:, :m2]
        right = torch.bmm(
            torch.transpose(generalized_coords, 1, 2),
            g_i1s[:, :, :self.m2]
        )

        D_i_s_reshaped = torch.bmm(left, right).reshape(coords.shape[0], self.m1 * self.m2)
```
ここでは、論文の
$$D_i = (g_{i1})^{T} \hat{R_i} \hat{R_i}^{T} g_{i2}$$
という部分で、N個の原子の$D_i$を同時に計算しています。

```python
        feature_vectors = torch.concat((
            D_i_s_reshaped,
            atom_type_embedded_matrix[atom_types]
        ), dim=1)

        total_potential_energy = torch.sum(
            self.fitting_net(feature_vectors)
        )
```
最後に、$D_i$をreshapeして、原子$i$のatom type embeddingを一緒にしてfitting netに入れることで、ポテンシャルエネルギーを予測しています。

```python
model = DeepPotSEModel(...)
total_potential_energy = model(...)
forces = -1 * torch.autograd.grad(
          total_potential_energy,
          coords,
      )[0]
```
とすることで、それぞれの原子にかかる力を求めることが出来ます。

# 精度
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/405525/1d1532d3-f6fa-fb36-9644-71433199974b.png)
学習させたものは、密度が１の水分子で、学習に用いていないものをプロットしました。横軸に第一原理計算の結果、縦軸にDeepPot-SEで予測した結果をプロットしました。きれいに予測できていることがわかります。


# 最後に
forward中にfor文が入ってしまうと、遅くなってしまうため、
なるべくpytorchの関数をうまく使うのが難しかった。