こんばんは。今日は、材料データベースを使った機械学習による物性予測をやってみたいと思います。いわゆるマテリアルズインフォマティクスと呼ばれる分野の話題です。既にQiitaなどを見るとランダムフォレストなどの手法を用いた実例があるわけですが、本稿ではElemNetと呼ばれる多層パーセプトロン型の機械学習モデルを用いたエネルギーの予測を行いました。使ったデータベースはMaterialsProjectで、pymatgenというAPIを用いました。機械学習部分はPyTorchで実装しました。コード付きなので、適宜ご参照下さい。
以下が目次です。
１．学習用データの準備
材料科学にはさまざまなデータベースが存在していますが、ここではMaterialsProjectという有名な無機材料データベースを活用します。Python用のAPIであるpymatgenというライブラリが存在しており、Pythonからこのデータベースの中の情報を簡単に抽出することができるので大変便利です。
materialsproject.org
上記のウェブページに飛ぶと、画面右上にAPIと表示されています。ここで、Pythonで操作するためのキーを獲得することができます。ただしGoogleやgithubなどのアカウントでログインする必要があります。
まずは必要なライブラリのインポートから。
from time import sleep import numpy as np import matplotlib.pyplot as plt import pickle import torch import torch.nn as nn import torch.nn.functional as F from torch.utils.data import DataLoader, Dataset from torch.utils.data.sampler import SubsetRandomSampler from pymatgen import MPRester from pymatgen.core.periodic_table import Element
実際にデータを取得するコードは以下のように大変シンプルです。ここでは元素の組み合わせをキーとした辞書形式でデータを抽出します。辞書の値としては、例えば鉄とヒ素（Fe-As）をキーとする場合にはFe, Asからなるすべての２元型化合物の情報が格納されます。取得できた化合物は全部で3089化合物でした。
api = "MaterialsProjectから取得したAPIキーをコピー" pt = Element._member_names_ mat_dic = {} with MPRester(api) as m: for i, a in enumerate(pt): for b in pt[i+1:]: mat_dic[(a, b)] = m.get_data(f"{a}-{b}") sleep(0.5) print(f"\rset the pair:{a} and {b}", end = "")
２．Datasetクラスの定義
PyTorchには、DataLoaderというミニバッチごとに塊としてデータを吐き出す便利機能があります。ここに引数としてDatasetクラスを渡す約束になっているので、これを実装します。Datasetクラスには、辞書形式で学習データに参照したり、イテレータとして参照する２通りの流儀があります。どちらを使ってもいいのですが、ここでは辞書形式のクラスを定義します。上でデータベースから情報を引っ張ってきましたが、そのままではニューラルネットワークに食べされられないので、データの加工もこのクラスの内部で行うことにします（クラスの外で加工しても全く問題ありません）。
class MPDataset(Dataset): def __init__(self, mat_dic): s = make_set(mat_dic) lst_inp, lst_out = [], [] for n, d in enumerate(mat_dic.values()): if len(d) == 0: continue sum_ = 0 arr = torch.zeros([len(s)]) energy = d[0]['formation_energy_per_atom'] lst_out.append(torch.Tensor([energy])) for k, v in d[0]["unit_cell_formula"].items(): idx = list(s).index(k) arr[idx] = v sum_ += v arr = arr / sum_ lst_inp.append(arr) self.X = lst_inp self.y = lst_out def __len__(self): return len(self.X) def __getitem__(self, idx): out_data = self.X[idx] out_label = self.y[idx] return out_data, out_label
あっさり書きましたが、そもそも材料データをどのようにニューラルネットワークに食べさせられるよう加工するのでしょう？ここでは、以下の図のように、学習データに含まれる元素数を入力ノードの個数とし、各ノードが元素を表すようにします。そして組成情報をノードの値として入力します。例えばフッ化カルシウムCaF2であれば、Ca, Fを表すノードにそれぞれ1, 2を格納します。また、足して1になるように標準化しましょう。すると、1/3, 2/3という数値がそれぞれのノードに格納されます。それ以外のノードは全てゼロなので、かなり疎な（スパースな）入力データになります。これは以下の論文で提案された、非常にシンプルですが予測性能の高いモデルとして知られています。上記のDatasetクラスではこのような入力になるようにデータを加工しているわけです。ここでは、入力ノード数は86となりました。
３．DataLoaderの加工
学習データは訓練、バリデーション、テストに３分割する必要があります。DataLoaderクラスの引数にSamplerというのがあり、そこで指定した範囲でミニバッチを生成してくれます。SubsetRandomSamplerというクラスを使えば指定した範囲のインデックスをランダムに割り振ってくれるので、この実装は重くはないです。以下のようになります。
def train_val_test_loader(dataset, batch_size=32, val_ratio=0.2, test_ratio=0.1): total_size = len(dataset) indices = list(range(total_size)) train_ratio = 1 - val_ratio - test_ratio train_size = int(train_ratio * total_size) test_size = int(test_ratio * total_size) val_size = int(val_ratio * total_size) train_sampler = SubsetRandomSampler(indices[:train_size]) val_sampler = SubsetRandomSampler(indices[-(val_size + test_size):-test_size]) test_sampler = SubsetRandomSampler(indices[-test_size:]) train_loader = DataLoader(dataset, batch_size=batch_size, sampler=train_sampler) val_loader = DataLoader(dataset, batch_size=batch_size, sampler=val_sampler) test_loader = DataLoader(dataset, batch_size=batch_size, sampler=test_sampler) return train_loader, val_loader, test_loader
４．モデルの作成
もうちょっとなので頑張りましょう！ここでは多層パーセプトロンモデルを定義します。入力となる特徴量ベクトルの長さと、隠れ層の長さを与えてやります。ここでは隠れ層３層で、バッチノーマライゼーションだとかドロップアウトなどを挟まない非常に素朴なモデルを構築しました。
class MLP(nn.Module): def __init__(self, fea_len=86, h1_fea_len=128): super(MLP, self).__init__() self.fc1 = nn.Linear(fea_len, h1_fea_len) self.fc2 = nn.Linear(h1_fea_len, h1_fea_len) self.fc3 = nn.Linear(h1_fea_len, 1) def forward(self, x): x = F.relu(self.fc1(x)) x = F.relu(self.fc2(x)) x = F.relu(self.fc2(x)) x = F.relu(self.fc2(x)) x = self.fc3(x) return x
５．学習ループ
これで準備が整ったので、学習を行いましょう。上で定義したデータセットやモデル、それと更新手法などを定義した後、エポックとミニバッチのループを回します。この、ミニバッチのループを回す際にDataLoaderクラスを活用するわけですね。実際に見てみましょうか。
epochs = 50 batch_size = 32 dataset = MPDataset(mat_dic) train_loader, val_loader, test_loader = train_val_test_loader( dataset=dataset, batch_size=32, val_ratio=0.2, test_ratio=0.1) model = MLP() loss_function = nn.MSELoss() optimizer = torch.optim.Adam(model.parameters(), lr=0.01) for i in range(epochs): model.train() for X, y in train_loader: y_pred = model(X) optimizer.zero_grad() loss = loss_function(y_pred, y) loss.backward() optimizer.step() model.eval() for X, y in val_loader: y_val = model(X) loss_val = loss_function(y_val, y) print(f'epoch: {i:3} loss(train): {loss.item():10.8f} loss(val): {loss_val.item():10.8f}')
６．学習結果
さて、それでは結果です。まずは学習曲線の推移から。訓練データに対しては適合できていますが、検証データの損失関数がいまいち下がっていませんね。どうやらこのデータからは汎化性能を獲得するには至らなかったように見えます。また突発的に損失関数が増大するのも気になりますね。本当はこれらの大きく外れている原因を解析したいところですが、とりあえず置いておきましょう。
おわりに
長くなりました。ここまで読んで下さりありがとうございます。マテリアルズインフォマティクスってこんな感じなんだというイメージを掴んで頂けたでしょうか。使うデータセットや機械学習モデルはもちろん問題によって変わりますが、上記のDatasetクラスやモデルクラスを書き換えれば他のパーツは割と共通して使えると思います。データ解析の各ステップを覚えることが重要だと思っています。
