matplotlibでfontが見つからないとき

matplotlibでrcParams['font.family'] = 'serif'と設定しているのに

font family ['serif'] not found. Falling back to DejaVu Sans

とエラーが出てserif体が使えないときは、以下のようにキャッシュを消してmsttcorefontsを入れると直る。

rm -r ~/.cache/matplotlib
sudo apt install msttcorefonts -qq

https://github.com/matplotlib/matplotlib/issues/13139/ https://stackoverflow.com/questions/42097053/matplotlib-cannot-find-basic-fonts/49884009#49884009

matplotlib≥3.4からはこんなことをしなくてもよくなるらしい。 https://github.com/matplotlib/matplotlib/pull/17521

JAX/HaikuでCrystal Graph Convolutional Neural Network (CGCNN)を実装した

Crystal Graph Convolutional Neural Network (CGCNN)

CGCNNは結晶にGraph Neural Network (GNN)を適用した(おそらく)初めての手法。GNNがやたら流行っている有機分子と違って、結晶の場合はそもそもグラフの作り方が自明でない、結晶の周期性のせいで隣接行列を陽に持てないなどの問題があるが、この手法はそのあたりに取り組んだ点で新しかったのだと思う。手法の詳細については論文を読んでもらうとして、今回はCGCNNのpytorch版公式実装をjax/haikuで書き直したので、ハマった点とメモを残しておく。

JAX

JAXは雑に言うと自動微分JITコンパイラが使えてバックエンドにGPUを指定できるNumpyみたいなやつ。去年くらいから流行っている気がする。最近はMDとかDFTをJAXで実装する話があって面白いなと思っている1

実装

CGCNNを実装したレポジトリには以下のリンクから飛べる。学習モデルの部分にはdm-haikuを使った。

github.com

データセットの収集

公式実装がデータセットに使ったMaterials Project上のmp-idを公開してくれていたので、そこからAPIを叩いてデータセットを集めた。ただ一部の物質でmaterial_idが変更されたらしくて大分ハマった。

lan496.hatenadiary.jp

前処理

結晶からグラフをつくるのには(元実装と同じく)pymatgen.core.Structure.get_all_neighborsを使った。最初はpymatgen.analysis.graphs.StructureGraphを使ってグラフの作り方をいろいろ試せるように書こうとしていたが、やたら重いので止めた。

ちなみにStructureGraphが重いのは自分が過去に投げたPRのせいで、pymatgen.analysis.local_env.NearNeighborsで隣接している各サイトの属するユニットセルの番号(jimage)を計算する部分がO(num_sites2)になっているから。なんでこんな実装になっているかというと、元の結晶のfractional coordinatesが[0, 1)に収まっていないエッジケースに対応する方法がこれしか思いつかなかったからだ(この修正をしたときにまさかGNNの前処理にこの部分が使われ得るとは思ってもいなかった)。上手い実装が思いつく人はPRを投げてください。

github.com

Pooling layer

モデルの実装はpytorch版とほぼ同じように書けるが、graph aggregation の部分だけは注意が必要。ひとつのバッチには次数の異なるグラフたちが並んでいるのでpoolingするのは次数に応じてsplitする必要がある。しかしjax.numpy.split(ary, indices_or_sections, axis)indices_or_sectionsには(jitするなら)定数しか渡せないし、無理やりstatic_argnumsに突っ込んでもバッチごとにjit recompileされて激重になる。解決策としてはsegment_sumを使えばjitできるようになる。

hyperparameter

argparseでhyperparameterを弄るのが好きでないのでconfig用のdataclassを作ってdataclass-jsonJSONをパースすることにした。

ベンチマーク

元のpytorch実装と同じhyperparameterで論文と同じデータセットでformation energyを学習させたらMAEが116 meV/atomのモデルを作ることをできた。しかし、論文では39 meV/atomまで下がると言っているので完全には結果を再現できていない。自前で用意したデータセットで元実装を学習させたら90 meV/atomで、これでも論文の値より大きいのでhyperparameter tuningとかの問題なのか?

Future works

いろいろやれることが思いつくけど、モチベが続かない気がするのでとりあえず区切りとする。

  • jax-mdとの連携: 夢がある。現状pymatgenでグラフを作る部分が無視できない重さなので粒子の隣接リストもjaxで書く必要がありそう。
  • 頂点の次数を可変にする: 元実装と自分の実装ではグラフの各頂点の次数が同じになるようにpaddingしている。これに関してはこことかここで議論されている。

  1. 既存のコードを置き換えるとは思えないが

