はじめに
前回の記事（機械学習ポテンシャル実装入門）では機械学習ポテンシャルを実装する方法について紹介しました。ここではさらに踏み込んで、結晶構造の構造緩和をするために必要なstress計算を実装します。
既存の機械学習ポテンシャルパッケージの多くは、 stress計算を実装していませんが、以下の実装のようにすれば、既存パッケージも最小限の変更でstress計算を実装できます。今回は前回と同様、CGCNNにstress計算を実装します。
コード
前回作成したCGCNNを以下のように変更します。他の部分は使い回しです。
class CGCNN(nn.Module):
def __init__(
self,
node_dim: int,
edge_dim: int,
cutoff: float = 6.0,
n_layers: int = 3,
graph_reduce: str = "mean",
properties: tuple = ("energy", "forces")
):
super().__init__()
input_dim = node_dim * 2 + edge_dim
self.lin_out = nn.Linear(node_dim, 1)
self.edge_featurizer = GaussianSmearing(0.0, cutoff, edge_dim)
self.node_embedding = nn.Embedding(118, node_dim)
self.graph_reduce = graph_reduce
self._properties = self._check_properties(properties)
layers = []
for i in range(n_layers):
layers.append(CGCNNLayer(node_dim, edge_dim, cutoff))
self.layers = nn.ModuleList(layers)
def _check_properties(self, props):
available_types = (str, list, tuple)
assert isinstance(props, available_types), (
f"`props` type must be in {available_types}"
)
if isinstance(props, str):
props = (props, )
else:
props = tuple(props)
for p in props:
assert p in ["energy", "forces", "stress"]
return props
@property
def properties(self):
return self._properties
@properties.setter
def properties(self, props):
self._properties = self._check_properties(props)
def get_energy(self, batch: Batch):
node_features = self.node_embedding(batch.atomic_numbers - 1)
batch.x = node_features
distances = get_distances(batch)
edge_attr = self.edge_featurizer(distances)
batch.edge_attr = edge_attr
for layer in self.layers:
x_updated = layer(batch)
batch.x = x_updated
graph_features = scatter(batch.x, batch.batch, dim=0, reduce=self.graph_reduce)
energy = self.lin_out(graph_features)
return energy
def forward(self, batch: Batch):
self.assign_grad_mode(batch)
energy = self.get_energy(batch).view(-1)
out = {"energy": energy}
forces, stress = self.get_gradients(energy, batch)
out.update({"forces": forces, "stress": stress})
return out
def assign_grad_mode(self, batch: Batch) -> None:
"""
Change `requires_grad` parameters for force and stress calculations.
Args:
batch: :class:`torch_geometric.data.Batch` object.
"""
with torch.set_grad_enabled(True):
if "forces" in self.properties:
batch.pos.requires_grad_()
if "stress" in self.properties:
pos = torch.zeros_like(batch.pos)
cell_deform = torch.zeros_like(batch.cell, requires_grad=True)
for i in range(batch.cell.size(0)):
idx = torch.where(batch.batch==i)[0]
pos[idx] = (
batch.pos[idx]
+ torch.matmul(batch.pos[idx], cell_deform[i])
)
batch.pos = pos
batch.cell = batch.cell + torch.bmm(batch.cell, cell_deform)
batch.cell_deform = cell_deform
def get_gradients(self, energies: torch.Tensor, batch: Batch) -> tuple:
"""
Get energy gradients w.r.t. positions or cell deformation.
Args:
energies: Energy of structures.
batch: Batch object.
Returns:
tuple: Forces and stress tensors.
"""
forces, stress = None, None
wrt = []
if "forces" in self.properties:
wrt.append(batch["pos"])
if "stress" in self.properties:
wrt.append(batch["cell_deform"])
if wrt:
cell_grad = None
grads = torch.autograd.grad(
energies,
wrt,
grad_outputs=torch.ones_like(energies),
create_graph=self.training
)
if len(wrt) == 1:
if "forces" in self.properties:
forces = -1 * grads[0]
else:
cell_grad = grads[0]
else:
forces, cell_grad = grads
forces.mul_(-1)
if cell_grad is not None:
volume = torch.abs(batch["cell"].det()).detach().view(-1,1,1)
stress = cell_grad / volume
return forces, stress
propertiesという属性を追加しています。これはforceやstressの計算が必要な時と、そうでない時で計算を分岐するために使用します。energyだけで良い場合は微分計算を実行せず、計算時間を短縮できます。
また、stress計算をするために、assign_grad_mode, get_gradientsというメソッドを追加しました。
assign_grad_modeでは計算グラフの構築が必要なテンソルに対しrequires_gradをTrueにしています。具体的には、force計算時は原子位置、stress計算時はセルの変位です。
get_gradientsではself.propertiesの値に応じてエネルギーの勾配を計算します。forceとstressを同時に計算する場合は、引数にbatch["forces"]とbatch["cell_deform"]のリストを渡すことが大事です。別々に計算してしまうと、2回分のbackwardが発生し、計算時間が長くなってしまいます。リストを渡すことでほぼ1回分（＋α）の時間で計算できます。
なお、stressの計算はエネルギーを原子の変位ベクトルで微分をとることでも計算可能です。私が調べたかぎり、既存の公開されている機械学習ポテンシャルパッケージはforward内で原子の変位ベクトルを計算しているため、変位ベクトルで微分をとる実装だと、stressを計算するためにコード変更が多く必要ですので、セル変位で微分するのが良いと思います。計算時間はほぼ同じです。
計算は以下のように行います。
dataset = ToyDataset(atoms, 10)
loader = DataLoader(dataset, batch_size=4)
batch = next(iter(loader))
cgcnn = CGCNN(node_dim=100, edge_dim=50, properties=("energy", "forces", "stress"))
out = cgcnn(batch) # 4 structure batch
out.keys()
# dict_keys(['energy', 'forces', 'stress'])
out["stress"].size()
# torch.Size([4, 3, 3])
まとめ
今回は機械学習ポテンシャルモデルにstress計算を実装する方法を紹介しました。構造最適化を行うためには、torchの最適化アルゴリズムでforce, stressを最小化しても良いですし、aseのCalculatorとして実装しても良いです。
時間があればその実装方法についても紹介したいと思います。