概要
GeoDiff は、ICLR 2022 の論文で説明されている分子コンフォメーション生成のための幾何拡散モデルの公式実装です。文書化されたワークフローは、GEOM 分子データセットからの学習と、テストセット内の分子のコンフォメーション生成を中心としています。これは、学習、サンプリング、評価用スクリプトを備えた研究モデルの実装であり、上流の分子データベースや汎用化学ツールキットではありません。
リポジトリには、元の GEOM データセットと、別途提供されている前処理済みデータセットの両方へのリンクがあります。データセットの場所と学習設定は YAML ファイルで構成します。学習すると、選択した出力ディレクトリにチェックポイント、設定ファイル、ログが生成されます。生成用として、プロジェクトは qm9_default および drugs_default チェックポイントを提供し、テストセットの全体または一部を選択できます。生成されたコンフォメーションはその後、評価スクリプトで処理されます。例では sample_all.pkl ファイルが参照されています。文書化された評価には、GEOM のコンフォメーションカバレッジとマッチングスコア、および別個の小規模な QM9 分割を使った物性予測ワークフローが含まれます。
ブランチと依存関係の選択には重要な制約があります。README は、提供された事前学習済みチェックポイントを使うユーザーに pretrain ブランチを指定しています。これは、main の変更によってそれらのチェックポイントとの互換性に問題が生じたためです。また、コードベースと処理済み PyD データオブジェクトの両方が、より新しい torch-geometric のバージョンと互換性がないと報告されています。したがって、提供資料が裏付けるのは特定の研究ワークフローであり、現在の環境との互換性や、任意の分子で性能が検証済みであることを示すものではありません。
主な機能
- QM9、薬物様分子、および文書化された時間ステップ削減のアブレーション設定を含む、GEOM データセットでの設定駆動型学習。
- `qm9_default` および `drugs_default` の事前学習済みチェックポイントを提供し、使用するブランチとして `pretrain` を指定。
- テストセット全体またはインデックスで選択したサブセットのコンフォメーション生成に対応し、サンプリング設定は `test.py` で指定可能。
- `eval_covmat.py` によるコンフォメーション評価。GEOM データセット上の `COV` および `MAT` スコアを報告。
- 生成したコンフォメーションと `eval_prop.py` を用いた、別個の QM9 物性予測評価ワークフロー。
- PyMol で分子コンフォメーションを表示・アラインメントするためのガイダンス。
用途
- 想定される評価:提供された GEOM テストデータの限られたサブセットについてコンフォメーションを生成し、文書化された `COV` および `MAT` のワークフローで評価する。
- 想定される評価:時間ステップを削減した薬物モデルのアブレーションを含め、設定で制御される学習条件を比較する。
- 想定される評価:生成コンフォメーションが、別個の QM9 分割におけるリポジトリの物性予測ベンチマークをどのように支えるかを調べる。
- 想定される評価:提供された PyMol の表示およびアラインメントのガイダンスを使い、生成コンフォメーションを視覚的に確認する。
使い方
- 公式 READMEを読み、ローカル環境を選ぶ前に、環境に関する説明と既知の torch-geometric 互換性の問題を確認します。
- リンク先のダウンロードフォルダーから前処理済み GEOM データを入手します。設定ファイルの
dataset設定で指定されている場所に配置し、元の GEOM データセットとは区別してください。 - 学習とチェックポイントを用いた生成のいずれかを選びます。提供された事前学習済みチェックポイントを使う場合は、pretrain ブランチを使用し、文書化されたチェックポイント/設定のディレクトリ構成を維持してください。
- 学習する場合は、リポジトリの YAML 設定と
train.pyの例を確認します。先に進む前に設定パスを確認してください。README には./configs/*.ymlと./config/の両方の例が記載されています。 - 初期評価では、文書化された
test.pyのワークフローで小さなテストセット範囲を選び、生成結果をeval_covmat.pyで確認します。別個の QM9 物性分割は独立した評価として扱い、使用前にコードと照合して例を確認してください。