Materials Projectのmaterial_idとtask_id

Materials Project に登録されているデータはmaterial_id(mp-****みたいなやつ)で指定される。だが、2018年の半ばからmaterial_idが一部の結晶では変更されたようなので注意が必要。

それぞれの計算系にはmaterial_idが振られている。そして各計算系には(single-shotとか構造緩和とかの)複数の計算条件(task_id)が対応している。material_idtask_idの中で一番精度の高い構造緩和に対応するものが選ばれる(らしい)。で、2018年にデータベースがアップデートされて、それに伴って一部のmaterial_idが以前とは違うtask_idを指すようになった。

例えばCs2CrH2Cl5O(material_id=mp-542535)は以前はmp-633688がmaterial_idとして使われていた。Materials ExplorerでみるときはどちらのIDを使っても同じページにリダイレクトされるので問題ないが、古い方のIDしか知らないときにAPIを叩きたい時は以下のようにする必要がある。

m = MPRester()
material_id_old = 'mp-633688'

material_id_new = m.get_materials_id_from_task_id(material_id_old)  # -> 'mp-542535'

# get property
formation_energy_per_atom = m.get_data(material_id_new, prop='energy_per_atom')
# get structure
structure = m.get_structure_by_material_id(material_id_new)

参考

今年に読んだ記事とかだいたい2020

年末だしPocketでクリップした記事の中で面白かったものを見返した。あと今年聴いた曲も振り返った。


invenia.github.io scheme手習いを読んでいてYコンビネータで詰まったので読んだ。

arxiv.org JAXを使った分子動力学法のナウい実装。JAXを使うために粒子の隣接リストも関数型っぽく書かなきゃいけないのがしんどそうだった。

www.youtube.com www.youtube.com Lex Fridmanのポッドキャストは胡散臭い回もあるけど↑の2つは面白かった。

yksn25.hatenablog.com nomolk.hatenablog.com dtdwtd.hatenablog.com xcloche.hateblo.jp baku89.com Titans of Tech Testify in Their Trust-Me Suits - The New York Timeswww.nytimes.com


Myd

open.spotify.com Washed out のmixで知った。曲は好きだけどアートワークは下品で好みでない

電気グルーヴ

open.spotify.com 配信再開おめでとう。ハロー! ミスターモンキーマジックオーケストラが多分サンプリング元の版権の関係で聴けないのが残念。

What Kinda Music - Tom Misch

open.spotify.com Tom Misch といえばお気楽なイメージがあったが印象が変わった。Yussef Dayes という好きなドラマーが参加している。

Vulfpeck

www.youtube.com www.youtube.com

Cory Henry

www.youtube.com open.spotify.com 今年は聴く曲がファンクに寄っていった。

High Heart - Ben Wendel

open.spotify.com 前のアルバムのThe Seasonsも好きだった。散歩しながら聴くのに丁度いい。

NGLViewで結晶構造をJupyterLab上で表示する

nglviewを使うとASEのAtoms(とpymatgenのStructure)をjupyterlab上で可視化できるのだが、nglviewはドキュメントが少なくてどうやればユニットセルを表示できるのか謎だった。

pyironで同じことをしようとしているコードをさっき見つけたので、それを参考にいい感じにpymatgenのStructureオブジェクトを表示するスクリプト↓を書いてみた。 構造が変になっていないかサクッと確認できるようになったのでとりあえず満足。

Visualize pymatgen's structure with NGLView · GitHub

f:id:lan496:20201209014116p:plain

numbaでscipy.specialの関数をつかう

numbaでscipy.specialの関数を使うには、現状自分で拡張を書く必要がある。

support for scipy.special functions : feature request · Issue #3086 · numba/numba · GitHub

↓のノートブックでは球面調和関数用のAPIを作っている。 球面調和関数の返り値は複素数で、関数の返り値がcomplexの場合は実装がnumbaだけで閉じなくて難しいらしい。 ただ今の場合、陪ルジャンドル多項式(これは実)さえ求まればそこから球面調和関数も出せるので、上のissueで示されている方法に従えばよい。

pymatgenでDOSとCOHPをプロット

さくっとプロットするには便利だがmatplotlib.pyploy.Axesを使ってくれないので細かい調整をするのは難しそう。 細かい調整をしたかったらplt.gcf()からaxesを取得すればいい。

# plotter: DosPlotter
plotter.get_plot()
fig = plt.gcf()
axes = fig.get_axes()

gist.github.com