diff --git a/.gitattributes b/.gitattributes
index a6344aac8c09253b3b630fb776ae94478aa0275b..53d7257ace88cda1144353e6bbdd2fe8dba8b9a5 100644
--- a/.gitattributes
+++ b/.gitattributes
@@ -1,35 +1,47 @@
*.7z filter=lfs diff=lfs merge=lfs -text
*.arrow filter=lfs diff=lfs merge=lfs -text
*.bin filter=lfs diff=lfs merge=lfs -text
+*.bin.* filter=lfs diff=lfs merge=lfs -text
*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
*.ftz filter=lfs diff=lfs merge=lfs -text
*.gz filter=lfs diff=lfs merge=lfs -text
*.h5 filter=lfs diff=lfs merge=lfs -text
*.joblib filter=lfs diff=lfs merge=lfs -text
*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
*.model filter=lfs diff=lfs merge=lfs -text
*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
*.onnx filter=lfs diff=lfs merge=lfs -text
*.ot filter=lfs diff=lfs merge=lfs -text
*.parquet filter=lfs diff=lfs merge=lfs -text
*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
*.pt filter=lfs diff=lfs merge=lfs -text
*.pth filter=lfs diff=lfs merge=lfs -text
*.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
saved_model/**/* filter=lfs diff=lfs merge=lfs -text
*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
*.tflite filter=lfs diff=lfs merge=lfs -text
*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
*.xz filter=lfs diff=lfs merge=lfs -text
*.zip filter=lfs diff=lfs merge=lfs -text
+*.zstandard filter=lfs diff=lfs merge=lfs -text
+*.tfevents* filter=lfs diff=lfs merge=lfs -text
+*.db* filter=lfs diff=lfs merge=lfs -text
+*.ark* filter=lfs diff=lfs merge=lfs -text
+**/*ckpt*data* filter=lfs diff=lfs merge=lfs -text
+**/*ckpt*.meta filter=lfs diff=lfs merge=lfs -text
+**/*ckpt*.index filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.gguf* filter=lfs diff=lfs merge=lfs -text
+*.ggml filter=lfs diff=lfs merge=lfs -text
+*.llamafile* filter=lfs diff=lfs merge=lfs -text
+*.pt2 filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
\ No newline at end of file
diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000000000000000000000000000000000000..95e55dd03e05d600969fe9e2358bc60ad6678ff8
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1,6 @@
+__pycache__/
+*.py[cod]
+
+outputs/
+demo/slurm_*.out
+demo/slurm_*.err
diff --git a/README.md b/README.md
new file mode 100644
index 0000000000000000000000000000000000000000..07ca6caef474628629621b3aa3e4711bff958394
--- /dev/null
+++ b/README.md
@@ -0,0 +1,236 @@
+---
+frameworks:
+ - ""
+language:
+ - en
+license: mit
+tags:
+ - OneScience
+ - MatterGen
+ - materials-science
+ - crystal-generation
+ - diffusion-model
+ - graph-neural-network
+ - conditional-generation
+ - training
+ - fine-tuning
+tasks: []
+---
+
+
+
+ MatterGen
+
+
+
+# Model Introduction
+
+MatterGen is a generative model for inorganic materials developed by Microsoft Research. It jointly generates elemental compositions, unit cells, and periodic atomic coordinates, and supports candidate crystal generation conditioned on material properties.
+
+Paper: *MatterGen: a generative model for inorganic materials design*
+Reference implementation: https://github.com/microsoft/mattergen
+
+# Model Description
+
+MatterGen uses a diffusion model to learn distributions over atomic species, fractional coordinates, and crystal lattices. It supports unconditional crystal generation and conditional generation based on properties such as chemical system, space group, magnetic density, band gap, and bulk modulus.
+
+This repository contains the OneScience-adapted MatterGen model, diffusion and sampling code, and entry points for training, property fine-tuning, crystal generation, and data conversion. Model code is located in `model/`; data processing, general-purpose network layers, and property embeddings are provided by OneScience MatChem.
+
+# Use Cases
+
+| Use case | Description |
+| :---: | :--- |
+| Unconditional crystal generation | Generate new candidate crystal structures with the base checkpoint |
+| Property-conditioned generation | Generate structures for target properties such as magnetic density, band gap, and bulk modulus |
+| Fixed-composition structure prediction | Search for possible crystal structures for a specified composition |
+| Training from scratch | Train with MP-20 or data compatible with the MatterGen cache format |
+| Property fine-tuning | Add a property adapter to a pretrained model and fine-tune it |
+| Environment connectivity check | Verify OneScience, the checkpoint, and GPU/DCU availability by generating one sample |
+
+# Usage
+
+## 1. Using OneCode
+
+Try intelligent, one-click AI4S programming in the OneCode online environment:
+
+[Try intelligent, one-click AI4S programming](https://web-2069360198568017922-iaaj.ksai.scnet.cn:58043/home)
+
+## 2. Manual Installation and Usage
+
+**Hardware requirements**
+
+- A GPU or DCU is recommended.
+- A CPU can be used for import and configuration connectivity checks, but full generation and training will be slow.
+- A DCU requires DTK to match the PyTorch build. Single-sample generation has been verified with DTK 25.04.2 and `torch 2.5.1+das.opt1.dtk25042`.
+
+### Download the Model Package
+
+```bash
+hf download --model OneScience-Sugon/Mattergen --local-dir ./Mattergen
+cd Mattergen
+```
+
+### Install the Runtime Environment
+
+**DCU environment**
+
+```bash
+# Activate DTK and conda first
+module load compiler/dtk/25.04.2
+conda create -n onescience311 python=3.11 -y
+conda activate onescience311
+# uv installation is also supported
+pip install onescience[matchem-dcu] \
+ -i http://mirrors.onescience.ai:3141/pypi/simple/ \
+ --trusted-host mirrors.onescience.ai
+```
+
+**GPU environment**
+
+```bash
+# Activate conda first
+conda create -n onescience311 python=3.11 -y \
+ libstdcxx-ng=12 libgcc-ng=12 gcc_linux-64=12 gxx_linux-64=12
+conda activate onescience311
+# uv installation is also supported
+pip install onescience[matchem-gpu] \
+ -i http://mirrors.onescience.ai:3141/pypi/simple/ \
+ --trusted-host mirrors.onescience.ai
+```
+
+### Training Data
+
+MatterGen training uses a converted data cache. The MP-20 dataset is available on Hugging Face and can be downloaded directly:
+
+```bash
+hf download --dataset OneScience-Sugon/mp20 --local-dir ./datasets/mp20
+```
+
+The downloaded MatterGen cache is located at:
+
+```text
+datasets/mp20/data/MP20/cache/mp_20/
+├── train/
+├── val/
+└── test/
+```
+
+The default training YAML cannot use the downloaded data until you add its actual path under `args`. For example, update `demo/configs/train_8dcu.yaml` as follows:
+
+```yaml
+args:
+ data_module: mp_20
+ data_module.root_dir: ./datasets/mp20/data/MP20/cache/mp_20
+```
+
+For property fine-tuning, also add `data_module.root_dir` under `args` in `demo/configs/finetune_dft_mag_density_smoke.yaml`. If you downloaded the data elsewhere, use the corresponding absolute path.
+
+Convert custom CSV data first:
+
+```bash
+python csv_to_dataset.py \
+ --csv-folder /path/to/csv_folder \
+ --dataset-name my_dataset \
+ --cache-folder ./datasets/cache
+```
+
+### Trained Weights
+
+After downloading the model package with the Hugging Face command above, the pretrained checkpoints are located in `weight/`:
+
+```text
+weight/
+├── mattergen_base/
+│ ├── config.yaml
+│ └── checkpoints/
+│ └── last.ckpt
+└── dft_mag_density/
+ ├── config.yaml
+ └── checkpoints/
+ └── last.ckpt
+```
+
+The repository also provides property-conditioned checkpoints for `chemical_system`, `dft_band_gap`, `dft_mag_density`, `ml_bulk_modulus`, and `space_group`. When using a demo YAML, set `checkpoint` in `demo/configs/generate_base.yaml` or `demo/configs/generate_dft_mag_density.yaml` to the corresponding `./weight/` path. For fine-tuning, update `adapter.model_path` in `demo/configs/finetune_dft_mag_density_smoke.yaml`.
+
+### Inference
+
+Generate one crystal structure unconditionally:
+
+```bash
+python generate.py \
+ --checkpoint ./weight/mattergen_base \
+ --output outputs/generate/mattergen_base \
+ --batch-size 1 \
+ --num-batches 1
+```
+
+The generated outputs include:
+
+```text
+outputs/generate/mattergen_base/
+├── generated_crystals.extxyz
+└── generated_crystals_cif.zip
+```
+
+Magnetic-density-conditioned generation requires the matching property checkpoint:
+
+```bash
+python generate.py \
+ --checkpoint ./weight/dft_mag_density \
+ --output outputs/generate/dft_mag_density \
+ --batch-size 1 \
+ --num-batches 1 \
+ --properties '{"dft_mag_density": 0.15}'
+```
+
+You can also use the demo YAML files:
+
+```bash
+cd demo
+bash run.sh --config configs/generate_base.yaml
+bash run.sh --config configs/generate_dft_mag_density.yaml
+```
+
+### Training
+
+Submit an eight-DCU training job with the MP-20 cache:
+
+```bash
+cd demo
+bash run.sh --config configs/train_8dcu.yaml --submit
+```
+
+You can also launch training directly from the repository root with Hydra arguments:
+
+```bash
+python train.py \
+ data_module=mp_20 \
+ data_module.root_dir=./datasets/mp20/data/MP20/cache/mp_20 \
+ trainer.devices=1 \
+ data_module.batch_size.train=4
+```
+
+### Fine-Tuning
+
+The repository provides a smoke-test configuration for magnetic-density property fine-tuning. By default, it runs only one training batch and one validation batch:
+
+```bash
+cd demo
+bash run.sh --config configs/finetune_dft_mag_density_smoke.yaml
+```
+
+For a full fine-tuning run, copy the YAML and update `adapter.model_path`, `data_module.properties`, the property embedding, batch size, and number of epochs. Remove `trainer.limit_train_batches` and `trainer.limit_val_batches`.
+
+# Official OneScience Resources
+
+| Platform | OneScience Main Repository | Skills Repository |
+| --- | --- | --- |
+| Gitee | https://gitee.com/onescience-ai/onescience | https://gitee.com/onescience-ai/oneskills |
+| GitHub | https://github.com/onescience-ai/OneScience | https://github.com/onescience-ai/oneskills |
+
+---
+
+# Citation and License
+
+- The MatterGen-related code comes from the MatChem examples in the OneScience project and refers to the upstream MatterGen project (https://github.com/microsoft/mattergen). The upstream MatterGen code is released under the [MIT License](https://github.com/microsoft/mattergen/blob/main/LICENSE).
+- If you use MatterGen training or generation results in research, please cite the original MatterGen paper, the relevant OneScience projects, and the datasets used.
diff --git a/conf/adapter/default.yaml b/conf/adapter/default.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..c30e6d0cd4345e0ece307a44e2f06acd4268589e
--- /dev/null
+++ b/conf/adapter/default.yaml
@@ -0,0 +1,16 @@
+pretrained_name: null
+model_path: ${oc.env:ONESCIENCE_MODELS_DIR,/public/share/sugonhpcapp01/onestore/onemodels}/mattergen/mattergen_base
+load_epoch: last
+full_finetuning: true
+
+adapter:
+ # these arguments are used to initialize GemNetTAdapter
+ # more args are added by the finetuning script during runtime
+ _target_: model.adapter.GemNetTAdapter
+ property_embeddings_adapt: {}
+
+defaults: []
+ # path/to/config_dir@attribute.name: config_file_name
+ ## e.g., insert values from dft_bulk_modulus.yaml in /lightning_module/diffusion_module/model/property_embeddings/
+ ## into adapter.property_embeddings_adapt[dft_bulk_modulus]
+ # - /lightning_module/diffusion_module/model/property_embeddings@adapter.property_embeddings_adapt.dft_bulk_modulus: dft_bulk_modulus
diff --git a/conf/csp.yaml b/conf/csp.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..fa9d880df3903a59c920e35863ea52a076e754e1
--- /dev/null
+++ b/conf/csp.yaml
@@ -0,0 +1,14 @@
+hydra:
+ run:
+ dir: ${oc.env:OUTPUT_DIR,outputs/singlerun/${now:%Y-%m-%d}/${now:%H-%M-%S}}
+
+
+auto_resume: true
+
+defaults:
+ - data_module: mp_20
+ - trainer: default
+ - lightning_module: default
+ - lightning_module/diffusion_module: csp
+ - lightning_module/diffusion_module/model: mattergen
+ - lightning_module/diffusion_module/corruption: csp
diff --git a/conf/data_module/alex_mp_20.yaml b/conf/data_module/alex_mp_20.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..fbcc187ab1933b739cee043fbabe9174dd75379b
--- /dev/null
+++ b/conf/data_module/alex_mp_20.yaml
@@ -0,0 +1,49 @@
+_target_: onescience.datapipes.materials.mattergen.datamodule.CrystDataModule
+_recursive_: true
+properties: []
+ # Supported properties:
+ # - dft_bulk_modulus
+ # - dft_band_gap
+ # - dft_mag_density
+ # - ml_bulk_modulus
+ # - hhi_score
+ # - space_group
+ # - energy_above_hull
+
+dataset_transforms:
+ - _target_: onescience.datapipes.materials.mattergen.dataset_transform.filter_sparse_properties
+ _partial_: true
+
+transforms:
+- _target_: onescience.datapipes.materials.mattergen.transform.symmetrize_lattice
+ _partial_: true
+- _target_: onescience.datapipes.materials.mattergen.transform.set_chemical_system_string
+ _partial_: true
+
+average_density: 0.05771451654022283 # atoms/Angstrom**3 : this is used in models/scripts/run.py to set lattice_limit_density
+root_dir: ${oc.env:ONESCIENCE_DATASETS_DIR,/public/share/sugonhpcapp01/onestore/onedatasets}/matchem/mattergen/cache/alex_mp_20
+
+train_dataset:
+ _target_: onescience.datapipes.materials.mattergen.dataset.CrystalDataset.from_cache_path
+ cache_path: ${data_module.root_dir}/train
+ properties: ${data_module.properties}
+ transforms: ${data_module.transforms}
+ dataset_transforms: ${data_module.dataset_transforms}
+
+val_dataset:
+ _target_: onescience.datapipes.materials.mattergen.dataset.CrystalDataset.from_cache_path
+ cache_path: ${data_module.root_dir}/val
+ properties: ${data_module.properties}
+ transforms: ${data_module.transforms}
+ dataset_transforms: ${data_module.dataset_transforms}
+
+num_workers:
+ train: 0
+ val: 0
+
+batch_size:
+ # total batch size of 512, adjust for number of devices, nodes, and gradient accumulation
+ train: ${eval:'(512 // ${trainer.accumulate_grad_batches}) // (${trainer.devices} * ${trainer.num_nodes})'}
+ val: ${eval:'(512 // ${trainer.accumulate_grad_batches}) // (${trainer.devices} * ${trainer.num_nodes})'}
+
+max_epochs: 2200
diff --git a/conf/data_module/mp_20.yaml b/conf/data_module/mp_20.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..4fe5eaf16430180848cc22868719bb5d987b0e1e
--- /dev/null
+++ b/conf/data_module/mp_20.yaml
@@ -0,0 +1,54 @@
+_target_: onescience.datapipes.materials.mattergen.datamodule.CrystDataModule
+_recursive_: true
+properties: []
+ # Supported properties:
+ # - dft_bulk_modulus
+ # - dft_band_gap
+ # - dft_mag_density
+
+transforms:
+- _target_: onescience.datapipes.materials.mattergen.transform.symmetrize_lattice
+ _partial_: true
+- _target_: onescience.datapipes.materials.mattergen.transform.set_chemical_system_string
+ _partial_: true
+
+dataset_transforms:
+ - _target_: onescience.datapipes.materials.mattergen.dataset_transform.filter_sparse_properties
+ _partial_: true
+
+average_density: 0.05771451654022283 # atoms/Angstrom**3 : this is used in models/scripts/run.py to set lattice_limit_density
+root_dir: ${oc.env:ONESCIENCE_DATASETS_DIR}/matchem/mattergen/cache/mp_20
+
+train_dataset:
+ _target_: onescience.datapipes.materials.mattergen.dataset.CrystalDataset.from_cache_path
+ cache_path: ${data_module.root_dir}/train
+ properties: ${data_module.properties}
+ transforms: ${data_module.transforms}
+ dataset_transforms: ${data_module.dataset_transforms}
+
+val_dataset:
+ _target_: onescience.datapipes.materials.mattergen.dataset.CrystalDataset.from_cache_path
+ cache_path: ${data_module.root_dir}/val
+ properties: ${data_module.properties}
+ transforms: ${data_module.transforms}
+ dataset_transforms: ${data_module.dataset_transforms}
+
+test_dataset:
+ _target_: onescience.datapipes.materials.mattergen.dataset.CrystalDataset.from_cache_path
+ cache_path: ${data_module.root_dir}/test
+ properties: ${data_module.properties}
+ transforms: ${data_module.transforms}
+ dataset_transforms: ${data_module.dataset_transforms}
+
+num_workers:
+ train: 0
+ val: 0
+ test: 0
+
+batch_size:
+ # total batch size of 512, adjust for number of devices, nodes, and gradient accumulation
+ train: ${eval:'(512 // ${trainer.accumulate_grad_batches}) // (${trainer.devices} * ${trainer.num_nodes})'}
+ val: ${eval:'(64 // (${trainer.devices} * ${trainer.num_nodes}))'}
+ test: ${eval:'(64 // (${trainer.devices} * ${trainer.num_nodes}))'}
+
+max_epochs: 900
diff --git a/conf/default.yaml b/conf/default.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..bb8893f8b4b0fb343ae673642fc4951f577262cc
--- /dev/null
+++ b/conf/default.yaml
@@ -0,0 +1,13 @@
+hydra:
+ run:
+ dir: ${oc.env:OUTPUT_DIR,outputs/singlerun/${now:%Y-%m-%d}/${now:%H-%M-%S}}
+
+auto_resume: True
+
+defaults:
+ - data_module: mp_20
+ - trainer: default
+ - lightning_module: default
+ - lightning_module/diffusion_module: default
+ - lightning_module/diffusion_module/model: mattergen
+ - lightning_module/diffusion_module/corruption: default
diff --git a/conf/finetune.yaml b/conf/finetune.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..9a9c53fa537a1205cc61f711c4162d9d01897d0f
--- /dev/null
+++ b/conf/finetune.yaml
@@ -0,0 +1,18 @@
+hydra:
+ run:
+ dir: ${oc.env:OUTPUT_DIR,outputs/singlerun/${now:%Y-%m-%d}/${now:%H-%M-%S}}
+
+defaults:
+ - data_module: mp_20
+ - trainer: default
+ - lightning_module: default
+ - adapter: default
+
+trainer:
+ max_epochs: 200
+ logger:
+ job_type: train_finetune # override default defined in defaults.trainer yaml file
+
+lightning_module:
+ optimizer_partial:
+ lr: 5e-6
diff --git a/conf/lightning_module/default.yaml b/conf/lightning_module/default.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..3b302c5e4e2bc4c299f3e83cce31f1445b5b68a8
--- /dev/null
+++ b/conf/lightning_module/default.yaml
@@ -0,0 +1,17 @@
+_target_: model.diffusion.lightning_module.DiffusionLightningModule
+optimizer_partial:
+ lr: 1e-4
+ _target_: torch.optim.Adam
+ _partial_: true
+scheduler_partials:
+ - scheduler:
+ _target_: torch.optim.lr_scheduler.ReduceLROnPlateau
+ factor: 0.6
+ patience: 100
+ min_lr: 1e-6
+ verbose: true
+ _partial_: true
+ interval: epoch
+ frequency: 1
+ monitor: loss_train
+ strict: true
diff --git a/conf/lightning_module/diffusion_module/corruption/csp.yaml b/conf/lightning_module/diffusion_module/corruption/csp.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..eabae62bb8b3386ff9f4fd92f09baf1a2100f15e
--- /dev/null
+++ b/conf/lightning_module/diffusion_module/corruption/csp.yaml
@@ -0,0 +1,15 @@
+_target_: model.diffusion.corruption.multi_corruption.MultiCorruption
+sdes:
+ pos:
+ _target_: model.common.diffusion.corruption.NumAtomsVarianceAdjustedWrappedVESDE
+ wrapping_boundary: 1.0
+ sigma_max: 5.0
+ limit_info_key: num_atoms
+
+ cell:
+ _target_: model.common.diffusion.corruption.LatticeVPSDE.from_vpsde_config
+ vpsde_config:
+ beta_min: 0.1
+ beta_max: 20
+ limit_density: ${data_module.average_density}
+ limit_var_scaling_constant: 0.25
diff --git a/conf/lightning_module/diffusion_module/corruption/default.yaml b/conf/lightning_module/diffusion_module/corruption/default.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..308d4c0b05693fd90277656d0d18db23fe9bcd3f
--- /dev/null
+++ b/conf/lightning_module/diffusion_module/corruption/default.yaml
@@ -0,0 +1,27 @@
+_target_: model.diffusion.corruption.multi_corruption.MultiCorruption
+sdes:
+ pos:
+ _target_: model.common.diffusion.corruption.NumAtomsVarianceAdjustedWrappedVESDE
+ wrapping_boundary: 1.0
+ sigma_max: 5.0
+ limit_info_key: num_atoms
+
+ cell:
+ _target_: model.common.diffusion.corruption.LatticeVPSDE.from_vpsde_config
+ vpsde_config:
+ beta_min: 0.1
+ beta_max: 20
+ limit_density: ${data_module.average_density}
+ limit_var_scaling_constant: 0.25
+
+discrete_corruptions:
+ atomic_numbers:
+ _target_: model.diffusion.corruption.d3pm_corruption.D3PMCorruption
+ offset: 1
+ d3pm:
+ _target_: model.diffusion.d3pm.d3pm.MaskDiffusion
+ dim: 101
+ schedule:
+ _target_: model.diffusion.d3pm.d3pm.create_discrete_diffusion_schedule
+ kind: standard
+ num_steps: 1000
diff --git a/conf/lightning_module/diffusion_module/csp.yaml b/conf/lightning_module/diffusion_module/csp.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..0cbc03cc8eb538c85531defda30e06b992ae495f
--- /dev/null
+++ b/conf/lightning_module/diffusion_module/csp.yaml
@@ -0,0 +1,19 @@
+_target_: model.diffusion.diffusion_module.DiffusionModule
+model: mattergen
+corruption: csp
+
+loss_fn:
+ _target_: model.common.loss.MaterialsLoss
+ reduce: sum
+ include_pos: True
+ include_cell: True
+ include_atomic_numbers: False
+ weights:
+ cell: 1.0
+ pos: 0.1
+
+
+pre_corruption_fn:
+ _target_: onescience.modules.embedding.mattergen_property_embeddings.SetEmbeddingType
+ p_unconditional: 0.2
+ dropout_fields_iid: false
\ No newline at end of file
diff --git a/conf/lightning_module/diffusion_module/default.yaml b/conf/lightning_module/diffusion_module/default.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..b75f1ad872b0e2886694b5f4e15f9d89c1fdb29b
--- /dev/null
+++ b/conf/lightning_module/diffusion_module/default.yaml
@@ -0,0 +1,18 @@
+_target_: model.diffusion.diffusion_module.DiffusionModule
+loss_fn:
+ _target_: model.common.loss.MaterialsLoss
+ reduce: sum
+ include_pos: True
+ include_cell: True
+ include_atomic_numbers: True
+ d3pm_hybrid_lambda: 0.01
+ weights:
+ cell: 1.0
+ pos: 0.1
+ atomic_numbers: 1.0
+model: mattergen
+corruption: default
+pre_corruption_fn:
+ _target_: onescience.modules.embedding.mattergen_property_embeddings.SetEmbeddingType
+ p_unconditional: 0.2
+ dropout_fields_iid: false
diff --git a/conf/lightning_module/diffusion_module/model/mattergen.yaml b/conf/lightning_module/diffusion_module/model/mattergen.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..4f8c04d189b60bc76218d0a9a862e3b2b002ef07
--- /dev/null
+++ b/conf/lightning_module/diffusion_module/model/mattergen.yaml
@@ -0,0 +1,29 @@
+_target_: model.denoiser.GemNetTDenoiser
+hidden_dim: 512
+gemnet:
+ _target_: model.common.gemnet.gemnet.GemNetT
+ num_targets: 1
+ latent_dim: ${eval:'${..hidden_dim} * (1 + len(${..property_embeddings}))'} # 1 is for time encoding.
+ atom_embedding:
+ _target_: onescience.modules.layer.mattergen.embedding_block.AtomEmbedding
+ emb_size: ${...hidden_dim}
+ with_mask_type: ${eval:'${...denoise_atom_types} and "${...atom_type_diffusion}" == "mask"'}
+ emb_size_atom: ${..hidden_dim}
+ emb_size_edge: ${..hidden_dim}
+ max_neighbors: 50
+ max_cell_images_per_dim: 5
+ cutoff: 7.
+ num_blocks: 4
+ regress_stress: true
+ otf_graph: true
+ scale_file: ${oc.env:PROJECT_ROOT}/common/gemnet/gemnet-dT.json
+denoise_atom_types: true
+atom_type_diffusion: mask
+property_embeddings_adapt: {}
+property_embeddings: {}
+defaults: [] # NOTE: to train a conditional model, unccoment entries such as property_embeddings@property_embeddings.chemical_system: chemical_system below and edit/add properties to the defaults list as desired.
+ # see https://stackoverflow.com/questions/71356361/selecting-multiple-configs-from-a-config-group-in-hydra-without-using-an-explici
+ # add via config override: +lightning_module/diffusion_module/model/property_embeddings@lightning_module.diffusion_module.model.property_embeddings.dft_bulk_modulus=dft_bulk_modulus
+ # delete via config override: ~lightning_module/diffusion_module/model/property_embeddings@lightning_module.diffusion_module.model.property_embeddings.chemical_system
+ # - property_embeddings@property_embeddings.chemical_system: chemical_system
+ # - property_embeddings@property_embeddings.dft_bulk_modulus: dft_bulk_modulus
diff --git a/conf/lightning_module/diffusion_module/model/property_embeddings/chemical_system.yaml b/conf/lightning_module/diffusion_module/model/property_embeddings/chemical_system.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..5a13ab75a0bc3fd4dbb2c133d6766df1f16bd9f3
--- /dev/null
+++ b/conf/lightning_module/diffusion_module/model/property_embeddings/chemical_system.yaml
@@ -0,0 +1,10 @@
+_target_: onescience.modules.embedding.mattergen_property_embeddings.PropertyEmbedding
+name: chemical_system
+unconditional_embedding_module:
+ _target_: onescience.modules.embedding.mattergen_property_embeddings.EmbeddingVector
+ hidden_dim: ${lightning_module.diffusion_module.model.hidden_dim}
+conditional_embedding_module:
+ _target_: onescience.modules.embedding.mattergen_property_embeddings.ChemicalSystemMultiHotEmbedding
+ hidden_dim: ${lightning_module.diffusion_module.model.hidden_dim}
+scaler:
+ _target_: torch.nn.Identity
diff --git a/conf/lightning_module/diffusion_module/model/property_embeddings/dft_band_gap.yaml b/conf/lightning_module/diffusion_module/model/property_embeddings/dft_band_gap.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..561bcfca48f81584db911315b549f4398284d552
--- /dev/null
+++ b/conf/lightning_module/diffusion_module/model/property_embeddings/dft_band_gap.yaml
@@ -0,0 +1,10 @@
+_target_: onescience.modules.embedding.mattergen_property_embeddings.PropertyEmbedding
+name: dft_band_gap
+unconditional_embedding_module:
+ _target_: onescience.modules.embedding.mattergen_property_embeddings.EmbeddingVector
+ hidden_dim: ${lightning_module.diffusion_module.model.hidden_dim}
+conditional_embedding_module:
+ _target_: model.diffusion.model_utils.NoiseLevelEncoding
+ d_model: ${lightning_module.diffusion_module.model.hidden_dim}
+scaler:
+ _target_: model.common.utils.data_utils.StandardScalerTorch
diff --git a/conf/lightning_module/diffusion_module/model/property_embeddings/dft_bulk_modulus.yaml b/conf/lightning_module/diffusion_module/model/property_embeddings/dft_bulk_modulus.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..71c33118ef4c39526279fdca732be02e028be14e
--- /dev/null
+++ b/conf/lightning_module/diffusion_module/model/property_embeddings/dft_bulk_modulus.yaml
@@ -0,0 +1,11 @@
+_target_: onescience.modules.embedding.mattergen_property_embeddings.PropertyEmbedding
+name: dft_bulk_modulus
+unconditional_embedding_module:
+ _target_: onescience.modules.embedding.mattergen_property_embeddings.EmbeddingVector
+ hidden_dim: ${lightning_module.diffusion_module.model.hidden_dim}
+conditional_embedding_module:
+ _target_: model.diffusion.model_utils.NoiseLevelEncoding
+ d_model: ${lightning_module.diffusion_module.model.hidden_dim}
+scaler:
+ _target_: model.common.utils.data_utils.StandardScalerTorch
+ log10_transform: true
diff --git a/conf/lightning_module/diffusion_module/model/property_embeddings/dft_mag_density.yaml b/conf/lightning_module/diffusion_module/model/property_embeddings/dft_mag_density.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..571c78ae0130207a5f5157f98e41cad1f3a7ebfe
--- /dev/null
+++ b/conf/lightning_module/diffusion_module/model/property_embeddings/dft_mag_density.yaml
@@ -0,0 +1,10 @@
+_target_: onescience.modules.embedding.mattergen_property_embeddings.PropertyEmbedding
+name: dft_mag_density
+unconditional_embedding_module:
+ _target_: onescience.modules.embedding.mattergen_property_embeddings.EmbeddingVector
+ hidden_dim: ${lightning_module.diffusion_module.model.hidden_dim}
+conditional_embedding_module:
+ _target_: model.diffusion.model_utils.NoiseLevelEncoding
+ d_model: ${lightning_module.diffusion_module.model.hidden_dim}
+scaler:
+ _target_: model.common.utils.data_utils.StandardScalerTorch
diff --git a/conf/lightning_module/diffusion_module/model/property_embeddings/energy_above_hull.yaml b/conf/lightning_module/diffusion_module/model/property_embeddings/energy_above_hull.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..777e6dc6608a8309fde0edd201ada2f9e57aeb4c
--- /dev/null
+++ b/conf/lightning_module/diffusion_module/model/property_embeddings/energy_above_hull.yaml
@@ -0,0 +1,10 @@
+_target_: onescience.modules.embedding.mattergen_property_embeddings.PropertyEmbedding
+name: energy_above_hull
+unconditional_embedding_module:
+ _target_: onescience.modules.embedding.mattergen_property_embeddings.EmbeddingVector
+ hidden_dim: ${lightning_module.diffusion_module.model.hidden_dim}
+conditional_embedding_module:
+ _target_: model.diffusion.model_utils.NoiseLevelEncoding
+ d_model: ${lightning_module.diffusion_module.model.hidden_dim}
+scaler:
+ _target_: model.common.utils.data_utils.StandardScalerTorch
diff --git a/conf/lightning_module/diffusion_module/model/property_embeddings/hhi_score.yaml b/conf/lightning_module/diffusion_module/model/property_embeddings/hhi_score.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..430b6cf416695ec013fa83ab523a7af3c46381f5
--- /dev/null
+++ b/conf/lightning_module/diffusion_module/model/property_embeddings/hhi_score.yaml
@@ -0,0 +1,10 @@
+_target_: onescience.modules.embedding.mattergen_property_embeddings.PropertyEmbedding
+name: hhi_score
+unconditional_embedding_module:
+ _target_: onescience.modules.embedding.mattergen_property_embeddings.EmbeddingVector
+ hidden_dim: ${lightning_module.diffusion_module.model.hidden_dim}
+conditional_embedding_module:
+ _target_: model.diffusion.model_utils.NoiseLevelEncoding
+ d_model: ${lightning_module.diffusion_module.model.hidden_dim}
+scaler:
+ _target_: model.common.utils.data_utils.StandardScalerTorch
diff --git a/conf/lightning_module/diffusion_module/model/property_embeddings/ml_bulk_modulus.yaml b/conf/lightning_module/diffusion_module/model/property_embeddings/ml_bulk_modulus.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..5a49e02e8b99f068046d3840a606ac86d02c0fde
--- /dev/null
+++ b/conf/lightning_module/diffusion_module/model/property_embeddings/ml_bulk_modulus.yaml
@@ -0,0 +1,11 @@
+_target_: onescience.modules.embedding.mattergen_property_embeddings.PropertyEmbedding
+name: ml_bulk_modulus
+unconditional_embedding_module:
+ _target_: onescience.modules.embedding.mattergen_property_embeddings.EmbeddingVector
+ hidden_dim: ${lightning_module.diffusion_module.model.hidden_dim}
+conditional_embedding_module:
+ _target_: model.diffusion.model_utils.NoiseLevelEncoding
+ d_model: ${lightning_module.diffusion_module.model.hidden_dim}
+scaler:
+ _target_: model.common.utils.data_utils.StandardScalerTorch
+ log10_transform: true
\ No newline at end of file
diff --git a/conf/lightning_module/diffusion_module/model/property_embeddings/space_group.yaml b/conf/lightning_module/diffusion_module/model/property_embeddings/space_group.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..38a814ca27a7ac4b0c2ab6295c92b91dae62a764
--- /dev/null
+++ b/conf/lightning_module/diffusion_module/model/property_embeddings/space_group.yaml
@@ -0,0 +1,10 @@
+_target_: onescience.modules.embedding.mattergen_property_embeddings.PropertyEmbedding
+name: space_group
+unconditional_embedding_module:
+ _target_: onescience.modules.embedding.mattergen_property_embeddings.EmbeddingVector
+ hidden_dim: ${lightning_module.diffusion_module.model.hidden_dim}
+conditional_embedding_module:
+ _target_: onescience.modules.embedding.mattergen_property_embeddings.SpaceGroupEmbeddingVector
+ hidden_dim: ${lightning_module.diffusion_module.model.hidden_dim}
+scaler:
+ _target_: torch.nn.Identity
diff --git a/conf/trainer/default.yaml b/conf/trainer/default.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..11c351761133a7bfd9077627fe20936c645ae524
--- /dev/null
+++ b/conf/trainer/default.yaml
@@ -0,0 +1,38 @@
+_target_: pytorch_lightning.Trainer
+accelerator: 'gpu'
+devices: 1
+num_nodes: 1
+precision: 32
+max_epochs: ${data_module.max_epochs}
+accumulate_grad_batches: 1
+gradient_clip_val: 0.5
+gradient_clip_algorithm: value
+check_val_every_n_epoch: 5
+strategy:
+ _target_: pytorch_lightning.strategies.ddp.DDPStrategy
+ find_unused_parameters: true
+
+logger:
+ _target_: pytorch_lightning.loggers.WandbLogger
+ project: crystal-generation
+ job_type: train
+ settings:
+ _target_: wandb.Settings
+ start_method: fork
+ _save_requirements: False
+
+callbacks:
+ - _target_: pytorch_lightning.callbacks.LearningRateMonitor
+ logging_interval: step
+ log_momentum: False
+ - _target_: pytorch_lightning.callbacks.ModelCheckpoint
+ monitor: loss_val
+ mode: min
+ save_top_k: 1
+ save_last: True
+ verbose: false
+ every_n_epochs: 1
+ filename: "{epoch}-{loss_val:.2f}"
+ - _target_: pytorch_lightning.callbacks.TQDMProgressBar
+ refresh_rate: 50
+ - _target_: onescience.datapipes.materials.mattergen.callback.SetPropertyScalers
diff --git a/configuration.json b/configuration.json
new file mode 100644
index 0000000000000000000000000000000000000000..3fc809d275f08e4fbca6f5d2b7ff3fb1073aeef9
--- /dev/null
+++ b/configuration.json
@@ -0,0 +1 @@
+{"framework":"Pytorch","task":"other"}
\ No newline at end of file
diff --git a/csv_to_dataset.py b/csv_to_dataset.py
new file mode 100644
index 0000000000000000000000000000000000000000..fbda81f99d68db793c9f05e315a694bdb9267d22
--- /dev/null
+++ b/csv_to_dataset.py
@@ -0,0 +1,47 @@
+"""Convert MatterGen CSV splits into cached CrystalDataset directories."""
+
+import argparse
+from pathlib import Path
+
+from onescience.datapipes.materials.mattergen import CrystalDataset
+
+
+def main() -> None:
+ parser = argparse.ArgumentParser(
+ description="Convert every CSV file in a directory to MatterGen cache format."
+ )
+ parser.add_argument(
+ "--csv-folder",
+ type=Path,
+ required=True,
+ help="Directory containing split files such as train.csv and val.csv.",
+ )
+ parser.add_argument(
+ "--dataset-name",
+ required=True,
+ help="Dataset directory name created below --cache-folder.",
+ )
+ parser.add_argument(
+ "--cache-folder",
+ type=Path,
+ required=True,
+ help="Parent directory in which the dataset cache is created.",
+ )
+ args = parser.parse_args()
+
+ if not args.csv_folder.is_dir():
+ parser.error(f"CSV directory does not exist: {args.csv_folder}")
+
+ csv_files = sorted(args.csv_folder.glob("*.csv"))
+ if not csv_files:
+ parser.error(f"No CSV files found in: {args.csv_folder}")
+
+ dataset_root = args.cache_folder / args.dataset_name
+ for csv_path in csv_files:
+ cache_path = dataset_root / csv_path.stem
+ print(f"Processing {csv_path} -> {cache_path}")
+ CrystalDataset.from_csv(csv_path=csv_path, cache_path=cache_path)
+
+
+if __name__ == "__main__":
+ main()
diff --git a/demo/configs/finetune_dft_mag_density_smoke.yaml b/demo/configs/finetune_dft_mag_density_smoke.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..d2bd154304bb645119c5a25346bdabce562ce208
--- /dev/null
+++ b/demo/configs/finetune_dft_mag_density_smoke.yaml
@@ -0,0 +1,18 @@
+# 属性微调 smoke test。修改属性和 batch size。
+task: finetune
+args:
+ adapter.model_path: ${ONESCIENCE_MODELS_DIR}/mattergen/mattergen_base
+ data_module: mp_20
+ data_module.properties: '[dft_mag_density]'
+ +lightning_module/diffusion_module/model/property_embeddings@adapter.adapter.property_embeddings_adapt.dft_mag_density: dft_mag_density
+ trainer.devices: 1
+ trainer.num_nodes: 1
+ trainer.max_epochs: 1
+ trainer.check_val_every_n_epoch: 1
+ +trainer.limit_train_batches: 1
+ +trainer.limit_val_batches: 1
+ +trainer.num_sanity_val_steps: 0
+ data_module.batch_size.train: 1
+ data_module.batch_size.val: 1
+ ~trainer.logger: true
+ ~trainer.strategy: true
diff --git a/demo/configs/generate_base.yaml b/demo/configs/generate_base.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..19d5eedc9f56c42379de3ee811e505ea7ad5cd8b
--- /dev/null
+++ b/demo/configs/generate_base.yaml
@@ -0,0 +1,7 @@
+# 基础模型单样本生成。修改 checkpoint、output、batch_size 和 num_batches。
+task: generate
+args:
+ checkpoint: ${ONESCIENCE_MODELS_DIR}/mattergen/mattergen_base
+ output: outputs/generate/mattergen_base
+ batch_size: 1
+ num_batches: 1
diff --git a/demo/configs/generate_dft_mag_density.yaml b/demo/configs/generate_dft_mag_density.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..ba701a116a45c9e462f8c064dd2bc8b9d335b386
--- /dev/null
+++ b/demo/configs/generate_dft_mag_density.yaml
@@ -0,0 +1,9 @@
+# 磁密度条件生成。checkpoint 必须与 properties 中的属性匹配。
+task: generate
+args:
+ checkpoint: ${ONESCIENCE_MODELS_DIR}/mattergen/dft_mag_density
+ output: outputs/generate/dft_mag_density
+ batch_size: 1
+ num_batches: 1
+ properties:
+ dft_mag_density: 0.15
diff --git a/demo/configs/train_8dcu.yaml b/demo/configs/train_8dcu.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..ed7fec75c3ba7e608e84df4c547cd5e054727a6b
--- /dev/null
+++ b/demo/configs/train_8dcu.yaml
@@ -0,0 +1,14 @@
+# 8 卡从头训练。优先修改 batch size 和梯度累积。
+task: train
+args:
+ data_module: mp_20
+ trainer.devices: 8
+ trainer.num_nodes: 1
+ trainer.accumulate_grad_batches: 16
+ data_module.batch_size.train: 4
+ data_module.batch_size.val: 4
+ data_module.batch_size.test: 4
+ data_module.num_workers.train: 2
+ data_module.num_workers.val: 2
+ data_module.num_workers.test: 2
+ ~trainer.logger: true
diff --git a/demo/run.sh b/demo/run.sh
new file mode 100644
index 0000000000000000000000000000000000000000..1c4a0fd42669aa9faf92fa0984ca2254c4e2deba
--- /dev/null
+++ b/demo/run.sh
@@ -0,0 +1,61 @@
+#!/bin/bash
+set -euo pipefail
+DEMO_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+CONFIG=""
+MODE="run"
+while [[ $# -gt 0 ]]; do
+ case "$1" in
+ --config) CONFIG="$2"; shift 2 ;;
+ --config=*) CONFIG="${1#*=}"; shift ;;
+ --submit) MODE="submit"; shift ;;
+ -h|--help) echo "Usage: bash run.sh --config configs/.yaml [--submit]"; exit 0 ;;
+ *) echo "Unknown argument: $1" >&2; exit 2 ;;
+ esac
+done
+[[ -n "$CONFIG" ]] || { echo "Please specify --config configs/.yaml" >&2; exit 2; }
+[[ "$CONFIG" = /* ]] || CONFIG="$DEMO_DIR/$CONFIG"
+[[ -f "$CONFIG" ]] || { echo "Config not found: $CONFIG" >&2; exit 2; }
+if [[ -z "${CONDA_PREFIX:-}" || -z "${ONESCIENCE_MODELS_DIR:-}" ]]; then
+ export MATCHEM_CONDA_NAME="${MATCHEM_CONDA_NAME:-onescience-mattergen-source}"
+ source "$DEMO_DIR/../../matchem_env.sh"
+fi
+CMD=$(python3 - "$CONFIG" "$DEMO_DIR/.." <<'PY'
+import os, shlex, sys, yaml
+import json
+cfg = yaml.safe_load(open(sys.argv[1]))
+root = os.path.abspath(sys.argv[2])
+task = cfg.get("task", "train")
+script = {"train": "train.py", "finetune": "finetune.py", "generate": "generate.py"}.get(task)
+if not script: raise SystemExit(f"Unsupported task: {task}")
+args = []
+for key, value in cfg.get("args", {}).items():
+ if isinstance(value, bool):
+ if value: args.append(f"--{key.replace('_', '-')}")
+ elif value is not None:
+ value = json.dumps(value) if isinstance(value, (dict, list)) else str(value)
+ value = os.path.expandvars(value)
+ args.extend([f"--{key.replace('_', '-')}", value] if task == "generate" else [f"{key}={value}"])
+print(" ".join(map(shlex.quote, [sys.executable, os.path.join(root, script), *args])))
+PY
+)
+echo "Running MatterGen: $CMD"
+if [[ "$MODE" == "submit" ]]; then
+ TASK=$(python3 - "$CONFIG" <<'PY'
+import sys, yaml
+print(yaml.safe_load(open(sys.argv[1])).get("task", "train"))
+PY
+)
+ [[ "$TASK" == "train" ]] || { echo "--submit currently supports train configs only" >&2; exit 2; }
+ DEVICES=$(python3 - "$CONFIG" <<'PY'
+import sys, yaml
+print(yaml.safe_load(open(sys.argv[1])).get("args", {}).get("trainer.devices", 8))
+PY
+)
+ sbatch \
+ --gres="dcu:$DEVICES" \
+ --export="ALL,CONFIG_PATH=$CONFIG,SCRIPT_DIR=$DEMO_DIR/.." \
+ "$DEMO_DIR/../submit_train.sh"
+ exit 0
+fi
+cd "$DEMO_DIR/.."
+eval "$CMD"
diff --git a/finetune.py b/finetune.py
new file mode 100644
index 0000000000000000000000000000000000000000..424c5f52b9f52f516de6ebbbad8245b327114261
--- /dev/null
+++ b/finetune.py
@@ -0,0 +1,64 @@
+"""Fine-tune MatterGen with a pretrained checkpoint and property adapter."""
+
+import json
+import os
+from datetime import datetime
+from pathlib import Path
+
+import hydra
+import omegaconf
+import pytorch_lightning as pl
+import torch
+from omegaconf import OmegaConf, open_dict
+from pytorch_lightning.cli import SaveConfigCallback
+
+from model.diffusion.run import (
+ AddConfigCallback,
+ SimpleParser,
+ maybe_instantiate,
+)
+from model.finetune import (
+ init_adapter_lightningmodule_from_pretrained,
+)
+
+EXAMPLE_DIR = Path(__file__).resolve().parent
+os.environ.setdefault(
+ "OUTPUT_DIR",
+ str(EXAMPLE_DIR / "outputs" / "finetune" / datetime.now().strftime("%Y-%m-%d_%H-%M-%S")),
+)
+
+
+@hydra.main(
+ config_path=str(EXAMPLE_DIR / "conf"),
+ config_name="finetune",
+ version_base="1.1",
+)
+def main(cfg: omegaconf.DictConfig) -> None:
+ """Build the data, adapter model, and Trainer, then start fine-tuning."""
+ torch.set_float32_matmul_precision("high")
+ trainer: pl.Trainer = maybe_instantiate(cfg.trainer, pl.Trainer)
+ datamodule: pl.LightningDataModule = maybe_instantiate(
+ cfg.data_module, pl.LightningDataModule
+ )
+
+ model, lightning_module_cfg = init_adapter_lightningmodule_from_pretrained(
+ cfg.adapter, cfg.lightning_module
+ )
+ with open_dict(cfg):
+ cfg.lightning_module = lightning_module_cfg
+
+ resolved_config = OmegaConf.to_container(cfg, resolve=True)
+ print(json.dumps(resolved_config, indent=4))
+ trainer.callbacks.append(
+ SaveConfigCallback(
+ parser=SimpleParser(),
+ config=resolved_config,
+ overwrite=True,
+ )
+ )
+ trainer.callbacks.append(AddConfigCallback(resolved_config))
+ trainer.fit(model=model, datamodule=datamodule, ckpt_path=None)
+
+
+if __name__ == "__main__":
+ main()
diff --git a/generate.py b/generate.py
new file mode 100644
index 0000000000000000000000000000000000000000..7df04e3c3fff48ecd7c0c672d22c0ba51471e6e6
--- /dev/null
+++ b/generate.py
@@ -0,0 +1,44 @@
+import argparse
+import json
+import os
+from pathlib import Path
+
+from model.common.utils.data_classes import MatterGenCheckpointInfo
+from model.generator import CrystalGenerator
+
+
+def main():
+ parser = argparse.ArgumentParser(description="Generate crystals with MatterGen")
+ parser.add_argument("--checkpoint", required=True)
+ parser.add_argument("--output", default="outputs/mattergen")
+ parser.add_argument("--batch-size", type=int, default=1)
+ parser.add_argument("--num-batches", type=int, default=1)
+ parser.add_argument(
+ "--properties",
+ type=json.loads,
+ default=None,
+ help='Condition values as JSON, for example {"dft_mag_density": 0.15}.',
+ )
+ parser.add_argument("--record-trajectories", action="store_true")
+ args = parser.parse_args()
+ if not os.path.isdir(args.checkpoint):
+ parser.error(f"checkpoint directory does not exist: {args.checkpoint}")
+ checkpoint_info = MatterGenCheckpointInfo(
+ model_path=Path(args.checkpoint).expanduser().resolve(),
+ load_epoch="last",
+ )
+ generator = CrystalGenerator(
+ checkpoint_info=checkpoint_info,
+ batch_size=args.batch_size,
+ num_batches=args.num_batches,
+ properties_to_condition_on=args.properties,
+ record_trajectories=args.record_trajectories,
+ )
+ structures = generator.generate(
+ output_dir=Path(args.output).expanduser().resolve()
+ )
+ print(f"Generated {len(structures)} structures in {args.output}")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/model/__init__.py b/model/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..14ad7c27d2edba8411508d450abd0ddaae547601
--- /dev/null
+++ b/model/__init__.py
@@ -0,0 +1,19 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+__version__ = "1.0.0"
+"""MatterGen crystal structure generation model."""
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+ from .generator import CrystalGenerator
+
+__all__ = ["CrystalGenerator"]
+
+
+def __getattr__(name: str):
+ if name == "CrystalGenerator":
+ from .generator import CrystalGenerator
+
+ return CrystalGenerator
+ raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
diff --git a/model/adapter.py b/model/adapter.py
new file mode 100644
index 0000000000000000000000000000000000000000..8b2eb8dc0e26defee086102d9ef88008c98a344f
--- /dev/null
+++ b/model/adapter.py
@@ -0,0 +1,128 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+from typing import Callable
+
+import torch
+
+from onescience.datapipes.materials.mattergen.chemgraph import ChemGraph
+from onescience.datapipes.materials.mattergen.types import PropertySourceId
+from .denoiser import GemNetTDenoiser, get_chemgraph_from_denoiser_output
+from onescience.modules.embedding.mattergen_property_embeddings import (
+ ZerosEmbedding,
+ get_property_embeddings,
+ get_use_unconditional_embedding,
+)
+
+BatchTransform = Callable[[ChemGraph], ChemGraph]
+
+
+class GemNetTAdapter(GemNetTDenoiser):
+ """
+ Denoiser layerwise adapter with GemNetT. On top of a model.denoiser.GemNetTDenoiser,
+ additionally inputs that specifies extra conditions to be conditioned on.
+ """
+
+ def __init__(self, property_embeddings_adapt: torch.nn.ModuleDict, *args, **kwargs):
+ super().__init__(*args, **kwargs)
+
+ # ModuleDict[PropertyName, PropertyEmbedding] -- conditions adding by this adapter
+ self.property_embeddings_adapt = torch.nn.ModuleDict(property_embeddings_adapt)
+
+ # sanity check keys are required by the adapter that already exist in the base model
+ assert all(
+ [
+ k not in self.property_embeddings.keys()
+ for k in self.property_embeddings_adapt.keys()
+ ]
+ ), f"One of adapter conditions {self.property_embeddings_adapt.keys()} already exists in base model {self.property_embeddings.keys()}, please remove."
+
+ # we make the choice that new adapter fields do not alter the unconditional score
+ # we therefore need the unconditional embedding for all properties added in the adapter
+ # to return 0. We hack the unconditional embedding module here to achieve that
+ for property_embedding in self.property_embeddings_adapt.values():
+ property_embedding.unconditional_embedding_module = ZerosEmbedding(
+ hidden_dim=property_embedding.unconditional_embedding_module.hidden_dim,
+ )
+
+ def forward(
+ self,
+ x: ChemGraph,
+ t: torch.Tensor,
+ ) -> ChemGraph:
+ """
+ augment with .
+ """
+ (frac_coords, lattice, atom_types, num_atoms, batch,) = (
+ x["pos"],
+ x["cell"],
+ x["atomic_numbers"],
+ x["num_atoms"],
+ x.get_batch_idx("pos"),
+ )
+ # (num_atoms, hidden_dim) (num_crysts, 3)
+ t_enc = self.noise_level_encoding(t).to(lattice.device)
+ z_per_crystal = t_enc
+
+ # shape = (Nbatch, sum(hidden_dim of all properties in condition_on_adapt))
+ conditions_base_model: torch.Tensor = get_property_embeddings(
+ property_embeddings=self.property_embeddings, batch=x
+ )
+
+ if len(conditions_base_model) > 0:
+ z_per_crystal = torch.cat([z_per_crystal, conditions_base_model], dim=-1)
+
+ # compose into a dict
+ conditions_adapt_dict = {}
+ conditions_adapt_mask_dict = {}
+ for cond_field, property_embedding in self.property_embeddings_adapt.items():
+ conditions_adapt_dict[cond_field] = property_embedding.forward(batch=x)
+ try:
+ conditions_adapt_mask_dict[cond_field] = get_use_unconditional_embedding(
+ batch=x, cond_field=cond_field
+ )
+ except KeyError:
+ # no values have been provided for the conditional field,
+ # interpret this as the user wanting an unconditional score
+ conditions_adapt_mask_dict[cond_field] = torch.ones_like(
+ x["num_atoms"], dtype=torch.bool
+ ).reshape(-1, 1)
+
+ output = self.gemnet(
+ z=z_per_crystal,
+ frac_coords=frac_coords,
+ atom_types=atom_types,
+ num_atoms=num_atoms,
+ batch=batch,
+ lengths=None,
+ angles=None,
+ lattice=lattice,
+ # we construct the graph on the fly, hence pass None for these:
+ edge_index=None,
+ to_jimages=None,
+ num_bonds=None,
+ cond_adapt=conditions_adapt_dict,
+ cond_adapt_mask=conditions_adapt_mask_dict, # when True use unconditional embedding
+ )
+
+ pred_atom_types = self.fc_atom(output.node_embeddings)
+
+ return get_chemgraph_from_denoiser_output(
+ pred_atom_types=pred_atom_types,
+ pred_lattice_eps=output.stress,
+ pred_cart_pos_eps=output.forces,
+ training=self.training,
+ element_mask_func=self.element_mask_func,
+ x_input=x,
+ )
+
+ @property
+ def cond_fields_model_was_trained_on(self) -> list[PropertySourceId]:
+ """
+ We adopt the convention that all property embeddings are stored in torch.nn.ModuleDicts of
+ name property_embeddings or property_embeddings_adapt in the case of a fine tuned model.
+
+ This function returns the list of all field names that a given score model was trained to
+ condition on.
+ """
+ return list(self.property_embeddings) + list(self.property_embeddings_adapt)
diff --git a/model/common/__init__.py b/model/common/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..f63beed9d314ccd7ca1ae4705b6701b98dc89b92
--- /dev/null
+++ b/model/common/__init__.py
@@ -0,0 +1,3 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+__version__ = "0.1.0"
diff --git a/model/common/diffusion/__init__.py b/model/common/diffusion/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/model/common/diffusion/corruption.py b/model/common/diffusion/corruption.py
new file mode 100644
index 0000000000000000000000000000000000000000..252b2c9383358d167777ab71a88a2eeed14eea77
--- /dev/null
+++ b/model/common/diffusion/corruption.py
@@ -0,0 +1,284 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+import torch
+from omegaconf import DictConfig
+
+from ...diffusion.corruption.corruption import B, BatchedData, maybe_expand
+from ...diffusion.corruption.sde_lib import SDE as DiffSDE
+from ...diffusion.corruption.sde_lib import VESDE as DiffVESDE
+from ...diffusion.corruption.sde_lib import VPSDE
+from ...diffusion.wrapped.wrapped_sde import WrappedVESDE
+
+
+def expand(a, x_shape, left=False):
+ a_dim = len(a.shape)
+ if left:
+ return a.reshape(*(((1,) * (len(x_shape) - a_dim)) + a.shape))
+ else:
+ return a.reshape(*(a.shape + ((1,) * (len(x_shape) - a_dim))))
+
+
+def make_noise_symmetric_preserve_variance(noise: torch.Tensor) -> torch.Tensor:
+ """Makes the noise matrix symmetric, preserving the variance. Assumes i.i.d. noise for each dimension.
+
+ Args:
+ noise (torch.Tensor): Input noise matrix, must be a batched square matrix, i.e., have shape (batch_size, dim, dim).
+
+ Returns:
+ torch.Tensor: The symmetric noise matrix, with the same variance as the input.
+ """
+ assert (
+ len(noise.shape) == 3 and noise.shape[1] == noise.shape[2]
+ ), "Symmetric noise only works for square-matrix-shaped data."
+ # Var[1/sqrt(2) * (eps_i + eps_j)] = 0.5 Var[eps_i] + 0.5 Var[eps_j] = Var[noise]
+ # Special treatment of the diagonal elements, i.e., those we leave unchanged via masking.
+ return (1 / (2**0.5)) * (1 - torch.eye(3, device=noise.device)[None]) * (
+ noise + noise.transpose(1, 2)
+ ) + torch.eye(3, device=noise.device)[None] * noise
+
+
+class LatticeVPSDE(VPSDE):
+ @staticmethod
+ def from_vpsde_config(vpsde_config: DictConfig):
+ return LatticeVPSDE(
+ **vpsde_config,
+ )
+
+ def __init__(
+ self,
+ beta_min: float = 0.1,
+ beta_max: float = 20,
+ limit_density: float | None = 0.05,
+ limit_var_scaling_constant: float = 0.25,
+ **kwargs,
+ ):
+ """Variance-preserving SDE with drift coefficient changing linearly over time."""
+ super().__init__()
+ self.beta_0 = beta_min
+ self.beta_1 = beta_max
+
+ # each crystal is diffused to have expected lattice vectors
+ # based on the number of atoms per crystal and self.limit_density
+ # units=(atoms/Angstrom**3)
+ self.limit_density = limit_density
+ self.limit_var_scaling_constant = limit_var_scaling_constant
+
+ self._limit_info_key = "num_atoms"
+
+ @property
+ def limit_info_key(self) -> str:
+ return self._limit_info_key
+
+ def beta(self, t: torch.Tensor) -> torch.Tensor:
+ return self.beta_0 + t * (self.beta_1 - self.beta_0)
+
+ def _marginal_mean_coeff(self, t: torch.Tensor) -> torch.Tensor:
+ log_mean_coeff = -0.25 * t**2 * (self.beta_1 - self.beta_0) - 0.5 * t * self.beta_0
+ return torch.exp(log_mean_coeff)
+
+ def marginal_prob(
+ self,
+ x: torch.Tensor,
+ t: torch.Tensor,
+ batch_idx: B = None,
+ batch: BatchedData | None = None,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ assert batch is not None
+ mean_coeff = self._marginal_mean_coeff(t)
+ # x: shape [batch_size, *x.shape[1:]]
+ # t, limit_info: shape [batch_size,]
+ limit_mean = self.get_limit_mean(x=x, batch=batch)
+ limit_var = self.get_limit_var(x=x, batch=batch)
+ mean_coeff_expanded = maybe_expand(mean_coeff, batch_idx, x)
+ mean = mean_coeff_expanded * x + (1 - mean_coeff_expanded) * limit_mean
+ std = torch.sqrt((1.0 - mean_coeff_expanded**2) * limit_var)
+ return mean, std
+
+ def mean_coeff_and_std(
+ self,
+ x: torch.Tensor,
+ t: torch.Tensor,
+ batch_idx: B = None,
+ batch: BatchedData | None = None,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ """Returns mean coefficient and standard deviation of marginal distribution at time t."""
+ mean_coeff = self._marginal_mean_coeff(t)
+ std = self.marginal_prob(x, t, batch_idx, batch)[1]
+ return maybe_expand(mean_coeff, batch=None, like=x), std
+
+ def get_limit_mean(self, x: torch.Tensor, batch: BatchedData) -> torch.Tensor:
+ # x: shape [batch_size, *x.shape[1:]]
+ # limit_info: shape [batch_size,], a 1d tensor containing number of atoms per crystal
+ # self.limit_density = limit_info / mean lattice vector length**3
+
+ # shape=[Ncrystals,]
+ n_atoms = batch[self.limit_info_key]
+
+ # shape=[Ncrystals, 3, 3]
+ return torch.pow(
+ torch.eye(3, device=x.device).expand(len(n_atoms), 3, 3)
+ * n_atoms[:, None, None]
+ / self.limit_density,
+ 1.0 / 3,
+ ).to(x.device)
+
+ def get_limit_var(self, x: torch.Tensor, batch: BatchedData) -> torch.Tensor:
+ """
+ Returns the element-wise variance of the limit distribution.
+ NOTE: even though we have a different limit variance per data
+ dimension we still sample IID for each element per data point.
+ We do NOT do any correlated sampling over data dimensions per
+ data point.
+
+ Return shape=x.shape
+ """
+
+ # x: shape [batch_size, *x.shape[1:]]
+ # limit_info: shape [batch_size,]
+ # necessary for mypy
+ n_atoms = batch[self.limit_info_key]
+
+ # expand to fit shape of data, shape = (n_crystals, 1, 1)
+ n_atoms_expanded = expand(n_atoms, x.shape)
+
+ # shape = (n_crystals, 3, 3)
+ n_atoms_expanded = torch.tile(n_atoms_expanded, (1, 3, 3))
+
+ # scale limit standard deviation to be proportional to number atoms = n_atoms**(1/3)
+ # per lattice vector. We hope that prod_i std_i scales as the standard deviation
+ # of the actual volume. NOTE: we return variance here, hence 2 in the power
+ # shape=(Ncrystals, 3, 3) for limit_info.shape=[Ncrystals,]
+ out = torch.pow(n_atoms_expanded, 2.0 / 3).to(x.device) * self.limit_var_scaling_constant
+
+ return out
+
+ def sample_marginal(
+ self,
+ x: torch.Tensor,
+ t: torch.Tensor,
+ batch_idx: B = None,
+ batch: BatchedData | None = None,
+ ) -> torch.Tensor:
+ mean, std = self.marginal_prob(x=x, t=t, batch=batch)
+ z = torch.randn_like(x)
+ z = make_noise_symmetric_preserve_variance(z)
+ return mean + expand(std, z.shape) * z
+
+ def prior_sampling(
+ self,
+ shape: torch.Size | tuple,
+ conditioning_data: BatchedData | None = None,
+ batch_idx: B = None,
+ ) -> torch.Tensor:
+ x_sample = torch.randn(*shape)
+ x_sample = make_noise_symmetric_preserve_variance(x_sample)
+ assert conditioning_data is not None
+ limit_info = conditioning_data[self.limit_info_key]
+ x_sample = x_sample.to(limit_info.device)
+ limit_mean = self.get_limit_mean(x=x_sample, batch=conditioning_data)
+ limit_var = self.get_limit_var(x=x_sample, batch=conditioning_data)
+ # shape=[Nbatch,...] for shape[0]=Nbatch
+ return x_sample * limit_var.sqrt() + limit_mean
+
+ def sde(
+ self,
+ x: torch.Tensor,
+ t: torch.Tensor,
+ batch_idx: B = None,
+ batch: BatchedData | None = None,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ assert batch is not None
+ # x: shape [batch_size, *x.shape[1:]]
+ # t, limit_info: shape [batch_size,]
+ # if a per data-point limit mean is supplied, expand to shape of data
+ # shape=x.shape
+ limit_mean = self.get_limit_mean(x=x, batch=batch)
+
+ # if a per data-point limit variance is supplied, shape=[x.shape[0], ]
+ limit_var = self.get_limit_var(x=x, batch=batch)
+
+ beta_t = self.beta(t)
+ drift = (
+ -0.5
+ * expand(
+ beta_t,
+ x.shape,
+ )
+ * (x - limit_mean)
+ )
+ diffusion = torch.sqrt(expand(beta_t, limit_var.shape) * limit_var)
+ # drift.shape=[Nbatch,...], diffusion.shape=[Nbatch,] for x.shape[0]=Nbatch
+ return maybe_expand(drift, batch_idx), maybe_expand(diffusion, batch_idx)
+
+
+class NumAtomsVarianceAdjustedWrappedVESDE(WrappedVESDE):
+ """Wrapped VESDE with variance adjusted by number of atoms. We divide the standard deviation by the cubic root of the number of atoms.
+ The goal is to reduce the influence by the cell size on the variance of the fractional coordinates.
+ """
+
+ def __init__(
+ self,
+ wrapping_boundary: float | torch.Tensor = 1.0,
+ sigma_min: float = 0.01,
+ sigma_max: float = 5.0,
+ limit_info_key: str = "num_atoms",
+ ):
+ super().__init__(
+ sigma_min=sigma_min, sigma_max=sigma_max, wrapping_boundary=wrapping_boundary
+ )
+ self.limit_info_key = limit_info_key
+
+ def std_scaling(self, batch: BatchedData) -> torch.Tensor:
+ return batch[self.limit_info_key] ** (-1 / 3)
+
+ def marginal_prob(
+ self,
+ x: torch.Tensor,
+ t: torch.Tensor,
+ batch_idx: B = None,
+ batch: BatchedData | None = None,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ mean, std = super().marginal_prob(x, t, batch_idx, batch)
+ assert (
+ batch is not None
+ ), "batch must be provided when using NumAtomsVarianceAdjustedWrappedVESDEMixin"
+ std_scale = self.std_scaling(batch)
+ std = std * maybe_expand(std_scale, batch_idx, like=std)
+ return mean, std
+
+ def prior_sampling(
+ self,
+ shape: torch.Size | tuple,
+ conditioning_data: BatchedData | None = None,
+ batch_idx=None,
+ ) -> torch.Tensor:
+ _super = super()
+ assert isinstance(self, DiffSDE) and hasattr(_super, "prior_sampling")
+ assert (
+ conditioning_data is not None
+ ), "batch must be provided when using NumAtomsVarianceAdjustedWrappedVESDEMixin"
+ num_atoms = conditioning_data[self.limit_info_key]
+ batch_idx = torch.repeat_interleave(
+ torch.arange(num_atoms.shape[0], device=num_atoms.device), num_atoms, dim=0
+ )
+ std_scale = self.std_scaling(conditioning_data)
+ # prior sample is randn() * sigma_max, so we need additionally multiply by std_scale to get the correct variance.
+ # We call VESDE.prior_sampling (a "grandparent" function) because the super() prior_sampling already does the wrapping,
+ # which means we couldn't do the variance adjustment here anymore otherwise.
+ prior_sample = DiffVESDE.prior_sampling(self, shape=shape).to(num_atoms.device)
+ return self.wrap(prior_sample * maybe_expand(std_scale, batch_idx, like=prior_sample))
+
+ def sde(
+ self,
+ x: torch.Tensor,
+ t: torch.Tensor,
+ batch_idx: B = None,
+ batch: BatchedData | None = None,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ sigma = self.marginal_prob(x, t, batch_idx, batch)[1]
+ sigma_min = self.marginal_prob(x, torch.zeros_like(t), batch_idx, batch)[1]
+ sigma_max = self.marginal_prob(x, torch.ones_like(t), batch_idx, batch)[1]
+ drift = torch.zeros_like(x)
+ diffusion = sigma * torch.sqrt(2 * (sigma_max.log() - sigma_min.log()))
+ return drift, diffusion
diff --git a/model/common/diffusion/predictors_correctors.py b/model/common/diffusion/predictors_correctors.py
new file mode 100644
index 0000000000000000000000000000000000000000..d69451bdfed3c13caeeb6fcae01dde68c71dfa85
--- /dev/null
+++ b/model/common/diffusion/predictors_correctors.py
@@ -0,0 +1,98 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+import torch
+
+from ...common.diffusion import corruption as sde_lib
+from ...common.utils.data_utils import compute_lattice_polar_decomposition
+from ...diffusion.corruption.corruption import Corruption, maybe_expand
+from ...diffusion.data.batched_data import BatchedData
+from ...diffusion.sampling import predictors_correctors as pc
+from ...diffusion.sampling.predictors import AncestralSamplingPredictor
+
+SampleAndMean = tuple[torch.Tensor, torch.Tensor]
+
+
+class LatticeAncestralSamplingPredictor(AncestralSamplingPredictor):
+ @classmethod
+ def is_compatible(cls, corruption: Corruption) -> bool:
+ _super = super()
+ assert hasattr(_super, "is_compatible")
+ return _super.is_compatible(corruption) or isinstance(corruption, sde_lib.LatticeVPSDE)
+
+ def update_given_score(
+ self,
+ *,
+ x: torch.Tensor,
+ t: torch.Tensor,
+ dt: torch.Tensor,
+ batch_idx: torch.LongTensor,
+ score: torch.Tensor,
+ batch: BatchedData | None,
+ ) -> SampleAndMean:
+ x_coeff, score_coeff, std = self._get_coeffs(
+ x=x,
+ t=t,
+ dt=dt,
+ batch_idx=batch_idx,
+ batch=batch,
+ )
+ # mean = (x + score * beta**2 - limit_mean)/(1-beta) + limit_mean
+ # <=> mean = x / (1-beta) + score * beta**2 / (1-beta) + limit_mean * (1 - 1/(1-beta))
+ # => mean_coeff = 1 - x_coeff = 1 - 1/(1-beta)
+ mean_coeff = 1 - x_coeff
+ # Sample random noise.
+ z = sde_lib.make_noise_symmetric_preserve_variance(torch.randn_like(x_coeff))
+ assert hasattr(self.corruption, "get_limit_mean") # mypy
+ mean = (
+ x_coeff * x
+ + score_coeff * score
+ + mean_coeff * self.corruption.get_limit_mean(x=x, batch=batch)
+ )
+ sample = mean + std * z
+ return sample, mean
+
+
+# create a langevin corrector that accepts LatticeVPSDE
+class LatticeLangevinDiffCorrector(pc.LangevinCorrector):
+ @classmethod
+ def is_compatible(cls, corruption: Corruption) -> bool:
+ _super = super()
+ assert hasattr(_super, "is_compatible")
+ return _super.is_compatible(corruption) or isinstance(corruption, sde_lib.LatticeVPSDE)
+
+ def step_given_score(
+ self,
+ *,
+ x: torch.Tensor,
+ batch_idx: torch.LongTensor | None,
+ score: torch.Tensor,
+ t: torch.Tensor,
+ dt: torch.Tensor,
+ ) -> SampleAndMean:
+ assert isinstance(self.corruption, sde_lib.LatticeVPSDE)
+ alpha = self.get_alpha(t, dt=dt)
+ snr = self.snr
+ noise = torch.randn_like(x)
+ noise = sde_lib.make_noise_symmetric_preserve_variance(noise)
+
+ # [batch_size, ] or [num_atoms, ] if batch_idx is not None
+ grad_norm_square = torch.square(score).reshape(score.shape[0], -1).sum(dim=1)
+ noise_norm_square = torch.square(noise).reshape(noise.shape[0], -1).sum(dim=1)
+ # Average over items, leading to scalars.
+ grad_norm = grad_norm_square.sqrt().mean()
+ noise_norm = noise_norm_square.sqrt().mean()
+
+ # If gradient is zero (i.e., we are sampling from an improper distribution that's flat over the whole of R^n)
+ # the step_size blows up. Clip step_size to avoid this.
+ # The EGNN reports zero scores when there are no edges between nodes.
+ step_size = (snr * noise_norm / grad_norm) ** 2 * 2 * alpha
+ step_size = torch.minimum(step_size, self.max_step_size)
+ step_size[grad_norm == 0, :] = self.max_step_size
+ step_size = maybe_expand(step_size, batch_idx, score)
+ mean = x + step_size * score
+ x = mean + torch.sqrt(step_size * 2) * noise
+
+ x = compute_lattice_polar_decomposition(x)
+ mean = compute_lattice_polar_decomposition(mean)
+ return x, mean
diff --git a/model/common/gemnet/__init__.py b/model/common/gemnet/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/model/common/gemnet/cgmanifest.json b/model/common/gemnet/cgmanifest.json
new file mode 100644
index 0000000000000000000000000000000000000000..9e50fd5999ef2b57958f824ee84cefd327243085
--- /dev/null
+++ b/model/common/gemnet/cgmanifest.json
@@ -0,0 +1,16 @@
+{
+ "$schema": "https://json.schemastore.org/component-detection-manifest.json",
+ "version": 1,
+ "registrations":[
+ {
+ "component": {
+ "type": "git",
+ "git": {
+ "repositoryUrl": "https://github.com/FAIR-Chem/fairchem",
+ "commitHash": "65c2d6246e69169f43949858d39550d2a635c7e0"
+ }
+ },
+ "developmentDependency" : false
+ }
+ ]
+}
\ No newline at end of file
diff --git a/model/common/gemnet/gemnet-dT.json b/model/common/gemnet/gemnet-dT.json
new file mode 100644
index 0000000000000000000000000000000000000000..98aac8b55c33dc42064a4e7d740e5ec298a48109
--- /dev/null
+++ b/model/common/gemnet/gemnet-dT.json
@@ -0,0 +1,20 @@
+{
+ "comment": "tri_gaussian128, from https://github.com/FAIR-Chem/fairchem/blob/main/configs/s2ef/all/gemnet/scaling_factors/gemnet-dT.json",
+ "TripInteraction_1_had_rbf": 18.873615264892578,
+ "TripInteraction_1_sum_cbf": 7.996850490570068,
+ "AtomUpdate_1_sum": 1.220463752746582,
+ "TripInteraction_2_had_rbf": 16.10817527770996,
+ "TripInteraction_2_sum_cbf": 7.614634037017822,
+ "AtomUpdate_2_sum": 0.9690994620323181,
+ "TripInteraction_3_had_rbf": 15.01930046081543,
+ "TripInteraction_3_sum_cbf": 7.025179862976074,
+ "AtomUpdate_3_sum": 0.8903237581253052,
+ "OutBlock_0_sum": 1.6437848806381226,
+ "OutBlock_0_had": 16.161039352416992,
+ "OutBlock_1_sum": 1.1077653169631958,
+ "OutBlock_1_had": 13.54678726196289,
+ "OutBlock_2_sum": 0.9477927684783936,
+ "OutBlock_2_had": 12.754337310791016,
+ "OutBlock_3_sum": 0.9059251546859741,
+ "OutBlock_3_had": 13.484951972961426
+}
diff --git a/model/common/gemnet/gemnet.py b/model/common/gemnet/gemnet.py
new file mode 100644
index 0000000000000000000000000000000000000000..edf3c17e27bbb55c90496aed787d456d75e7a53f
--- /dev/null
+++ b/model/common/gemnet/gemnet.py
@@ -0,0 +1,778 @@
+"""
+Copyright (c) Facebook, Inc. and its affiliates.
+Copyright (c) Microsoft Corporation.
+Licensed under the MIT License.
+Adapted from https://github.com/FAIR-Chem/fairchem/blob/main/src/fairchem/core/models/gemnet/gemnet.py.
+"""
+
+from dataclasses import dataclass
+from typing import Optional, Tuple
+
+# import numpy as np
+import torch
+import torch.nn as nn
+from torch_scatter import scatter
+from torch_sparse import SparseTensor
+
+from onescience.modules.layer.mattergen.atom_update_block import OutputBlock
+from onescience.modules.layer.mattergen.base_layers import Dense
+from onescience.modules.layer.mattergen.efficient import EfficientInteractionDownProjection
+from onescience.modules.layer.mattergen.embedding_block import EdgeEmbedding
+from onescience.modules.layer.mattergen.interaction_block import InteractionBlockTripletsOnly
+from onescience.modules.layer.mattergen.radial_basis import RadialBasis
+from onescience.modules.layer.mattergen.scaling import AutomaticFit
+from onescience.modules.layer.mattergen.spherical_basis import CircularBasisLayer
+from ...common.gemnet.utils import (
+ inner_product_normalized,
+ mask_neighbors,
+ ragged_range,
+ repeat_blocks,
+)
+from ...common.utils.data_utils import (
+ frac_to_cart_coords_with_lattice,
+ get_pbc_distances,
+ lattice_params_to_matrix_torch,
+ radius_graph_pbc,
+)
+from ...common.utils.globals import MODELS_PROJECT_ROOT, get_device, get_pyg_device
+from ...common.utils.lattice_score import edge_score_to_lattice_score_frac_symmetric
+
+
+@dataclass(frozen=True)
+class ModelOutput:
+ energy: torch.Tensor
+ node_embeddings: torch.Tensor
+ forces: Optional[torch.Tensor] = None
+ stress: Optional[torch.Tensor] = None
+
+
+class RBFBasedLatticeUpdateBlock(torch.nn.Module):
+ # Lattice update block that mimics GemNet's edge processing, e.g., uses radial basis functions.
+ def __init__(
+ self,
+ emb_size: int,
+ activation: str,
+ emb_size_rbf: int,
+ emb_size_edge: int,
+ num_heads: int = 1,
+ ):
+ super().__init__()
+ self.num_out = num_heads
+ self.mlp = nn.Sequential(
+ Dense(emb_size, emb_size, activation=activation), Dense(emb_size, emb_size)
+ )
+ self.dense_rbf_F = Dense(emb_size_rbf, emb_size_edge, activation=None, bias=False)
+ self.out_forces = Dense(emb_size_edge, num_heads, bias=False, activation=None)
+
+ def compute_score_per_edge(
+ self,
+ edge_emb: torch.Tensor, # [Num_edges, emb_dim]
+ rbf: torch.Tensor, # [Num_edges, num_rbf_bases]
+ ) -> torch.Tensor:
+ x_F = self.mlp(edge_emb)
+ rbf_emb_F = self.dense_rbf_F(rbf) # (nEdges, emb_size_edge)
+ x_F_rbf = x_F * rbf_emb_F
+ # x_F = self.scale_rbf_F(x_F, x_F_rbf)
+ x_F = self.out_forces(x_F_rbf) # (nEdges, self.num_out)
+ return x_F
+
+
+class RBFBasedLatticeUpdateBlockFrac(RBFBasedLatticeUpdateBlock):
+ # Lattice update block that mimics GemNet's edge processing, e.g., uses radial basis functions.
+ def __init__(
+ self,
+ emb_size: int,
+ activation: str,
+ emb_size_rbf: int,
+ emb_size_edge: int,
+ num_heads: int = 1,
+ ):
+ super().__init__(
+ emb_size=emb_size,
+ activation=activation,
+ emb_size_rbf=emb_size_rbf,
+ emb_size_edge=emb_size_edge,
+ num_heads=num_heads,
+ )
+
+ def forward(
+ self,
+ edge_emb: torch.Tensor, # [Num_edges, emb_dim]
+ edge_index: torch.Tensor, # [2, Num_edges]
+ distance_vec: torch.Tensor, # [Num_edges, 3]
+ lattice: torch.Tensor, # [Num_crystals, 3, 3]
+ batch: torch.Tensor, # [Num_atoms, ]
+ rbf: torch.Tensor, # [Num_edges, num_rbf_bases]
+ normalize_score: bool = True,
+ ) -> torch.Tensor:
+ edge_scores = self.compute_score_per_edge(edge_emb=edge_emb, rbf=rbf)
+ if normalize_score:
+ num_edges = scatter(torch.ones_like(distance_vec[:, 0]), batch[edge_index[0]])
+ edge_scores /= num_edges[batch[edge_index[0]], None]
+ outs = []
+ for i in range(self.num_out):
+ lattice_update = edge_score_to_lattice_score_frac_symmetric(
+ score_d=edge_scores[:, i],
+ edge_index=edge_index,
+ edge_vectors=distance_vec,
+ batch=batch,
+ )
+ outs.append(lattice_update)
+ outs = torch.stack(outs, dim=-1).sum(-1)
+ # [Batch_size, 3, 3]
+ return outs
+
+
+class GemNetT(torch.nn.Module):
+ """
+ GemNet-T, triplets-only variant of GemNet
+
+ Parameters
+ ----------
+ num_targets: int
+ Number of prediction targets.
+
+ num_spherical: int
+ Controls maximum frequency.
+ num_radial: int
+ Controls maximum frequency.
+ num_blocks: int
+ Number of building blocks to be stacked.
+
+ atom_embedding: torch.nn.Module
+ a module that embeds atomic numbers into vectors of size emb_dim_atomic_number.
+ emb_size_atom: int
+ Embedding size of the atoms. This can be different from emb_dim_atomic_number.
+ emb_size_edge: int
+ Embedding size of the edges.
+ emb_size_trip: int
+ (Down-projected) Embedding size in the triplet message passing block.
+ emb_size_rbf: int
+ Embedding size of the radial basis transformation.
+ emb_size_cbf: int
+ Embedding size of the circular basis transformation (one angle).
+ emb_size_bil_trip: int
+ Embedding size of the edge embeddings in the triplet-based message passing block after the bilinear layer.
+ num_before_skip: int
+ Number of residual blocks before the first skip connection.
+ num_after_skip: int
+ Number of residual blocks after the first skip connection.
+ num_concat: int
+ Number of residual blocks after the concatenation.
+ num_atom: int
+ Number of residual blocks in the atom embedding blocks.
+ cutoff: float
+ Embedding cutoff for interactomic directions in Angstrom.
+ rbf: dict
+ Name and hyperparameters of the radial basis function.
+ envelope: dict
+ Name and hyperparameters of the envelope function.
+ cbf: dict
+ Name and hyperparameters of the cosine basis function.
+ output_init: str
+ Initialization method for the final dense layer.
+ activation: str
+ Name of the activation function.
+ scale_file: str
+ Path to the json file containing the scaling factors.
+ encoder_mode: bool
+ if , use the encoder mode of the model, i.e. only get the atom/edge embedddings.
+ """
+
+ def __init__(
+ self,
+ num_targets: int,
+ latent_dim: int,
+ atom_embedding: torch.nn.Module,
+ num_spherical: int = 7,
+ num_radial: int = 128,
+ num_blocks: int = 3,
+ emb_size_atom: int = 512,
+ emb_size_edge: int = 512,
+ emb_size_trip: int = 64,
+ emb_size_rbf: int = 16,
+ emb_size_cbf: int = 16,
+ emb_size_bil_trip: int = 64,
+ num_before_skip: int = 1,
+ num_after_skip: int = 2,
+ num_concat: int = 1,
+ num_atom: int = 3,
+ regress_stress: bool = False,
+ cutoff: float = 6.0,
+ max_neighbors: int = 50,
+ rbf: dict = {"name": "gaussian"},
+ envelope: dict = {"name": "polynomial", "exponent": 5},
+ cbf: dict = {"name": "spherical_harmonics"},
+ otf_graph: bool = False,
+ output_init: str = "HeOrthogonal",
+ activation: str = "swish",
+ max_cell_images_per_dim: int = 5,
+ encoder_mode: bool = False, #
+ **kwargs,
+ ):
+ super().__init__()
+ scale_file = f"{MODELS_PROJECT_ROOT}/common/gemnet/gemnet-dT.json"
+ assert scale_file is not None, "`scale_file` is required."
+
+ self.encoder_mode = encoder_mode
+ self.num_targets = num_targets
+ assert num_blocks > 0
+ self.num_blocks = num_blocks
+ emb_dim_atomic_number = getattr(atom_embedding, "emb_size")
+
+ self.cutoff = cutoff
+
+ self.max_neighbors = max_neighbors
+
+ self.max_cell_images_per_dim = max_cell_images_per_dim
+
+ self.otf_graph = otf_graph
+
+ self.regress_stress = regress_stress
+ # we might want to take care of permutation invariance w.r.t. the order of the lattice vectors, though I don't think this is critical.
+ self.angle_edge_emb = nn.Sequential(
+ nn.Linear(emb_size_edge + 3, emb_size_edge),
+ nn.ReLU(),
+ nn.Linear(emb_size_edge, emb_size_edge),
+ )
+
+ AutomaticFit.reset() # make sure that queue is empty (avoid potential error)
+
+ # ---------------------------------- Basis Functions ---------------------------------- ###
+ self.radial_basis = RadialBasis(
+ num_radial=num_radial,
+ cutoff=cutoff,
+ rbf=rbf,
+ envelope=envelope,
+ )
+
+ radial_basis_cbf3 = RadialBasis(
+ num_radial=num_radial,
+ cutoff=cutoff,
+ rbf=rbf,
+ envelope=envelope,
+ )
+ self.cbf_basis3 = CircularBasisLayer(
+ num_spherical,
+ radial_basis=radial_basis_cbf3,
+ cbf=cbf,
+ efficient=True,
+ )
+ # ------------------------------------------------------------------------------------- ###
+
+ # --------------------------------- Update lattice MLP -------------------------------- ###
+ self.regress_stress = regress_stress
+ self.lattice_out_blocks = nn.ModuleList(
+ [
+ RBFBasedLatticeUpdateBlockFrac(
+ emb_size_edge,
+ activation,
+ emb_size_rbf,
+ emb_size_edge,
+ )
+ for _ in range(num_blocks + 1)
+ ]
+ )
+ self.mlp_rbf_lattice = Dense(
+ num_radial,
+ emb_size_rbf,
+ activation=None,
+ bias=False,
+ )
+ # ------------------------------------------------------------------------------------- ###
+ # ------------------------------- Share Down Projections ------------------------------ ###
+ # Share down projection across all interaction blocks
+ self.mlp_rbf3 = Dense(
+ num_radial,
+ emb_size_rbf,
+ activation=None,
+ bias=False,
+ )
+ self.mlp_cbf3 = EfficientInteractionDownProjection(num_spherical, num_radial, emb_size_cbf)
+
+ # Share the dense Layer of the atom embedding block across the interaction blocks
+ self.mlp_rbf_h = Dense(
+ num_radial,
+ emb_size_rbf,
+ activation=None,
+ bias=False,
+ )
+ self.mlp_rbf_out = Dense(
+ num_radial,
+ emb_size_rbf,
+ activation=None,
+ bias=False,
+ )
+ # ------------------------------------------------------------------------------------- ###
+
+ self.atom_emb = atom_embedding
+ self.atom_latent_emb = nn.Linear(emb_dim_atomic_number + latent_dim, emb_size_atom)
+ self.edge_emb = EdgeEmbedding(
+ emb_size_atom, num_radial, emb_size_edge, activation=activation
+ )
+
+ out_blocks = []
+ int_blocks = []
+
+ # Interaction Blocks
+ interaction_block = InteractionBlockTripletsOnly # GemNet-(d)T
+ for i in range(num_blocks):
+ int_blocks.append(
+ interaction_block(
+ emb_size_atom=emb_size_atom,
+ emb_size_edge=emb_size_edge,
+ emb_size_trip=emb_size_trip,
+ emb_size_rbf=emb_size_rbf,
+ emb_size_cbf=emb_size_cbf,
+ emb_size_bil_trip=emb_size_bil_trip,
+ num_before_skip=num_before_skip,
+ num_after_skip=num_after_skip,
+ num_concat=num_concat,
+ num_atom=num_atom,
+ activation=activation,
+ scale_file=scale_file,
+ name=f"IntBlock_{i+1}",
+ )
+ )
+
+ for i in range(num_blocks + 1):
+ out_blocks.append(
+ OutputBlock(
+ emb_size_atom=emb_size_atom,
+ emb_size_edge=emb_size_edge,
+ emb_size_rbf=emb_size_rbf,
+ nHidden=num_atom,
+ num_targets=num_targets,
+ activation=activation,
+ output_init=output_init,
+ direct_forces=True,
+ scale_file=scale_file,
+ name=f"OutBlock_{i}",
+ )
+ )
+
+ self.out_blocks = torch.nn.ModuleList(out_blocks)
+ self.int_blocks = torch.nn.ModuleList(int_blocks)
+
+ self.shared_parameters = [
+ (self.mlp_rbf3, self.num_blocks),
+ (self.mlp_cbf3, self.num_blocks),
+ (self.mlp_rbf_h, self.num_blocks),
+ (self.mlp_rbf_out, self.num_blocks + 1),
+ ]
+
+ def get_triplets(
+ self, edge_index: torch.Tensor, num_atoms: int
+ ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+ """
+ Get all b->a for each edge c->a.
+ It is possible that b=c, as long as the edges are distinct.
+
+ Returns
+ -------
+ id3_ba: torch.Tensor, shape (num_triplets,)
+ Indices of input edge b->a of each triplet b->a<-c
+ id3_ca: torch.Tensor, shape (num_triplets,)
+ Indices of output edge c->a of each triplet b->a<-c
+ id3_ragged_idx: torch.Tensor, shape (num_triplets,)
+ Indices enumerating the copies of id3_ca for creating a padded matrix
+ """
+ idx_s, idx_t = edge_index # c->a (source=c, target=a)
+
+ value = torch.arange(idx_s.size(0), device=idx_s.device, dtype=idx_s.dtype)
+ # Possibly contains multiple copies of the same edge (for periodic interactions)
+ pyg_device = get_pyg_device() if idx_s.device != torch.device("cpu") else idx_s.device
+ torch_device = get_device() if idx_s.device != torch.device("cpu") else idx_s.device
+ adj = SparseTensor(
+ row=idx_t.to(pyg_device),
+ col=idx_s.to(pyg_device),
+ value=value.to(pyg_device),
+ sparse_sizes=(num_atoms.to(pyg_device), num_atoms.to(pyg_device)),
+ )
+ adj_edges = adj[idx_t.to(pyg_device)].to(torch_device)
+
+ # Edge indices (b->a, c->a) for triplets.
+ id3_ba = adj_edges.storage.value().to(torch_device)
+ id3_ca = adj_edges.storage.row().to(torch_device)
+
+ # Remove self-loop triplets
+ # Compare edge indices, not atom indices to correctly handle periodic interactions
+ mask = id3_ba != id3_ca
+ id3_ba = id3_ba[mask]
+ id3_ca = id3_ca[mask]
+
+ # Get indices to reshape the neighbor indices b->a into a dense matrix.
+ # id3_ca has to be sorted for this to work.
+ num_triplets = torch.bincount(id3_ca, minlength=idx_s.size(0))
+ id3_ragged_idx = ragged_range(num_triplets)
+
+ return id3_ba, id3_ca, id3_ragged_idx
+
+ def select_symmetric_edges(self, tensor, mask, reorder_idx, inverse_neg):
+ # Mask out counter-edges
+ tensor_directed = tensor[mask]
+ # Concatenate counter-edges after normal edges
+ sign = 1 - 2 * inverse_neg
+ tensor_cat = torch.cat([tensor_directed, sign * tensor_directed])
+ # Reorder everything so the edges of every image are consecutive
+ tensor_ordered = tensor_cat[reorder_idx]
+ return tensor_ordered
+
+ def reorder_symmetric_edges(
+ self,
+ edge_index: torch.Tensor,
+ cell_offsets: torch.Tensor,
+ neighbors: torch.Tensor,
+ edge_dist: torch.Tensor,
+ edge_vector: torch.Tensor,
+ ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+ """
+ Reorder edges to make finding counter-directional edges easier.
+
+ Some edges are only present in one direction in the data,
+ since every atom has a maximum number of neighbors. Since we only use i->j
+ edges here, we lose some j->i edges and add others by
+ making it symmetric.
+ We could fix this by merging edge_index with its counter-edges,
+ including the cell_offsets, and then running torch.unique.
+ But this does not seem worth it.
+ """
+
+ # Generate mask
+ mask_sep_atoms = edge_index[0] < edge_index[1]
+ # Distinguish edges between the same (periodic) atom by ordering the cells
+ cell_earlier = (
+ (cell_offsets[:, 0] < 0)
+ | ((cell_offsets[:, 0] == 0) & (cell_offsets[:, 1] < 0))
+ | ((cell_offsets[:, 0] == 0) & (cell_offsets[:, 1] == 0) & (cell_offsets[:, 2] < 0))
+ )
+ mask_same_atoms = edge_index[0] == edge_index[1]
+ mask_same_atoms &= cell_earlier
+ mask = mask_sep_atoms | mask_same_atoms
+
+ # Mask out counter-edges
+ edge_index_new = edge_index[mask[None, :].expand(2, -1)].view(2, -1)
+
+ # Concatenate counter-edges after normal edges
+ edge_index_cat = torch.cat(
+ [
+ edge_index_new,
+ torch.stack([edge_index_new[1], edge_index_new[0]], dim=0),
+ ],
+ dim=1,
+ )
+
+ # Count remaining edges per image
+ batch_edge = torch.repeat_interleave(
+ torch.arange(neighbors.size(0), device=edge_index.device),
+ neighbors,
+ )
+ batch_edge = batch_edge[mask]
+ neighbors_new = 2 * torch.bincount(batch_edge, minlength=neighbors.size(0))
+
+ # Create indexing array
+ edge_reorder_idx = repeat_blocks(
+ neighbors_new // 2,
+ repeats=2,
+ continuous_indexing=True,
+ repeat_inc=edge_index_new.size(1),
+ )
+
+ # Reorder everything so the edges of every image are consecutive
+ edge_index_new = edge_index_cat[:, edge_reorder_idx]
+ cell_offsets_new = self.select_symmetric_edges(cell_offsets, mask, edge_reorder_idx, True)
+ edge_dist_new = self.select_symmetric_edges(edge_dist, mask, edge_reorder_idx, False)
+ edge_vector_new = self.select_symmetric_edges(edge_vector, mask, edge_reorder_idx, True)
+
+ return (
+ edge_index_new,
+ cell_offsets_new,
+ neighbors_new,
+ edge_dist_new,
+ edge_vector_new,
+ )
+
+ def select_edges(
+ self,
+ edge_index: torch.Tensor,
+ cell_offsets: torch.Tensor,
+ neighbors: torch.Tensor,
+ edge_dist: torch.Tensor,
+ edge_vector: torch.Tensor,
+ cutoff: Optional[float] = None,
+ ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+ if cutoff is not None:
+ edge_mask = edge_dist <= cutoff
+
+ edge_index = edge_index[:, edge_mask]
+ cell_offsets = cell_offsets[edge_mask]
+ neighbors = mask_neighbors(neighbors, edge_mask)
+ edge_dist = edge_dist[edge_mask]
+ edge_vector = edge_vector[edge_mask]
+
+ return edge_index, cell_offsets, neighbors, edge_dist, edge_vector
+
+ def generate_interaction_graph(
+ self,
+ cart_coords: torch.Tensor,
+ lattice: torch.Tensor,
+ num_atoms: torch.Tensor,
+ edge_index: torch.Tensor,
+ to_jimages: torch.Tensor,
+ num_bonds: torch.Tensor,
+ ) -> Tuple[
+ Tuple[torch.Tensor, torch.Tensor],
+ torch.Tensor,
+ torch.Tensor,
+ torch.Tensor,
+ torch.Tensor,
+ torch.Tensor,
+ torch.Tensor,
+ torch.Tensor,
+ torch.Tensor,
+ ]:
+ if self.otf_graph:
+ edge_index, to_jimages, num_bonds = radius_graph_pbc(
+ cart_coords=cart_coords,
+ lattice=lattice,
+ num_atoms=num_atoms,
+ radius=self.cutoff,
+ max_num_neighbors_threshold=self.max_neighbors,
+ max_cell_images_per_dim=self.max_cell_images_per_dim,
+ )
+
+ # Switch the indices, so the second one becomes the target index,
+ # over which we can efficiently aggregate.
+ out = get_pbc_distances(
+ cart_coords,
+ edge_index,
+ lattice,
+ to_jimages,
+ num_atoms,
+ num_bonds,
+ coord_is_cart=True,
+ return_offsets=True,
+ return_distance_vec=True,
+ )
+
+ edge_index = out["edge_index"]
+ D_st = out["distances"]
+ # These vectors actually point in the opposite direction.
+ # But we want to use col as idx_t for efficient aggregation.
+ V_st = -out["distance_vec"] / D_st[:, None]
+
+ (
+ edge_index,
+ cell_offsets,
+ neighbors,
+ D_st,
+ V_st,
+ ) = self.reorder_symmetric_edges(edge_index, to_jimages, num_bonds, D_st, V_st)
+
+ # Indices for swapping c->a and a->c (for symmetric MP)
+ block_sizes = neighbors // 2
+
+ # Remove 0 sizes
+ block_sizes = torch.masked_select(block_sizes, block_sizes > 0)
+ id_swap = repeat_blocks(
+ block_sizes,
+ repeats=2,
+ continuous_indexing=False,
+ start_idx=block_sizes[0],
+ block_inc=block_sizes[:-1] + block_sizes[1:],
+ repeat_inc=-block_sizes,
+ )
+
+ id3_ba, id3_ca, id3_ragged_idx = self.get_triplets(
+ edge_index,
+ num_atoms=num_atoms.sum(),
+ )
+
+ return (
+ edge_index,
+ neighbors,
+ D_st,
+ V_st,
+ id_swap,
+ id3_ba,
+ id3_ca,
+ id3_ragged_idx,
+ cell_offsets,
+ )
+
+ def forward(
+ self,
+ z: torch.Tensor,
+ frac_coords: torch.Tensor,
+ atom_types: torch.Tensor,
+ num_atoms: torch.Tensor,
+ batch: torch.Tensor,
+ lengths: Optional[torch.Tensor] = None,
+ angles: Optional[torch.Tensor] = None,
+ edge_index: Optional[torch.Tensor] = None,
+ to_jimages: Optional[torch.Tensor] = None,
+ num_bonds: Optional[torch.Tensor] = None,
+ lattice: Optional[torch.Tensor] = None,
+ ) -> ModelOutput:
+ """
+ args:
+ z: (N_cryst, num_latent)
+ frac_coords: (N_atoms, 3)
+ atom_types: (N_atoms, ) with D3PM need to use atomic number
+ num_atoms: (N_cryst,)
+ lengths: (N_cryst, 3) (optional, either lengths and angles or lattice must be passed)
+ angles: (N_cryst, 3) (optional, either lengths and angles or lattice must be passed)
+ edge_index: (2, N_edge) (optional, only needed if self.otf_graph is False)
+ to_jimages: (N_edge, 3) (optional, only needed if self.otf_graph is False)
+ num_bonds: (N_cryst,) (optional, only needed if self.otf_graph is False)
+ lattice: (N_cryst, 3, 3) (optional, either lengths and angles or lattice must be passed)
+ returns:
+ atom_frac_coords: (N_atoms, 3)
+ atom_types: (N_atoms, MAX_ATOMIC_NUM)
+ """
+
+ if self.otf_graph:
+ assert all(
+ [edge_index is None, to_jimages is None, num_bonds is None]
+ ), "OTF graph construction is active but received input graph information."
+ else:
+ assert not any(
+ [edge_index is None, to_jimages is None, num_bonds is None]
+ ), "OTF graph construction is off but received no input graph information."
+
+ assert (angles is None and lengths is None) != (
+ lattice is None
+ ), "Either lattice or lengths and angles must be provided, not both or none."
+ if angles is not None and lengths is not None:
+ lattice = lattice_params_to_matrix_torch(lengths, angles)
+ assert lattice is not None
+ distorted_lattice = lattice
+
+ pos = frac_to_cart_coords_with_lattice(frac_coords, num_atoms, lattice=distorted_lattice)
+
+ atomic_numbers = atom_types
+
+ (
+ edge_index,
+ neighbors,
+ D_st,
+ V_st,
+ id_swap,
+ id3_ba,
+ id3_ca,
+ id3_ragged_idx,
+ to_jimages,
+ ) = self.generate_interaction_graph(
+ pos, distorted_lattice, num_atoms, edge_index, to_jimages, num_bonds
+ )
+ idx_s, idx_t = edge_index
+
+ # Calculate triplet angles
+ cosφ_cab = inner_product_normalized(V_st[id3_ca], V_st[id3_ba])
+ rad_cbf3, cbf3 = self.cbf_basis3(D_st, cosφ_cab, id3_ca)
+
+ rbf = self.radial_basis(D_st)
+
+ # Embedding block
+ h = self.atom_emb(atomic_numbers)
+ # Merge z and atom embedding
+ if z is not None:
+ z_per_atom = z[batch]
+ h = torch.cat([h, z_per_atom], dim=1)
+ # Combine all embeddings
+ h = self.atom_latent_emb(h)
+ # (nAtoms, emb_size_atom)
+ m = self.edge_emb(h, rbf, idx_s, idx_t) # (nEdges, emb_size_edge)
+ batch_edge = batch[edge_index[0]]
+ cosines = torch.cosine_similarity(V_st[:, None], distorted_lattice[batch_edge], dim=-1)
+ m = torch.cat([m, cosines], dim=-1)
+ m = self.angle_edge_emb(m)
+
+ rbf3 = self.mlp_rbf3(rbf)
+ cbf3 = self.mlp_cbf3(rad_cbf3, cbf3, id3_ca, id3_ragged_idx)
+
+ rbf_h = self.mlp_rbf_h(rbf)
+ rbf_out = self.mlp_rbf_out(rbf)
+
+ E_t, F_st = self.out_blocks[0](h, m, rbf_out, idx_t)
+
+ distance_vec = V_st * D_st[:, None]
+
+ lattice_update = None
+ rbf_lattice = self.mlp_rbf_lattice(rbf)
+ lattice_update = self.lattice_out_blocks[0](
+ edge_emb=m,
+ edge_index=edge_index,
+ distance_vec=distance_vec,
+ lattice=distorted_lattice,
+ batch=batch,
+ rbf=rbf_lattice,
+ normalize_score=True,
+ )
+ F_fully_connected = torch.tensor(0.0, device=distorted_lattice.device)
+ for i in range(self.num_blocks):
+ # Interaction block
+ h, m = self.int_blocks[i](
+ h=h,
+ m=m,
+ rbf3=rbf3,
+ cbf3=cbf3,
+ id3_ragged_idx=id3_ragged_idx,
+ id_swap=id_swap,
+ id3_ba=id3_ba,
+ id3_ca=id3_ca,
+ rbf_h=rbf_h,
+ idx_s=idx_s,
+ idx_t=idx_t,
+ ) # (nAtoms, emb_size_atom), (nEdges, emb_size_edge)
+
+ E, F = self.out_blocks[i + 1](h, m, rbf_out, idx_t)
+ # (nAtoms, num_targets), (nEdges, num_targets)
+ F_st += F
+ E_t += E
+ rbf_lattice = self.mlp_rbf_lattice(rbf)
+ lattice_update += self.lattice_out_blocks[i + 1](
+ edge_emb=m,
+ edge_index=edge_index,
+ distance_vec=distance_vec,
+ lattice=distorted_lattice,
+ batch=batch,
+ rbf=rbf_lattice,
+ normalize_score=True,
+ )
+
+ nMolecules = torch.max(batch) + 1
+
+ if self.encoder_mode:
+ return E_t
+ # always use sum aggregation
+ E_t = scatter(
+ E_t, batch, dim=0, dim_size=nMolecules, reduce="sum"
+ ) # (nMolecules, num_targets)
+
+ # always output energy, forces and node embeddings
+ output = dict(energy=E_t, node_embeddings=h)
+
+ # map forces in edge directions
+ F_st_vec = F_st[:, :, None] * V_st[:, None, :]
+ # (nEdges, num_targets, 3)
+ F_t = scatter(
+ F_st_vec,
+ idx_t,
+ dim=0,
+ dim_size=num_atoms.sum(),
+ reduce="add",
+ ) # (nAtoms, num_targets, 3)
+ F_t = F_t.squeeze(1) # (nAtoms, 3)
+ output["forces"] = F_t + F_fully_connected
+
+ if self.regress_stress:
+ # optionally get predicted stress tensor
+ # shape=(Nbatch, 3, 3)
+ output["stress"] = lattice_update
+
+ return ModelOutput(**output)
+
+ @property
+ def num_params(self):
+ return sum(p.numel() for p in self.parameters())
diff --git a/model/common/gemnet/gemnet_ctrl.py b/model/common/gemnet/gemnet_ctrl.py
new file mode 100644
index 0000000000000000000000000000000000000000..5e6a669e43fd4a32d8b333010bff309f0ee11328
--- /dev/null
+++ b/model/common/gemnet/gemnet_ctrl.py
@@ -0,0 +1,268 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+# Adapted from https://github.com/FAIR-Chem/fairchem/blob/main/src/fairchem/core/models/gemnet/gemnet.py.
+
+from typing import Dict, List, Optional
+
+# import numpy as np
+import torch
+import torch.nn as nn
+from torch_scatter import scatter
+
+from onescience.datapipes.materials.mattergen.types import PropertySourceId
+from ...common.gemnet.gemnet import GemNetT, ModelOutput
+from ...common.gemnet.utils import inner_product_normalized
+from ...common.utils.data_utils import (
+ frac_to_cart_coords_with_lattice,
+ lattice_params_to_matrix_torch,
+)
+
+
+class GemNetTCtrl(GemNetT):
+ """
+ GemNet-T, triplets-only variant of GemNet
+
+ This variation allows for layerwise conditional control for the purpose of
+ conditional finetuning. It adds the following on top of GemNetT:
+
+ for each condition in :
+
+ 1. a series of adapt layers that take the concatenation of the node embedding
+ and the condition embedding, process it with an MLP. There is one adapt layer
+ for each GemNetT message passing block.
+ 2. a series of mixin layers that take the output of the adapt layer and mix it in
+ to the atom embedding. There is one mixin layer for each GemNetT message passing block.
+ The mixin layers are initialized to zeros so at the beginning of training, the model
+ outputs exactly the same scores as the base GemNetT model.
+
+ """
+
+ def __init__(self, condition_on_adapt: List[PropertySourceId], *args, **kwargs):
+ super().__init__(*args, **kwargs)
+
+ self.condition_on_adapt = condition_on_adapt
+ self.cond_adapt_layers = nn.ModuleDict()
+ self.cond_mixin_layers = nn.ModuleDict()
+ # default value for emb_size_atom is 512
+ self.emb_size_atom = kwargs["emb_size_atom"] if "emb_size_atom" in kwargs else 512
+
+ for cond in condition_on_adapt:
+ adapt_layers = []
+ mixin_layers = []
+
+ for _ in range(self.num_blocks):
+ adapt_layers.append(
+ nn.Sequential(
+ nn.Linear(self.emb_size_atom * 2, self.emb_size_atom),
+ nn.ReLU(),
+ nn.Linear(self.emb_size_atom, self.emb_size_atom),
+ )
+ )
+ mixin_layers.append(nn.Linear(self.emb_size_atom, self.emb_size_atom, bias=False))
+ nn.init.zeros_(mixin_layers[-1].weight)
+
+ self.cond_adapt_layers[cond] = torch.nn.ModuleList(adapt_layers)
+ self.cond_mixin_layers[cond] = torch.nn.ModuleList(mixin_layers)
+
+ def forward(
+ self,
+ z: torch.Tensor,
+ frac_coords: torch.Tensor,
+ atom_types: torch.Tensor,
+ num_atoms: torch.Tensor,
+ batch: torch.Tensor,
+ lengths: Optional[torch.Tensor] = None,
+ angles: Optional[torch.Tensor] = None,
+ edge_index: Optional[torch.Tensor] = None,
+ to_jimages: Optional[torch.Tensor] = None,
+ num_bonds: Optional[torch.Tensor] = None,
+ lattice: Optional[torch.Tensor] = None,
+ charges: Optional[torch.Tensor] = None,
+ cond_adapt: Optional[Dict[PropertySourceId, torch.Tensor]] = None,
+ cond_adapt_mask: Optional[Dict[PropertySourceId, torch.Tensor]] = None,
+ ) -> ModelOutput:
+ """
+ args:
+ z: (N_cryst, num_latent)
+ frac_coords: (N_atoms, 3)
+ atom_types: (N_atoms, ) with D3PM need to use atomic number
+ num_atoms: (N_cryst,)
+ lengths: (N_cryst, 3) (optional, either lengths and angles or lattice must be passed)
+ angles: (N_cryst, 3) (optional, either lengths and angles or lattice must be passed)
+ edge_index: (2, N_edge) (optional, only needed if self.otf_graph is False)
+ to_jimages: (N_edge, 3) (optional, only needed if self.otf_graph is False)
+ num_bonds: (N_cryst,) (optional, only needed if self.otf_graph is False)
+ lattice: (N_cryst, 3, 3) (optional, either lengths and angles or lattice must be passed)
+ cond_adapt: (N_cryst, num_cond, dim_cond) (optional, conditional signal for score prediction)
+ cond_adapt_mask: (N_cryst, num_cond) (optional, mask for which data points receive conditional signal)
+ returns:
+ atom_frac_coords: (N_atoms, 3)
+ atom_types: (N_atoms, MAX_ATOMIC_NUM)
+ """
+
+ if self.otf_graph:
+ assert all(
+ [edge_index is None, to_jimages is None, num_bonds is None]
+ ), "OTF graph construction is active but received input graph information."
+ else:
+ assert not any(
+ [edge_index is None, to_jimages is None, num_bonds is None]
+ ), "OTF graph construction is off but received no input graph information."
+
+ assert (angles is None and lengths is None) != (
+ lattice is None
+ ), "Either lattice or lengths and angles must be provided, not both or none."
+ if angles is not None and lengths is not None:
+ lattice = lattice_params_to_matrix_torch(lengths, angles)
+ assert lattice is not None
+ distorted_lattice = lattice
+
+ pos = frac_to_cart_coords_with_lattice(frac_coords, num_atoms, lattice=distorted_lattice)
+
+ atomic_numbers = atom_types
+
+ (
+ edge_index,
+ neighbors,
+ D_st,
+ V_st,
+ id_swap,
+ id3_ba,
+ id3_ca,
+ id3_ragged_idx,
+ to_jimages,
+ ) = self.generate_interaction_graph(
+ pos, distorted_lattice, num_atoms, edge_index, to_jimages, num_bonds
+ )
+ idx_s, idx_t = edge_index
+
+ # Calculate triplet angles
+ cosφ_cab = inner_product_normalized(V_st[id3_ca], V_st[id3_ba])
+ rad_cbf3, cbf3 = self.cbf_basis3(D_st, cosφ_cab, id3_ca)
+
+ rbf = self.radial_basis(D_st)
+
+ # Embedding block
+ h = self.atom_emb(atomic_numbers)
+ # Merge z and atom embedding
+ if z is not None:
+ z_per_atom = z[batch]
+ h = torch.cat([h, z_per_atom], dim=1)
+ h = self.atom_latent_emb(h)
+ # (nAtoms, emb_size_atom)
+ m = self.edge_emb(h, rbf, idx_s, idx_t) # (nEdges, emb_size_edge)
+ batch_edge = batch[edge_index[0]]
+ cosines = torch.cosine_similarity(V_st[:, None], distorted_lattice[batch_edge], dim=-1)
+ m = torch.cat([m, cosines], dim=-1)
+ m = self.angle_edge_emb(m)
+
+ rbf3 = self.mlp_rbf3(rbf)
+ cbf3 = self.mlp_cbf3(rad_cbf3, cbf3, id3_ca, id3_ragged_idx)
+
+ rbf_h = self.mlp_rbf_h(rbf)
+ rbf_out = self.mlp_rbf_out(rbf)
+
+ E_t, F_st = self.out_blocks[0](h, m, rbf_out, idx_t)
+
+ distance_vec = V_st * D_st[:, None]
+
+ lattice_update = None
+ rbf_lattice = self.mlp_rbf_lattice(rbf)
+ lattice_update = self.lattice_out_blocks[0](
+ edge_emb=m,
+ edge_index=edge_index,
+ distance_vec=distance_vec,
+ lattice=distorted_lattice,
+ batch=batch,
+ rbf=rbf_lattice,
+ normalize_score=True,
+ )
+
+ # currently only working for a single cond adapt property.
+ # to extend to multi-properties,
+ # use a ModuleDict for adapt layers and mixin layers.
+ # use a dictionary to track the conditions?
+
+ if cond_adapt is not None and cond_adapt_mask is not None:
+ cond_adapt_per_atom = {}
+ cond_adapt_mask_per_atom = {}
+ for cond in self.condition_on_adapt:
+ cond_adapt_per_atom[cond] = cond_adapt[cond][batch]
+ # 1 = use conditional embedding, 0 = use unconditional embedding
+ cond_adapt_mask_per_atom[cond] = 1.0 - cond_adapt_mask[cond][batch].float()
+
+ for i in range(self.num_blocks):
+ h_adapt = torch.zeros_like(h)
+ for cond in self.condition_on_adapt:
+ h_adapt_cond = self.cond_adapt_layers[cond][i](
+ torch.cat([h, cond_adapt_per_atom[cond]], dim=-1)
+ )
+ h_adapt_cond = self.cond_mixin_layers[cond][i](h_adapt_cond)
+ # cond_adapt_mask_per_atom[cond] is 1.0 if we want to use conditional embedding and 0 for unconditional embedding
+ h_adapt += cond_adapt_mask_per_atom[cond] * h_adapt_cond
+ h = h + h_adapt
+
+ # Interaction block
+ h, m = self.int_blocks[i](
+ h=h,
+ m=m,
+ rbf3=rbf3,
+ cbf3=cbf3,
+ id3_ragged_idx=id3_ragged_idx,
+ id_swap=id_swap,
+ id3_ba=id3_ba,
+ id3_ca=id3_ca,
+ rbf_h=rbf_h,
+ idx_s=idx_s,
+ idx_t=idx_t,
+ ) # (nAtoms, emb_size_atom), (nEdges, emb_size_edge)
+
+ E, F = self.out_blocks[i + 1](h, m, rbf_out, idx_t)
+ # (nAtoms, num_targets), (nEdges, num_targets)
+ F_st += F
+ E_t += E
+ rbf_lattice = self.mlp_rbf_lattice(rbf)
+ lattice_update += self.lattice_out_blocks[i + 1](
+ edge_emb=m,
+ edge_index=edge_index,
+ distance_vec=distance_vec,
+ lattice=distorted_lattice,
+ batch=batch,
+ rbf=rbf_lattice,
+ normalize_score=True,
+ )
+
+ nMolecules = torch.max(batch) + 1
+
+ # always use sum aggregation
+ E_t = scatter(
+ E_t, batch, dim=0, dim_size=nMolecules, reduce="sum"
+ ) # (nMolecules, num_targets)
+
+ # always output energy, forces and node embeddings
+ output = dict(energy=E_t, node_embeddings=h)
+
+ # map forces in edge directions
+ F_st_vec = F_st[:, :, None] * V_st[:, None, :]
+ # (nEdges, num_targets, 3)
+ F_t = scatter(
+ F_st_vec,
+ idx_t,
+ dim=0,
+ dim_size=num_atoms.sum(),
+ reduce="add",
+ ) # (nAtoms, num_targets, 3)
+ F_t = F_t.squeeze(1) # (nAtoms, 3)
+ output["forces"] = F_t
+
+ if self.regress_stress:
+ # shape=(Nbatch, 3, 3)
+ output["stress"] = lattice_update
+
+ return ModelOutput(**output)
+
+ @property
+ def num_params(self):
+ return sum(p.numel() for p in self.parameters())
diff --git a/model/common/gemnet/initializers.py b/model/common/gemnet/initializers.py
new file mode 100644
index 0000000000000000000000000000000000000000..681c29d9c724a17255179cd96187398d4edbfba7
--- /dev/null
+++ b/model/common/gemnet/initializers.py
@@ -0,0 +1,49 @@
+"""
+Copyright (c) Facebook, Inc. and its affiliates.
+Copyright (c) Microsoft Corporation.
+Licensed under the MIT License.
+Adapted from https://github.com/FAIR-Chem/fairchem/blob/main/src/fairchem/core/models/gemnet/initializers.py.
+"""
+
+import torch
+
+
+# This function is not type annotated because mypy complains that axis could be either an integer or a tuple of integers,
+# even though this is precicely how torch.var_mean works
+def _standardize(kernel):
+ """
+ Makes sure that N*Var(W) = 1 and E[W] = 0
+ """
+ eps = 1e-6
+
+ if len(kernel.shape) == 3:
+ axis = (0, 1) # last dimension is output dimension
+ else:
+ axis = 1
+
+ var, mean = torch.var_mean(kernel, dim=axis, unbiased=True, keepdim=True)
+ kernel = (kernel - mean) / (var + eps) ** 0.5
+ return kernel
+
+
+def he_orthogonal_init(tensor: torch.Tensor) -> torch.Tensor:
+ """
+ Generate a weight matrix with variance according to He (Kaiming) initialization.
+ Based on a random (semi-)orthogonal matrix neural networks
+ are expected to learn better when features are decorrelated
+ (stated by eg. "Reducing overfitting in deep networks by decorrelating representations",
+ "Dropout: a simple way to prevent neural networks from overfitting",
+ "Exact solutions to the nonlinear dynamics of learning in deep linear neural networks")
+ """
+ tensor = torch.nn.init.orthogonal_(tensor)
+
+ if len(tensor.shape) == 3:
+ fan_in = tensor.shape[:-1].numel()
+ else:
+ fan_in = tensor.shape[1]
+
+ with torch.no_grad():
+ tensor.data = _standardize(tensor.data)
+ tensor.data *= (1 / fan_in) ** 0.5
+
+ return tensor
diff --git a/model/common/gemnet/utils.py b/model/common/gemnet/utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..b675da529a31f60095a774845aad265f9c9e2f31
--- /dev/null
+++ b/model/common/gemnet/utils.py
@@ -0,0 +1,299 @@
+"""
+Copyright (c) Facebook, Inc. and its affiliates.
+Copyright (c) Microsoft Corporation.
+Licensed under the MIT License.
+Adapted from https://github.com/FAIR-Chem/fairchem/blob/main/src/fairchem/core/models/gemnet/utils.py.
+"""
+
+import json
+from typing import Any, Dict, Optional, Tuple
+
+import torch
+from torch_scatter import segment_csr
+
+
+def read_json(path: str) -> Dict:
+ """"""
+ if not path.endswith(".json"):
+ raise UserWarning(f"Path {path} is not a json-path.")
+
+ with open(path, "r") as f:
+ content = json.load(f)
+ return content
+
+
+def update_json(path: str, data: Dict):
+ """"""
+ if not path.endswith(".json"):
+ raise UserWarning(f"Path {path} is not a json-path.")
+
+ content = read_json(path)
+ content.update(data)
+ write_json(path, content)
+
+
+def write_json(path: str, data: Dict):
+ """"""
+ if not path.endswith(".json"):
+ raise UserWarning(f"Path {path} is not a json-path.")
+
+ with open(path, "w", encoding="utf-8") as f:
+ json.dump(data, f, ensure_ascii=False, indent=4)
+
+
+def read_value_json(path: str, key: str) -> Optional[Any]:
+ """"""
+ content = read_json(path)
+
+ if key in content.keys():
+ return content[key]
+ else:
+ return None
+
+
+def ragged_range(sizes: torch.Tensor) -> torch.Tensor:
+ """Multiple concatenated ranges.
+
+ Examples
+ --------
+ sizes = [1 4 2 3]
+ Return: [0 0 1 2 3 0 1 0 1 2]
+ """
+ assert sizes.dim() == 1
+ if sizes.sum() == 0:
+ return sizes.new_empty(0)
+
+ # Remove 0 sizes
+ sizes_nonzero = sizes > 0
+ if not torch.all(sizes_nonzero):
+ sizes = torch.masked_select(sizes, sizes_nonzero)
+
+ # Initialize indexing array with ones as we need to setup incremental indexing
+ # within each group when cumulatively summed at the final stage.
+ id_steps = torch.ones(sizes.sum(), dtype=torch.long, device=sizes.device)
+ id_steps[0] = 0
+ insert_index = sizes[:-1].cumsum(0)
+ insert_val = (1 - sizes)[:-1]
+
+ # Assign index-offsetting values
+ id_steps[insert_index] = insert_val
+
+ # Finally index into input array for the group repeated o/p
+ res = id_steps.cumsum(0)
+ return res
+
+
+def repeat_blocks(
+ sizes: torch.Tensor,
+ repeats: torch.Tensor,
+ continuous_indexing: bool = True,
+ start_idx: int = 0,
+ block_inc: int = 0,
+ repeat_inc: int = 0,
+) -> torch.Tensor:
+ """Repeat blocks of indices.
+ Adapted from https://stackoverflow.com/questions/51154989/numpy-vectorized-function-to-repeat-blocks-of-consecutive-elements
+
+ continuous_indexing: Whether to keep increasing the index after each block
+ start_idx: Starting index
+ block_inc: Number to increment by after each block,
+ either global or per block. Shape: len(sizes) - 1
+ repeat_inc: Number to increment by after each repetition,
+ either global or per block
+
+ Examples
+ --------
+ sizes = [1,3,2] ; repeats = [3,2,3] ; continuous_indexing = False
+ Return: [0 0 0 0 1 2 0 1 2 0 1 0 1 0 1]
+ sizes = [1,3,2] ; repeats = [3,2,3] ; continuous_indexing = True
+ Return: [0 0 0 1 2 3 1 2 3 4 5 4 5 4 5]
+ sizes = [1,3,2] ; repeats = [3,2,3] ; continuous_indexing = True ;
+ repeat_inc = 4
+ Return: [0 4 8 1 2 3 5 6 7 4 5 8 9 12 13]
+ sizes = [1,3,2] ; repeats = [3,2,3] ; continuous_indexing = True ;
+ start_idx = 5
+ Return: [5 5 5 6 7 8 6 7 8 9 10 9 10 9 10]
+ sizes = [1,3,2] ; repeats = [3,2,3] ; continuous_indexing = True ;
+ block_inc = 1
+ Return: [0 0 0 2 3 4 2 3 4 6 7 6 7 6 7]
+ sizes = [0,3,2] ; repeats = [3,2,3] ; continuous_indexing = True
+ Return: [0 1 2 0 1 2 3 4 3 4 3 4]
+ sizes = [2,3,2] ; repeats = [2,0,2] ; continuous_indexing = True
+ Return: [0 1 0 1 5 6 5 6]
+ """
+ assert sizes.dim() == 1
+ assert all(sizes >= 0)
+
+ # Remove 0 sizes
+ sizes_nonzero = sizes > 0
+ if not torch.all(sizes_nonzero):
+ assert block_inc == 0 # Implementing this is not worth the effort
+ sizes = torch.masked_select(sizes, sizes_nonzero)
+ if isinstance(repeats, torch.Tensor):
+ repeats = torch.masked_select(repeats, sizes_nonzero)
+ if isinstance(repeat_inc, torch.Tensor):
+ repeat_inc = torch.masked_select(repeat_inc, sizes_nonzero)
+
+ if isinstance(repeats, torch.Tensor):
+ assert all(repeats >= 0)
+ insert_dummy = repeats[0] == 0
+ if insert_dummy:
+ one = sizes.new_ones(1)
+ zero = sizes.new_zeros(1)
+ sizes = torch.cat((one, sizes))
+ repeats = torch.cat((one, repeats))
+ if isinstance(block_inc, torch.Tensor):
+ block_inc = torch.cat((zero, block_inc))
+ if isinstance(repeat_inc, torch.Tensor):
+ repeat_inc = torch.cat((zero, repeat_inc))
+ else:
+ assert repeats >= 0
+ insert_dummy = False
+
+ # Get repeats for each group using group lengths/sizes
+ r1 = torch.repeat_interleave(torch.arange(len(sizes), device=sizes.device), repeats)
+
+ # Get total size of output array, as needed to initialize output indexing array
+ N = (sizes * repeats).sum()
+
+ # Initialize indexing array with ones as we need to setup incremental indexing
+ # within each group when cumulatively summed at the final stage.
+ # Two steps here:
+ # 1. Within each group, we have multiple sequences, so setup the offsetting
+ # at each sequence lengths by the seq. lengths preceding those.
+ id_ar = torch.ones(N, dtype=torch.long, device=sizes.device)
+ id_ar[0] = 0
+ insert_index = sizes[r1[:-1]].cumsum(0)
+ insert_val = (1 - sizes)[r1[:-1]]
+
+ if isinstance(repeats, torch.Tensor) and torch.any(repeats == 0):
+ diffs = r1[1:] - r1[:-1]
+ indptr = torch.cat((sizes.new_zeros(1), diffs.cumsum(0)))
+ if continuous_indexing:
+ # If a group was skipped (repeats=0) we need to add its size
+ insert_val += segment_csr(sizes[: r1[-1]], indptr, reduce="sum")
+
+ # Add block increments
+ if isinstance(block_inc, torch.Tensor):
+ insert_val += segment_csr(block_inc[: r1[-1]], indptr, reduce="sum")
+ else:
+ insert_val += block_inc * (indptr[1:] - indptr[:-1])
+ if insert_dummy:
+ insert_val[0] -= block_inc
+ else:
+ idx = r1[1:] != r1[:-1]
+ if continuous_indexing:
+ # 2. For each group, make sure the indexing starts from the next group's
+ # first element. So, simply assign 1s there.
+ insert_val[idx] = 1
+
+ # Add block increments
+ insert_val[idx] += block_inc
+
+ # Add repeat_inc within each group
+ if isinstance(repeat_inc, torch.Tensor):
+ insert_val += repeat_inc[r1[:-1]]
+ if isinstance(repeats, torch.Tensor):
+ repeat_inc_inner = repeat_inc[repeats > 0][:-1]
+ else:
+ repeat_inc_inner = repeat_inc[:-1]
+ else:
+ insert_val += repeat_inc
+ repeat_inc_inner = repeat_inc
+
+ # Subtract the increments between groups
+ if isinstance(repeats, torch.Tensor):
+ repeats_inner = repeats[repeats > 0][:-1]
+ else:
+ repeats_inner = repeats
+ insert_val[r1[1:] != r1[:-1]] -= repeat_inc_inner * repeats_inner
+
+ # Assign index-offsetting values
+ id_ar[insert_index] = insert_val
+
+ if insert_dummy:
+ id_ar = id_ar[1:]
+ if continuous_indexing:
+ id_ar[0] -= 1
+
+ # Set start index now, in case of insertion due to leading repeats=0
+ id_ar[0] += start_idx
+
+ # Finally index into input array for the group repeated o/p
+ res = id_ar.cumsum(0)
+ return res
+
+
+def calculate_interatomic_vectors(
+ R: torch.Tensor, id_s: torch.Tensor, id_t: torch.Tensor, offsets_st: torch.Tensor
+) -> Tuple[torch.Tensor, torch.Tensor]:
+ """
+ Calculate the vectors connecting the given atom pairs,
+ considering offsets from periodic boundary conditions (PBC).
+
+ Parameters
+ ----------
+ R: Tensor, shape = (nAtoms, 3)
+ Atom positions.
+ id_s: Tensor, shape = (nEdges,)
+ Indices of the source atom of the edges.
+ id_t: Tensor, shape = (nEdges,)
+ Indices of the target atom of the edges.
+ offsets_st: Tensor, shape = (nEdges,)
+ PBC offsets of the edges.
+ Subtract this from the correct direction.
+
+ Returns
+ -------
+ (D_st, V_st): tuple
+ D_st: Tensor, shape = (nEdges,)
+ Distance from atom t to s.
+ V_st: Tensor, shape = (nEdges,)
+ Unit direction from atom t to s.
+ """
+ Rs = R[id_s]
+ Rt = R[id_t]
+ # ReLU prevents negative numbers in sqrt
+ if offsets_st is None:
+ V_st = Rt - Rs # s -> t
+ else:
+ V_st = Rt - Rs + offsets_st # s -> t
+ D_st = torch.sqrt(torch.sum(V_st**2, dim=1))
+ V_st = V_st / D_st[..., None]
+ return D_st, V_st
+
+
+def inner_product_normalized(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
+ """
+ Calculate the inner product between the given normalized vectors,
+ giving a result between -1 and 1.
+ """
+ return torch.sum(x * y, dim=-1).clamp(min=-1, max=1)
+
+
+def mask_neighbors(neighbors: torch.Tensor, edge_mask: torch.Tensor) -> torch.Tensor:
+ neighbors_old_indptr = torch.cat([neighbors.new_zeros(1), neighbors])
+ neighbors_old_indptr = torch.cumsum(neighbors_old_indptr, dim=0)
+ neighbors = segment_csr(edge_mask.long(), neighbors_old_indptr)
+ return neighbors
+
+
+def get_k_index_product_set(
+ num_k_x: torch.LongTensor, num_k_y: torch.LongTensor, num_k_z: torch.LongTensor
+) -> tuple[torch.FloatTensor, int]:
+ # Get a box of k-lattice indices around (0,0,0)
+ k_index_sets = (
+ torch.arange(-num_k_x, num_k_x + 1, dtype=torch.float),
+ torch.arange(-num_k_y, num_k_y + 1, dtype=torch.float),
+ torch.arange(-num_k_z, num_k_z + 1, dtype=torch.float),
+ )
+ k_index_product_set = torch.cartesian_prod(*k_index_sets)
+ # Because our "signal" is real-valued, for the Fourier transform it holds that
+ # F(omega) = F*(-omega) (where F* is the complex conjugate of F). Thus, we only
+ # need to consider the positive half of the Fourier space.
+ k_index_product_set = k_index_product_set[k_index_product_set.shape[0] // 2 + 1 :]
+
+ # Amount of k-points
+ num_k_degrees_of_freedom = k_index_product_set.shape[0]
+
+ return k_index_product_set, num_k_degrees_of_freedom
diff --git a/model/common/globals.py b/model/common/globals.py
new file mode 100644
index 0000000000000000000000000000000000000000..911e89fddf2d67e5ac33b0173b900e175ff5deed
--- /dev/null
+++ b/model/common/globals.py
@@ -0,0 +1,8 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+from pathlib import Path
+
+PROJECT_ROOT = Path(__file__).resolve().parent.parent
+GENERATED_CRYSTALS_ZIP_FILE_NAME = "generated_crystals_cif.zip"
+GENERATED_CRYSTALS_EXTXYZ_FILE_NAME = "generated_crystals.extxyz"
diff --git a/model/common/loss.py b/model/common/loss.py
new file mode 100644
index 0000000000000000000000000000000000000000..91524fe806e0b0a87d801b1ca09835ed6bf4f4c6
--- /dev/null
+++ b/model/common/loss.py
@@ -0,0 +1,56 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+from functools import partial
+from typing import Dict, Literal, Optional
+
+from ..diffusion.losses import SummedFieldLoss, denoising_score_matching
+from ..diffusion.model_target import ModelTarget
+from ..diffusion.training.field_loss import FieldLoss, d3pm_loss
+from ..diffusion.wrapped.wrapped_normal_loss import wrapped_normal_loss
+
+
+class MaterialsLoss(SummedFieldLoss):
+ def __init__(
+ self,
+ reduce: Literal["sum", "mean"] = "mean",
+ d3pm_hybrid_lambda: float = 0.0,
+ include_pos: bool = True,
+ include_cell: bool = True,
+ include_atomic_numbers: bool = True,
+ weights: Optional[Dict[str, float]] = None,
+ ):
+ model_targets = {"pos": ModelTarget.score_times_std, "cell": ModelTarget.score_times_std}
+ self.fields_to_score = []
+ self.categorical_fields = []
+ loss_fns: Dict[str, FieldLoss] = {}
+ if include_pos:
+ self.fields_to_score.append("pos")
+ loss_fns["pos"] = partial(
+ wrapped_normal_loss,
+ reduce=reduce,
+ model_target=model_targets["pos"],
+ )
+ if include_cell:
+ self.fields_to_score.append("cell")
+ loss_fns["cell"] = partial(
+ denoising_score_matching,
+ reduce=reduce,
+ model_target=model_targets["cell"],
+ )
+ if include_atomic_numbers:
+ model_targets["atomic_numbers"] = ModelTarget.logits
+ self.fields_to_score.append("atomic_numbers")
+ self.categorical_fields.append("atomic_numbers")
+ loss_fns["atomic_numbers"] = partial(
+ d3pm_loss,
+ reduce=reduce,
+ d3pm_hybrid_lambda=d3pm_hybrid_lambda,
+ )
+ self.reduce = reduce
+ self.d3pm_hybrid_lambda = d3pm_hybrid_lambda
+ super().__init__(
+ loss_fns=loss_fns,
+ weights=weights,
+ model_targets=model_targets,
+ )
diff --git a/model/common/utils/__init__.py b/model/common/utils/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/model/common/utils/config_utils.py b/model/common/utils/config_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..b8c20764845fc0643750ba54a0fba4e32975d9cb
--- /dev/null
+++ b/model/common/utils/config_utils.py
@@ -0,0 +1,53 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+import argparse
+import sys
+from typing import Callable, TypeVar, cast
+
+from omegaconf import OmegaConf
+
+R = TypeVar("R")
+
+
+def get_config(argv: list[str] | None, config_cls: Callable[..., R]) -> R:
+ """
+ Utility function to get OmegaConf config options.
+
+ Args:
+ argv: Either a list of command line arguments to parse, or None.
+ If None, this argument is set from sys.argv.
+ config_cls: Dataclass object specifying config structure
+ (i.e. which fields to expect in the config).
+ It should be the class itself, NOT an instance of the class.
+
+ Returns:
+ Config object, which will pass as an instance of `config_cls` among other things.
+ Note: the type for this could be specified more carefully, but OmegaConf's typing
+ system is a bit complex. See OmegaConf's docs for "structured" for more info.
+ """
+
+ if argv is None:
+ argv = sys.argv[1:]
+ # Parse command line arguments
+ parser = argparse.ArgumentParser(allow_abbrev=False) # prevent prefix matching issues
+ parser.add_argument(
+ "--config",
+ type=str,
+ action="append",
+ default=list(),
+ help="Path to a yaml config file. "
+ "Argument can be repeated multiple times, with later configs overwriting previous ones.",
+ )
+ args, config_changes = parser.parse_known_args(argv)
+
+ # Read configs from file and command line
+ conf_yamls = [OmegaConf.load(c) for c in args.config]
+ conf_cli = OmegaConf.from_cli(config_changes)
+
+ # Make merged config options
+ # CLI options take priority over YAML file options
+ schema = OmegaConf.structured(config_cls)
+ config = OmegaConf.merge(schema, *conf_yamls, conf_cli)
+ OmegaConf.set_readonly(config, True) # should not be written to
+ return cast(R, config)
diff --git a/model/common/utils/data_classes.py b/model/common/utils/data_classes.py
new file mode 100644
index 0000000000000000000000000000000000000000..678ccc5f8382454de572646ddfb40c99d56e86a1
--- /dev/null
+++ b/model/common/utils/data_classes.py
@@ -0,0 +1,192 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+import fnmatch
+import os
+from dataclasses import asdict, dataclass, field
+from functools import cached_property
+from pathlib import Path
+from typing import Any, Literal, Protocol
+
+import numpy as np
+from huggingface_hub import hf_hub_download
+from hydra import compose, initialize_config_dir
+from omegaconf import DictConfig
+from omegaconf import ListConfig
+
+PRETRAINED_MODEL_NAME = Literal[
+ "mattergen_base",
+ "chemical_system",
+ "space_group",
+ "dft_mag_density",
+ "dft_band_gap",
+ "ml_bulk_modulus",
+ "dft_mag_density_hhi_score",
+ "chemical_system_energy_above_hull",
+ "mp_20_base",
+]
+
+
+def _rewrite_vendored_targets(value):
+ """Rewrite upstream Hydra targets when loading an external checkpoint config."""
+ if isinstance(value, DictConfig):
+ for key in value:
+ value[key] = _rewrite_vendored_targets(value[key])
+ return value
+ if isinstance(value, ListConfig):
+ for index in range(len(value)):
+ value[index] = _rewrite_vendored_targets(value[index])
+ return value
+ if isinstance(value, str):
+ if value.startswith("onescience.models.mattergen."):
+ return "model." + value[len("onescience.models.mattergen."):]
+ if value.startswith("mattergen."):
+ if value.startswith("mattergen.common.data."):
+ return "onescience.datapipes.materials.mattergen." + value[len("mattergen.common.data."):]
+ if value.startswith("mattergen.common.gemnet.layers."):
+ return "onescience.modules.layer.mattergen." + value[len("mattergen.common.gemnet.layers."):]
+ if value.startswith("mattergen.property_embeddings."):
+ return "onescience.modules.embedding.mattergen_property_embeddings." + value[len("mattergen.property_embeddings."):]
+ return "model." + value[len("mattergen."):]
+ return value
+ return value
+
+
+def find_local_files(local_path: str, glob: str = "*", relative: bool = False) -> list[str]:
+ """
+ Find files in the given directory or blob storage path, and return the list of files
+ matching the given glob pattern. If relative is True, the returned paths are relative
+ to the given directory or blob storage path.
+
+ Args:
+ blob_or_local_path: path to the directory or blob storage path
+ glob: glob pattern to match. By default, all files are returned.
+ relative: whether to return relative paths. By default, absolute paths are returned.
+
+ Returns:
+ list of paths to files matching the given glob pattern.
+ """
+ # list all files here, filtering happens in the `fnmatch.filter` step
+ local_files = [x for x in Path(local_path).rglob("*") if os.path.isfile(x)]
+ files_list = [str(x.relative_to(local_path)) if relative else str(x) for x in local_files]
+ return fnmatch.filter(files_list, glob)
+
+
+@dataclass(frozen=True)
+class MatterGenCheckpointInfo:
+ model_path: str
+ load_epoch: int | Literal["best", "last"] | None = "last"
+ config_overrides: list[str] = field(default_factory=list)
+ split: str = "val"
+ strict_checkpoint_loading: bool = True
+
+ @classmethod
+ def from_hf_hub(
+ cls,
+ model_name: PRETRAINED_MODEL_NAME,
+ repository_name: str = "microsoft/mattergen",
+ config_overrides: list[str] = None,
+ ):
+ """
+ Instantiate a MatterGenCheckpointInfo object from a model hosted on the Hugging Face Hub.
+
+ """
+ hf_hub_download(
+ repo_id=repository_name, filename=f"checkpoints/{model_name}/checkpoints/last.ckpt"
+ )
+ config_path = hf_hub_download(
+ repo_id=repository_name, filename=f"checkpoints/{model_name}/config.yaml"
+ )
+ return cls(
+ model_path=Path(config_path).parent,
+ config_overrides=config_overrides or [],
+ load_epoch="last",
+ )
+
+ def as_dict(self) -> dict[str, Any]:
+ d = asdict(self)
+ d["model_path"] = str(self.model_path) # we cannot put Path object in mongo DB
+ return d
+
+ @classmethod
+ def from_dict(cls, d) -> "MatterGenCheckpointInfo":
+ d = d.copy()
+ d["model_path"] = Path(d["model_path"])
+ # no longer used
+ if "load_data" in d:
+ del d["load_data"]
+ return cls(**d)
+
+ @property
+ def config(self) -> DictConfig:
+ with initialize_config_dir(str(self.model_path)):
+ cfg = compose(config_name="config", overrides=self.config_overrides)
+ return _rewrite_vendored_targets(cfg)
+
+ @cached_property
+ def checkpoint_path(self) -> str:
+ """
+ Search for checkpoint files in the given directory, and return the path
+ to the checkpoint with the given epoch number or the best checkpoint if load_epoch is "best".
+ "Best" is selected via the lowest validation loss, which is stored in the checkpoint filename.
+ Assumes that the checkpoint filenames are of the form "epoch=1-val_loss=0.1234.ckpt" or 'last.ckpt'.
+
+ Returns:
+ Path to the checkpoint file to load.
+ """
+ # look for checkpoints recursively in the given directory or blob storage path.
+ # I.e., if the path is '/path/', we will find .ckpt files in '/path/version_0/checkpoints'
+ # and '/path/version_1/checkpoints', and so on.
+ model_path = str(self.model_path)
+ ckpts = find_local_files(local_path=model_path, glob="*.ckpt")
+ assert len(ckpts) > 0, f"No checkpoints found at {model_path}"
+ if self.load_epoch == "last":
+ assert any(
+ [x.endswith("last.ckpt") for x in ckpts]
+ ), "No last.ckpt found in checkpoints."
+ return [x for x in ckpts if x.endswith("last.ckpt")][0]
+ # Drop last.ckpt to exclude it from the epoch selection
+ ckpts = [x for x in ckpts if not x.endswith("last.ckpt")]
+
+ # Convert strings to Path to be able to use the .parts attribute
+ ckpt_paths = [Path(x) for x in ckpts]
+ # Extract the epoch number and validation loss from the checkpoint filenames
+ ckpt_epochs = np.array(
+ [
+ int(ckpt.parts[-1].split(".ckpt")[0].split("-")[0].split("=")[1])
+ for ckpt in ckpt_paths
+ ]
+ )
+ ckpt_val_losses = np.array(
+ [
+ (
+ float(ckpt.parts[-1].replace(".ckpt", "").split("-")[1].split("=")[1])
+ if "loss_val" in ckpt.parts[-1]
+ else 99999999.9
+ )
+ for ckpt in ckpt_paths
+ ]
+ )
+
+ # Determine the matching checkpoint index.
+ if self.load_epoch == "best":
+ ckpt_ix = ckpt_val_losses.argmin()
+ elif isinstance(self.load_epoch, int):
+ assert (
+ self.load_epoch in ckpt_epochs
+ ), f"Epoch {self.load_epoch} not found in checkpoints."
+ ckpt_ix = (ckpt_epochs == self.load_epoch).nonzero()[0][0].item()
+ else:
+ raise ValueError(f"Unrecognized load_epoch {self.load_epoch}")
+ ckpt = ckpts[ckpt_ix]
+ return ckpt
+
+
+class ProgressCallback(Protocol):
+ def __call__(self, progress: float):
+ """Callback which can be used to report progress on long-running inference.
+
+ Args:
+ progress: Float between 0 and 1.
+ """
+ pass
diff --git a/model/common/utils/data_utils.py b/model/common/utils/data_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..81a49a1043343d63aa99a23f86714667db2cce50
--- /dev/null
+++ b/model/common/utils/data_utils.py
@@ -0,0 +1,399 @@
+# Copyright (c) 2021 Tian Xie, Xiang Fu
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+# Adapted from https://github.com/txie-93/cdvae/blob/main/cdvae/common/data_utils.py
+
+from functools import lru_cache
+
+import numpy as np
+import torch
+from pymatgen.core import Element
+
+from onescience.datapipes.materials.mattergen.chemgraph import ChemGraph
+from ...common.utils.ocp_graph_utils import radius_graph_pbc as radius_graph_pbc_ocp
+
+EPSILON = 1e-5
+
+
+@lru_cache
+def get_atomic_number(symbol: str) -> int:
+ # get atomic number from Element symbol
+ return Element(symbol).Z
+
+
+@lru_cache
+def get_element_symbol(Z: int) -> str:
+ # get Element symbol from atomic number
+ return str(Element.from_Z(Z=Z))
+
+
+def abs_cap(val: float, max_abs_val: float = 1.0) -> float:
+ """
+ Returns the value with its absolute value capped at max_abs_val.
+ Particularly useful in passing values to trigonometric functions where
+ numerical errors may result in an argument > 1 being passed in.
+ https://github.com/materialsproject/pymatgen/blob/b789d74639aa851d7e5ee427a765d9fd5a8d1079/pymatgen/util/num.py#L15
+ Args:
+ val (float): Input value.
+ max_abs_val (float): The maximum absolute value for val. Defaults to 1.
+ Returns:
+ val if abs(val) < 1 else sign of val * max_abs_val.
+ """
+ return max(min(val, max_abs_val), -max_abs_val)
+
+
+def lattice_params_to_matrix(
+ a: float, b: float, c: float, alpha: float, beta: float, gamma: float
+) -> np.ndarray:
+ """Converts lattice from abc, angles to matrix.
+ https://github.com/materialsproject/pymatgen/blob/b789d74639aa851d7e5ee427a765d9fd5a8d1079/pymatgen/core/lattice.py#L311
+ """
+ angles_r = np.radians([alpha, beta, gamma])
+ cos_alpha, cos_beta, cos_gamma = np.cos(angles_r)
+ sin_alpha, sin_beta, sin_gamma = np.sin(angles_r)
+
+ val = (cos_alpha * cos_beta - cos_gamma) / (sin_alpha * sin_beta)
+ # Sometimes rounding errors result in values slightly > 1.
+ val = abs_cap(val)
+ gamma_star = np.arccos(val)
+
+ vector_a = [a * sin_beta, 0.0, a * cos_beta]
+ vector_b = [
+ -b * sin_alpha * np.cos(gamma_star),
+ b * sin_alpha * np.sin(gamma_star),
+ b * cos_alpha,
+ ]
+ vector_c = [0.0, 0.0, float(c)]
+ return np.array([vector_a, vector_b, vector_c])
+
+
+def lattice_params_to_matrix_torch(
+ lengths: torch.Tensor, angles: torch.Tensor, eps: float = 0.0
+) -> torch.Tensor:
+ """Batched torch version to compute lattice matrix from params.
+
+ lengths: torch.Tensor of shape (N, 3), unit A
+ angles: torch.Tensor of shape (N, 3), unit degree
+ """
+ coses = torch.clamp(torch.cos(torch.deg2rad(angles)), -1.0, 1.0)
+ sins = (1 - coses**2).sqrt()
+
+ val = (coses[:, 0] * coses[:, 1] - coses[:, 2]) / (sins[:, 0] * sins[:, 1])
+ val = torch.clamp(val, -1.0 + eps, 1.0 - eps)
+
+ vector_a = torch.stack(
+ [
+ lengths[:, 0] * sins[:, 1],
+ torch.zeros(lengths.size(0), device=lengths.device),
+ lengths[:, 0] * coses[:, 1],
+ ],
+ dim=1,
+ )
+ vector_b = torch.stack(
+ [
+ -lengths[:, 1] * sins[:, 0] * val,
+ lengths[:, 1] * sins[:, 0] * (1 - val**2).sqrt(),
+ lengths[:, 1] * coses[:, 0],
+ ],
+ dim=1,
+ )
+ vector_c = torch.stack(
+ [
+ torch.zeros(lengths.size(0), device=lengths.device),
+ torch.zeros(lengths.size(0), device=lengths.device),
+ lengths[:, 2],
+ ],
+ dim=1,
+ )
+
+ return torch.stack([vector_a, vector_b, vector_c], dim=1)
+
+
+def lattice_matrix_to_params_torch(
+ matrix: torch.Tensor, eps: float = 0.0
+) -> tuple[torch.Tensor, torch.Tensor]:
+ """Convert a batch of lattice matrices into their corresponding unit cell vector lengths and angles.
+
+ Args:
+ matrix (torch.Tensor, [B, 3, 3]): The batch of lattice matrices.
+
+ Returns:
+ tuple[torch.Tensor], ([B, 3], [B, 3]): tuple whose first element is the lengths of the unit cell vectors, and the second one gives the angles between the vectors.
+ """
+ assert len(matrix.shape) == 3
+
+ # derivatives of arccos(cos(theta)) are undefined for abs(cos(theta))=1
+ # we should physically encounter lattices that have vectors that are
+ # parallel to one another. NOTE: the value of eps may need tuning
+ # if calculations are found to fail, reduce this magnitude
+
+ lengths = matrix.norm(p=2, dim=-1)
+ ix_j = torch.tensor([1, 2, 0], dtype=torch.long, device=matrix.device)
+ ix_k = torch.tensor([2, 0, 1], dtype=torch.long, device=matrix.device)
+ cos_angles = (torch.cosine_similarity(matrix[:, ix_j], matrix[:, ix_k], dim=-1)).clamp(
+ -1 + eps, 1 - eps
+ )
+ if len(matrix.shape) == 2:
+ cos_angles = cos_angles.squeeze(0)
+ lengths = lengths.squeeze(0)
+ return lengths, torch.arccos(cos_angles) * 180.0 / np.pi
+
+
+def lattice_matrix_to_params(matrix: np.ndarray) -> tuple[float, float, float, float, float, float]:
+ lengths = np.sqrt(np.sum(matrix**2, axis=1)).tolist()
+
+ angles = np.zeros(3)
+ for i in range(3):
+ j = (i + 1) % 3
+ k = (i + 2) % 3
+ angles[i] = abs_cap(np.dot(matrix[j], matrix[k]) / (lengths[j] * lengths[k]))
+ angles = np.arccos(angles) * 180.0 / np.pi
+ a, b, c = lengths
+ alpha, beta, gamma = angles
+ return a, b, c, alpha, beta, gamma
+
+
+def frac_to_cart_coords(
+ frac_coords: torch.Tensor, lengths: torch.Tensor, angles: torch.Tensor, num_atoms: torch.Tensor
+) -> torch.Tensor:
+ lattice = lattice_params_to_matrix_torch(lengths, angles)
+ return frac_to_cart_coords_with_lattice(frac_coords, num_atoms, lattice)
+
+
+def cart_to_frac_coords(
+ cart_coords: torch.Tensor, lengths: torch.Tensor, angles: torch.Tensor, num_atoms: torch.Tensor
+) -> torch.Tensor:
+ lattice = lattice_params_to_matrix_torch(lengths, angles)
+ return cart_to_frac_coords_with_lattice(cart_coords, num_atoms, lattice)
+
+
+def frac_to_cart_coords_with_lattice(
+ frac_coords: torch.Tensor, num_atoms: torch.Tensor, lattice: torch.Tensor
+) -> torch.Tensor:
+ lattice_nodes = torch.repeat_interleave(lattice, num_atoms, dim=0)
+ pos = torch.einsum("bi,bij->bj", frac_coords, lattice_nodes) # cart coords
+ return pos
+
+
+def cart_to_frac_coords_with_lattice(
+ cart_coords: torch.Tensor, num_atoms: torch.Tensor, lattice: torch.Tensor
+) -> torch.Tensor:
+ # use pinv in case the predicted lattice is not rank 3
+ inv_lattice = torch.linalg.pinv(lattice)
+ inv_lattice_nodes = torch.repeat_interleave(inv_lattice, num_atoms, dim=0)
+ frac_coords = torch.einsum("bi,bij->bj", cart_coords, inv_lattice_nodes)
+ return frac_coords % 1.0
+
+
+def get_pbc_distances(
+ coords: torch.Tensor,
+ edge_index: torch.Tensor,
+ lattice: torch.Tensor,
+ to_jimages: torch.Tensor,
+ num_atoms: torch.Tensor,
+ num_bonds: torch.Tensor,
+ coord_is_cart: bool = False,
+ return_offsets: bool = False,
+ return_distance_vec: bool = False,
+) -> torch.Tensor:
+ if coord_is_cart:
+ pos = coords
+ else:
+ lattice_nodes = torch.repeat_interleave(lattice, num_atoms, dim=0)
+ pos = torch.einsum("bi,bij->bj", coords, lattice_nodes) # cart coords
+
+ j_index, i_index = edge_index
+
+ distance_vectors = pos[j_index] - pos[i_index]
+
+ # correct for pbc
+ lattice_edges = torch.repeat_interleave(lattice, num_bonds, dim=0)
+ offsets = torch.einsum("bi,bij->bj", to_jimages.float(), lattice_edges)
+ distance_vectors += offsets
+
+ # compute distances
+ distances = distance_vectors.norm(dim=-1)
+
+ out = {
+ "edge_index": edge_index,
+ "distances": distances,
+ }
+
+ if return_distance_vec:
+ out["distance_vec"] = distance_vectors
+
+ if return_offsets:
+ out["offsets"] = offsets
+
+ return out
+
+
+def radius_graph_pbc(
+ cart_coords: torch.Tensor,
+ lattice: torch.Tensor,
+ num_atoms: torch.Tensor,
+ radius: float,
+ max_num_neighbors_threshold: int,
+ max_cell_images_per_dim: int = 10,
+ topk_per_pair: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+ """Computes pbc graph edges under pbc.
+
+ topk_per_pair: (num_atom_pairs,), select topk edges per atom pair
+
+ Note: topk should take into account self-self edge for (i, i)
+
+ Keyword arguments
+ -----------------
+ cart_cords.shape=[Ntotal, 3] -- concatenate all atoms over all crystals
+ lattice.shape=[Ncrystal, 3, 3]
+ num_atoms.shape=[Ncrystal]
+ max_cell_images_per_dim -- constrain the max. number of cell images per dimension in event
+ that infinitesimal angles between lattice vectors are encountered.
+
+ WARNING: It is possible (and has been observed) that for rare cases when periodic atom images are
+ on or close to the cut off radius boundary, doing these operations in 32 bit floating point can
+ lead to atoms being spuriously considered within or outside of the cut off radius. This can lead
+ to invariance of the neighbour list under global translation of all atoms in the unit cell. For
+ the rare cases where this was observed, switching to 64 bit precision solved the issue. Since all
+ graph embeddings should taper messages from neighbours to zero at the cut off radius, the effect
+ of these errors in 32-bit should be negligible in practice.
+ """
+ assert topk_per_pair is None, "non None values of topk_per_pair is not supported"
+ edge_index, unit_cell, num_neighbors_image, _, _ = radius_graph_pbc_ocp(
+ pos=cart_coords,
+ cell=lattice,
+ natoms=num_atoms,
+ pbc=torch.Tensor([True, True, True])
+ .to(torch.bool)
+ .to(cart_coords.device), # torch.BoolTensor([...],device='cuda') fails
+ radius=radius,
+ max_num_neighbors_threshold=max_num_neighbors_threshold,
+ max_cell_images_per_dim=max_cell_images_per_dim,
+ )
+ return edge_index, unit_cell, num_neighbors_image
+
+
+class StandardScalerTorch(torch.nn.Module):
+ """Normalizes the targets of a dataset."""
+
+ def __init__(
+ self,
+ means: torch.Tensor | None = None,
+ stds: torch.Tensor | None = None,
+ stats_dim: tuple[int] = (
+ 1,
+ ), # dimension of mean, std stats (= X.shape[1:] for some input tensor X)
+ log10_transform: bool = False, # whether to log10-transform the property before scaling
+ ):
+ super().__init__()
+ # we need to make sure that we initialize means and stds with the right shapes
+ # otherwise, we cannot load checkpoints of fitted means/stds.
+ # ignore stats_dim if means and stds are provided
+ self.register_buffer(
+ "means", torch.atleast_1d(means) if means is not None else torch.empty(stats_dim)
+ )
+ self.register_buffer(
+ "stds", torch.atleast_1d(stds) if stds is not None else torch.empty(stats_dim)
+ )
+ self.log10_transform = log10_transform
+
+ @property
+ def device(self) -> torch.device:
+ return self.means.device # type: ignore
+
+ def fit(self, X: torch.Tensor):
+ if self.log10_transform:
+ assert torch.all(X > 0), "All values must be positive for log10 transform"
+ X = torch.log10(X)
+
+ means: torch.Tensor = torch.atleast_1d(torch.nanmean(X, dim=0).to(self.device))
+ stds: torch.Tensor = torch.atleast_1d(
+ torch_nanstd(X, dim=0, unbiased=False).to(self.device) + EPSILON
+ )
+ # mypy gets really confused about variables registered via register_buffer,
+ # so we need to ignore a lot of type errors below
+ assert (
+ means.shape == self.means.shape # type: ignore
+ ), f"Mean shape mismatch: {means.shape} != {self.means.shape}" # type: ignore
+ assert (
+ stds.shape == self.stds.shape # type: ignore
+ ), f"Std shape mismatch: {stds.shape} != {self.stds.shape}" # type: ignore
+ self.means = means # type: ignore
+ self.stds = stds # type: ignore
+
+ def transform(self, X: torch.Tensor) -> torch.Tensor:
+ assert self.means is not None and self.stds is not None
+ if self.log10_transform:
+ assert torch.all(X > 0), "All values must be positive for log10 transform"
+ X = torch.log10(X)
+ return (X - self.means) / self.stds
+
+ def inverse_transform(self, X: torch.Tensor) -> torch.Tensor:
+ assert self.means is not None and self.stds is not None
+ X = X * self.stds + self.means
+ if self.log10_transform:
+ X = torch.pow(10, X)
+ return X
+
+ def match_device(self, X: torch.Tensor) -> torch.Tensor:
+ assert self.means.numel() > 0 and self.stds.numel() > 0
+ if self.means.device != X.device:
+ self.means = self.means.to(X.device)
+ self.stds = self.stds.to(X.device)
+
+ def copy(self) -> "StandardScalerTorch":
+ return StandardScalerTorch(
+ means=self.means.clone().detach(),
+ stds=self.stds.clone().detach(),
+ log10_transform=self.log10_transform,
+ )
+
+ def forward(self, X: torch.Tensor) -> torch.Tensor:
+ return self.transform(X)
+
+ def __repr__(self) -> str:
+ return (
+ f"{self.__class__.__name__}("
+ f"means: {self.means.tolist() if self.means is not None else None}, "
+ f"stds: {self.stds.tolist() if self.stds is not None else None})"
+ f"log10_transform: {self.log10_transform}"
+ )
+
+
+def torch_nanstd(x: torch.Tensor, dim: int, unbiased: bool) -> torch.Tensor:
+ data_is_present = torch.all(
+ torch.reshape(torch.logical_not(torch.isnan(x)), (x.shape[0], -1)),
+ dim=1,
+ )
+ # https://github.com/pytorch/pytorch/issues/29372
+ return torch.std(x[data_is_present], dim=dim, unbiased=unbiased)
+
+
+def compute_lattice_polar_decomposition(lattice_matrix: torch.Tensor) -> torch.Tensor:
+ # Polar decomposition via SVD, see https://en.wikipedia.org/wiki/Polar_decomposition
+ # lattice_matrix: [batch_size, 3, 3]
+ # Computes the (unique) symmetric lattice matrix that is equivalent (up to rotation) to the input lattice.
+
+ W, S, V_transp = torch.linalg.svd(lattice_matrix)
+ S_square = torch.diag_embed(S)
+ V = V_transp.transpose(1, 2)
+ U = W @ V_transp
+ P = V @ S_square @ V_transp
+ P_prime = U @ P @ U.transpose(1, 2)
+ # symmetrized lattice matrix
+ symm_lattice_matrix = P_prime
+ return symm_lattice_matrix
+
+
+def create_chem_graph_from_composition(target_composition_dict: dict[str, float]) -> ChemGraph:
+ atomic_numbers = []
+ for element_name, number_of_atoms in target_composition_dict.items():
+ atomic_numbers += [Element(element_name).Z] * int(number_of_atoms)
+
+ return ChemGraph(
+ atomic_numbers=torch.tensor(atomic_numbers, dtype=torch.long),
+ num_atoms=torch.tensor([len(atomic_numbers)], dtype=torch.long),
+ cell=torch.eye(3, dtype=torch.float).reshape(1, 3, 3),
+ pos=torch.zeros((len(atomic_numbers), 3), dtype=torch.float),
+ )
diff --git a/model/common/utils/eval_utils.py b/model/common/utils/eval_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..02f5e78e64019f599cba85bea343c88b799c4900
--- /dev/null
+++ b/model/common/utils/eval_utils.py
@@ -0,0 +1,167 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+import logging
+import os
+from pathlib import Path
+from tempfile import TemporaryDirectory
+from typing import Sequence
+from zipfile import ZipFile
+
+import ase.io
+import hydra
+import numpy as np
+import torch
+from pymatgen.core import Lattice, Structure
+from pymatgen.io.ase import AseAtomsAdaptor
+
+from ...common.globals import (
+ GENERATED_CRYSTALS_EXTXYZ_FILE_NAME,
+ GENERATED_CRYSTALS_ZIP_FILE_NAME,
+)
+from ...common.utils.data_classes import MatterGenCheckpointInfo
+from ...common.utils.globals import get_device
+from ...diffusion.lightning_module import DiffusionLightningModule
+
+# logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+def make_structure(
+ lengths: torch.Tensor,
+ angles: torch.Tensor,
+ atom_types: torch.Tensor,
+ frac_coords: torch.Tensor,
+) -> Structure:
+ return Structure(
+ lattice=Lattice.from_parameters(
+ **{a: v for a, v in zip(["a", "b", "c"], lengths)},
+ **{a: v for a, v in zip(["alpha", "beta", "gamma"], angles)},
+ ),
+ species=atom_types,
+ coords=frac_coords,
+ coords_are_cartesian=False,
+ )
+
+
+def load_model_diffusion(
+ args: MatterGenCheckpointInfo,
+) -> DiffusionLightningModule:
+ assert args.load_epoch is not None
+ ckpt = args.checkpoint_path
+ logger.info(f"Loading model from checkpoint: {ckpt}")
+ cfg = args.config
+ try:
+ model, incompatible_keys = DiffusionLightningModule.load_from_checkpoint_and_config(
+ ckpt,
+ map_location=get_device(),
+ config=cfg.lightning_module,
+ strict=args.strict_checkpoint_loading,
+ )
+ except hydra.errors.HydraException as e:
+ raise
+ if len(incompatible_keys.unexpected_keys) > 0:
+ raise ValueError(f"Unexpected keys in checkpoint: {incompatible_keys.unexpected_keys}.")
+ if len(incompatible_keys.missing_keys) > 0:
+ raise ValueError(f"Missing keys in checkpoint: {incompatible_keys.missing_keys}.")
+
+ return model
+
+
+def get_crystals_list(
+ frac_coords, atom_types, lengths, angles, num_atoms
+) -> list[dict[str, np.ndarray]]:
+ """
+ args:
+ frac_coords: (num_atoms, 3)
+ atom_types: (num_atoms)
+ lengths: (num_crystals)
+ angles: (num_crystals)
+ num_atoms: (num_crystals)
+ """
+ assert frac_coords.size(0) == atom_types.size(0) == num_atoms.sum()
+ assert lengths.size(0) == angles.size(0) == num_atoms.size(0)
+
+ start_idx = 0
+ crystal_array_list = []
+ for batch_idx, num_atom in enumerate(num_atoms.tolist()):
+ cur_frac_coords = frac_coords.narrow(0, start_idx, num_atom)
+ cur_atom_types = atom_types.narrow(0, start_idx, num_atom)
+ cur_lengths = lengths[batch_idx]
+ cur_angles = angles[batch_idx]
+
+ crystal_array_list.append(
+ {
+ "frac_coords": cur_frac_coords.detach().cpu().numpy(),
+ "atom_types": cur_atom_types.detach().cpu().numpy(),
+ "lengths": cur_lengths.detach().cpu().numpy(),
+ "angles": cur_angles.detach().cpu().numpy(),
+ }
+ )
+ start_idx = start_idx + num_atom
+ return crystal_array_list
+
+
+def save_structures(output_path: Path, structures: Sequence[Structure]) -> None:
+ """Save structures to disk in a extxyz file and a compressed zip file containing cif files.
+
+ Args:
+ output_path: path to a directory where the results are written.
+ structures: sequence of structures.
+ """
+ ase_atoms = [AseAtomsAdaptor.get_atoms(x) for x in structures]
+ try:
+ ase.io.write(output_path / GENERATED_CRYSTALS_EXTXYZ_FILE_NAME, ase_atoms)
+
+ with ZipFile(output_path / GENERATED_CRYSTALS_ZIP_FILE_NAME, "w") as zip_obj:
+ for ix, ase_atom in enumerate(ase_atoms):
+ ase.io.write(f"/tmp/gen_{ix}.cif", ase_atom, format="cif")
+ zip_obj.write(f"/tmp/gen_{ix}.cif", arcname=f"gen_{ix}.cif")
+ except IOError as e:
+ print(f"Got error {e} writing the generated structures to disk.")
+
+
+def load_structures(input_path: Path) -> Sequence[Structure]:
+ """Load structures from disk.
+
+ Args:
+ output_path: path to a file or directory where the results are written.
+
+ Returns:
+ sequence of structures.
+ """
+ # if the path is an xyz or extxyz file, read it directly
+ if input_path.suffix == ".xyz" or input_path.suffix == ".extxyz":
+ ase_atoms = ase.io.read(input_path, ":")
+ return [AseAtomsAdaptor.get_structure(x) for x in ase_atoms]
+
+ # if the path is a zipped folder, extract it into a temporary directory
+ elif input_path.suffix == ".zip":
+ with TemporaryDirectory() as tmpdirname:
+ with ZipFile(input_path, "r") as zip_obj:
+ zip_obj.extractall(tmpdirname)
+ return extract_structures_from_folder(tmpdirname)
+
+ # if the path is a directory, read all files in it
+ elif input_path.is_dir():
+ return extract_structures_from_folder(input_path)
+
+ else:
+ raise ValueError(f"Invalid input path {input_path}")
+
+
+def extract_structures_from_folder(dirname: str) -> Sequence[Structure]:
+ structures = []
+ for filename in os.listdir(dirname):
+ if filename.endswith(".cif"):
+ try:
+ structures.append(Structure.from_file(f"{dirname}/{filename}"))
+ except ValueError as e:
+ logger.warning(f"Failed to read {filename} as a CIF file: {e}")
+ elif filename.endswith(".extxyz") or filename.endswith(".xyz"):
+ ase_atoms = ase.io.read(
+ f"{dirname}/{filename}", 0
+ ) # We assume that the file contains only one structure
+ structures.append(AseAtomsAdaptor.get_structure(ase_atoms))
+ return structures
diff --git a/model/common/utils/globals.py b/model/common/utils/globals.py
new file mode 100644
index 0000000000000000000000000000000000000000..9988adbab057d8cb9f60edac03fddaef43db678b
--- /dev/null
+++ b/model/common/utils/globals.py
@@ -0,0 +1,152 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+"""Note that importing this module has two side effects:
+1. It sets the environment variable `PROJECT_ROOT` to the root of the explorers project.
+2. It registers a new resolver for OmegaConf, `eval`, which allows us to use `eval` in our config files.
+"""
+import os
+from functools import lru_cache
+from pathlib import Path
+
+import torch
+from omegaconf import OmegaConf
+
+
+@lru_cache
+def get_device() -> torch.device:
+ if torch.cuda.is_available():
+ return torch.device("cuda")
+ if torch.backends.mps.is_available():
+ return torch.device("mps")
+ return torch.device("cpu")
+
+
+@lru_cache
+def get_pyg_device() -> torch.device:
+ """
+ Some operations of pyg don't work on MPS, so fall back to CPU.
+ """
+ if torch.cuda.is_available():
+ return torch.device("cuda")
+ return torch.device("cpu")
+
+
+MODELS_PROJECT_ROOT = Path(__file__).resolve().parents[2]
+print(f"MODELS_PROJECT_ROOT: {MODELS_PROJECT_ROOT}")
+
+# Set environment variable PROJECT_ROOT so that hydra / OmegaConf can access it.
+os.environ.setdefault("PROJECT_ROOT", str(MODELS_PROJECT_ROOT)) # for hydra
+
+DEFAULT_SAMPLING_CONFIG_PATH = MODELS_PROJECT_ROOT / "sampling_conf"
+PROPERTY_SOURCE_IDS = [
+ "dft_mag_density",
+ "dft_bulk_modulus",
+ "dft_shear_modulus",
+ "energy_above_hull",
+ "formation_energy_per_atom",
+ "space_group",
+ "hhi_score",
+ "ml_bulk_modulus",
+ "chemical_system",
+ "dft_band_gap",
+]
+
+SELECTED_ATOMIC_NUMBERS = [
+ 1,
+ 3,
+ 4,
+ 5,
+ 6,
+ 7,
+ 8,
+ 9,
+ 11,
+ 12,
+ 13,
+ 14,
+ 15,
+ 16,
+ 17,
+ 19,
+ 20,
+ 21,
+ 22,
+ 23,
+ 24,
+ 25,
+ 26,
+ 27,
+ 28,
+ 29,
+ 30,
+ 31,
+ 32,
+ 33,
+ 34,
+ 35,
+ 37,
+ 38,
+ 39,
+ 40,
+ 41,
+ 42,
+ 44,
+ 45,
+ 46,
+ 47,
+ 48,
+ 49,
+ 50,
+ 51,
+ 52,
+ 53,
+ 55,
+ 56,
+ 57,
+ 58,
+ 59,
+ 60,
+ 62,
+ 63,
+ 64,
+ 65,
+ 66,
+ 67,
+ 68,
+ 69,
+ 70,
+ 71,
+ 72,
+ 73,
+ 74,
+ 75,
+ 76,
+ 77,
+ 78,
+ 79,
+ 80,
+ 81,
+ 82,
+ 83,
+]
+MAX_ATOMIC_NUM = 100
+
+
+# Set `eval` resolver
+def try_eval(s):
+ """This is a custom resolver for OmegaConf that allows us to use `eval` in our config files
+ with the syntax `${eval:'${foo} + ${bar}'}
+
+ See:
+ https://omegaconf.readthedocs.io/en/2.3_branch/how_to_guides.html#id1
+ """
+ try:
+ return eval(s)
+ except Exception as e:
+ print(f"Calling eval on string {s} raised exception {e}")
+ raise
+
+
+if not OmegaConf.has_resolver("eval"):
+ OmegaConf.register_new_resolver("eval", try_eval)
diff --git a/model/common/utils/lattice_score.py b/model/common/utils/lattice_score.py
new file mode 100644
index 0000000000000000000000000000000000000000..a4bc6086d0d4ced14acebab474c1fdd7e265f8d9
--- /dev/null
+++ b/model/common/utils/lattice_score.py
@@ -0,0 +1,37 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+import torch
+from torch_scatter import scatter_add
+
+
+def edge_score_to_lattice_score_frac_symmetric(
+ score_d: torch.Tensor,
+ edge_index: torch.Tensor,
+ edge_vectors: torch.Tensor,
+ batch: torch.Tensor,
+) -> torch.Tensor:
+ """Converts a score per edge into a score for the atom coordinates and/or the lattice matrix via the chain rule.
+ This method explicitly takes into account the fact that the cartesian coordinates depend on the lattice via the fractional coordinates.
+ Moreover, we make sure to get a symmetric update: D_cart_norm @ Phi @ D_cart_norm^T, where Phi is a |E| x |E| diagonal matrix with the predicted edge scores
+
+ Args:
+ score_d (torch.Tensor, [num_edges,]): A score per edge in the graph.
+ edge_index (torch.Tensor, [2, num_edges]): The edge indices in the graph.
+ edge_vectors (torch.Tensor, [num_edges, 3]): The vectors connecting the source of each edge to the target.
+ lattice_matrix (torch.Tensor, [num_nodes, 3, 3]): The lattice matrices for each crystal in num_nodes.
+ batch (torch.Tensor, [num_nodes,]): The pointer indicating for each atom which molecule in the batch it belongs to.
+
+ Returns:
+ torch.Tensor: The predicted lattice score.
+ """
+ batch_edge = batch[edge_index[0]]
+ unit_edge_vectors_cart = edge_vectors / edge_vectors.norm(dim=-1, keepdim=True)
+ score_lattice = scatter_add(
+ score_d[:, None, None]
+ * (unit_edge_vectors_cart[:, :, None] @ unit_edge_vectors_cart[:, None, :]),
+ batch_edge,
+ dim=0,
+ dim_size=batch.max() + 1,
+ ).transpose(-1, -2)
+ return score_lattice
diff --git a/model/common/utils/ocp_graph_utils.py b/model/common/utils/ocp_graph_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..e7dc72c3b3885ef08e32de12d0c3afc66c222ca6
--- /dev/null
+++ b/model/common/utils/ocp_graph_utils.py
@@ -0,0 +1,334 @@
+"""
+Copyright (c) Facebook, Inc. and its affiliates.
+Copyright (c) Microsoft Corporation.
+Licensed under the MIT License.
+Code derived from the OCP codebase:
+https://github.com/Open-Catalyst-Project/ocp
+"""
+
+import sys
+
+import numpy as np
+import torch
+from torch_scatter import segment_coo, segment_csr
+
+from ...common.utils.globals import get_pyg_device
+
+
+def get_pbc_distances(
+ pos: torch.Tensor,
+ edge_index: torch.Tensor,
+ cell: torch.Tensor,
+ cell_offsets: torch.Tensor,
+ neighbors: torch.Tensor,
+ return_offsets: bool = False,
+ return_distance_vec: bool = False,
+) -> dict:
+ row, col = edge_index
+
+ distance_vectors = pos[row] - pos[col]
+
+ # correct for pbc
+ neighbors = neighbors.to(cell.device)
+ cell = torch.repeat_interleave(cell, neighbors, dim=0)
+ offsets = cell_offsets.float().view(-1, 1, 3).bmm(cell.float()).view(-1, 3)
+ distance_vectors += offsets
+
+ # compute distances
+ distances = distance_vectors.norm(dim=-1)
+
+ # redundancy: remove zero distances
+ nonzero_idx = torch.arange(len(distances))[distances > 0]
+ edge_index = edge_index[:, nonzero_idx]
+ distances = distances[nonzero_idx]
+
+ out = {
+ "edge_index": edge_index,
+ "distances": distances,
+ }
+
+ if return_distance_vec:
+ out["distance_vec"] = distance_vectors[nonzero_idx]
+
+ if return_offsets:
+ out["offsets"] = offsets[nonzero_idx]
+
+ return out
+
+
+def radius_graph_pbc(
+ pos: torch.Tensor,
+ pbc: torch.Tensor | None,
+ natoms: torch.Tensor,
+ cell: torch.Tensor,
+ radius: float,
+ max_num_neighbors_threshold: int,
+ max_cell_images_per_dim: int = sys.maxsize,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+ """Function computing the graph in periodic boundary conditions on a (batched) set of
+ positions and cells.
+
+ This function is copied from
+ https://github.com/Open-Catalyst-Project/ocp/blob/main/ocpmodels/common/utils.py,
+ commit 480eb9279ec4a5885981f1ee588c99dcb38838b5
+
+ Args:
+ pos (LongTensor): Atomic positions in cartesian coordinates
+ :obj:`[n, 3]`
+ pbc (BoolTensor): indicates periodic boundary conditions per structure.
+ :obj:`[n_structures, 3]`
+ natoms (IntTensor): number of atoms per structure. Has shape
+ :obj:`[n_structures]`
+ cell (Tensor): atomic cell. Has shape
+ :obj:`[n_structures, 3, 3]`
+ radius (float): cutoff radius distance
+ max_num_neighbors_threshold (int): Maximum number of neighbours to consider.
+
+ Returns:
+ edge_index (IntTensor): index of atoms in edges. Has shape
+ :obj:`[n_edges, 2]`
+ cell_offsets (IntTensor): cell displacement w.r.t. their original position of atoms in edges. Has shape
+ :obj:`[n_edges, 3, 3]`
+ num_neighbors_image (IntTensor): Number of neighbours per cell image.
+ :obj:`[n_structures]`
+ offsets (LongTensor): cartesian displacement w.r.t. their original position of atoms in edges. Has shape
+ :obj:`[n_edges, 3, 3]`
+ atom_distance (LongTensor): edge length. Has shape
+ :obj:`[n_edges]`
+ """
+ device = pos.device
+ batch_size = len(natoms)
+ pbc_ = [False, False, False]
+
+ if pbc is not None:
+ pbc = torch.atleast_2d(pbc)
+ for i in range(3):
+ if not torch.any(pbc[:, i]).item():
+ pbc_[i] = False
+ elif torch.all(pbc[:, i]).item():
+ pbc_[i] = True
+ else:
+ raise RuntimeError(
+ "Different structures in the batch have different PBC configurations. This is not currently supported."
+ )
+
+ natoms_squared = (natoms**2).long()
+
+ # index offset between images
+ index_offset = torch.cumsum(natoms, dim=0) - natoms
+
+ index_offset_expand = torch.repeat_interleave(index_offset, natoms_squared)
+ natoms_expand = torch.repeat_interleave(natoms, natoms_squared)
+
+ # Compute a tensor containing sequences of numbers that range from 0 to num_atoms_per_image_squared for each image
+ # that is used to compute indices for the pairs of atoms. This is a very convoluted way to implement
+ # the following (but 10x faster since it removes the for loop)
+ # for batch_idx in range(batch_size):
+ # batch_count = torch.cat([batch_count, torch.arange(num_atoms_per_image_squared[batch_idx], device=device)], dim=0)
+ num_atom_pairs = torch.sum(natoms_squared)
+ index_squared_offset = torch.cumsum(natoms_squared, dim=0) - natoms_squared
+ index_squared_offset = torch.repeat_interleave(index_squared_offset, natoms_squared)
+ atom_count_squared = torch.arange(num_atom_pairs, device=device) - index_squared_offset
+
+ # Compute the indices for the pairs of atoms (using division and mod)
+ # If the systems get too large this approach could run into numerical precision issues
+ index1 = (
+ torch.div(atom_count_squared, natoms_expand, rounding_mode="floor")
+ ) + index_offset_expand
+ index2 = (atom_count_squared % natoms_expand) + index_offset_expand
+ # Get the positions for each atom
+ pos1 = torch.index_select(pos, 0, index1)
+ pos2 = torch.index_select(pos, 0, index2)
+
+ # Calculate required number of unit cells in each direction.
+ # Smallest distance between planes separated by a1 is
+ # 1 / ||(a2 x a3) / V||_2, since a2 x a3 is the area of the plane.
+ # Note that the unit cell volume V = a1 * (a2 x a3) and that
+ # (a2 x a3) / V is also the reciprocal primitive vector
+ # (crystallographer's definition).
+
+ cross_a2a3 = torch.cross(cell[:, 1], cell[:, 2], dim=-1)
+ cell_vol = torch.sum(cell[:, 0] * cross_a2a3, dim=-1, keepdim=True)
+
+ if pbc_[0]:
+ inv_min_dist_a1 = torch.norm(cross_a2a3 / cell_vol, p=2, dim=-1)
+ rep_a1 = torch.ceil(radius * inv_min_dist_a1)
+ else:
+ rep_a1 = cell.new_zeros(1)
+
+ if pbc_[1]:
+ cross_a3a1 = torch.cross(cell[:, 2], cell[:, 0], dim=-1)
+ inv_min_dist_a2 = torch.norm(cross_a3a1 / cell_vol, p=2, dim=-1)
+ rep_a2 = torch.ceil(radius * inv_min_dist_a2)
+ else:
+ rep_a2 = cell.new_zeros(1)
+
+ if pbc_[2]:
+ cross_a1a2 = torch.cross(cell[:, 0], cell[:, 1], dim=-1)
+ inv_min_dist_a3 = torch.norm(cross_a1a2 / cell_vol, p=2, dim=-1)
+ rep_a3 = torch.ceil(radius * inv_min_dist_a3)
+ else:
+ rep_a3 = cell.new_zeros(1)
+
+ # Take the max over all images for uniformity. This is essentially padding.
+ # Note that this can significantly increase the number of computed distances
+ # if the required repetitions are very different between images
+ # (which they usually are). Changing this to sparse (scatter) operations
+ # might be worth the effort if this function becomes a bottleneck.
+ #
+ # max_cell_images_per_dim limits the number of periodic
+ # cell images that are considered per lattice vector dimension. This is
+ # useful in case we encounter an extremely skewed or small lattice that
+ # results in an explosion of the number of images considered.
+ max_rep = [
+ min(int(rep_a1.max()), max_cell_images_per_dim),
+ min(int(rep_a2.max()), max_cell_images_per_dim),
+ min(int(rep_a3.max()), max_cell_images_per_dim),
+ ]
+
+ # Tensor of unit cells
+ cells_per_dim = [
+ torch.arange(-rep, rep + 1, device=device, dtype=torch.float) for rep in max_rep
+ ]
+ cell_offsets = torch.cartesian_prod(*cells_per_dim)
+ num_cells = len(cell_offsets)
+ cell_offsets_per_atom = cell_offsets.view(1, num_cells, 3).repeat(len(index2), 1, 1)
+ cell_offsets = torch.transpose(cell_offsets, 0, 1)
+ cell_offsets_batch = cell_offsets.view(1, 3, num_cells).expand(batch_size, -1, -1)
+
+ # Compute the x, y, z positional offsets for each cell in each image
+ data_cell = torch.transpose(cell, 1, 2)
+ pbc_offsets = torch.bmm(data_cell, cell_offsets_batch)
+ pbc_offsets_per_atom = torch.repeat_interleave(pbc_offsets, natoms_squared, dim=0)
+
+ # Expand the positions and indices for the 9 cells
+ pos1 = pos1.view(-1, 3, 1).expand(-1, -1, num_cells)
+ pos2 = pos2.view(-1, 3, 1).expand(-1, -1, num_cells)
+ index1 = index1.view(-1, 1).repeat(1, num_cells).view(-1)
+ index2 = index2.view(-1, 1).repeat(1, num_cells).view(-1)
+ # Add the PBC offsets for the second atom
+ pos2 = pos2 + pbc_offsets_per_atom
+
+ # Compute the squared distance between atoms
+ atom_distance_squared = torch.sum((pos1 - pos2) ** 2, dim=1)
+ atom_distance_squared = atom_distance_squared.view(-1)
+
+ # Remove pairs that are too far apart
+ mask_within_radius = torch.le(atom_distance_squared, radius * radius)
+ # Remove pairs with the same atoms (distance = 0.0)
+ mask_not_same = torch.gt(atom_distance_squared, 0.0001)
+ mask = torch.logical_and(mask_within_radius, mask_not_same)
+ index1 = torch.masked_select(index1, mask)
+ index2 = torch.masked_select(index2, mask)
+ cell_offsets = torch.masked_select(
+ cell_offsets_per_atom.view(-1, 3), mask.view(-1, 1).expand(-1, 3)
+ )
+ cell_offsets = cell_offsets.view(-1, 3)
+ atom_distance_squared = torch.masked_select(atom_distance_squared, mask)
+
+ mask_num_neighbors, num_neighbors_image = get_max_neighbors_mask(
+ natoms=natoms,
+ index=index1,
+ atom_distance_squared=atom_distance_squared,
+ max_num_neighbors_threshold=max_num_neighbors_threshold,
+ )
+
+ if not torch.all(mask_num_neighbors):
+ # Mask out the atoms to ensure each atom has at most max_num_neighbors_threshold neighbors
+ index1 = torch.masked_select(index1, mask_num_neighbors)
+ index2 = torch.masked_select(index2, mask_num_neighbors)
+ atom_distance_squared = torch.masked_select(atom_distance_squared, mask_num_neighbors)
+ cell_offsets = torch.masked_select(
+ cell_offsets.view(-1, 3), mask_num_neighbors.view(-1, 1).expand(-1, 3)
+ )
+ cell_offsets = cell_offsets.view(-1, 3)
+
+ edge_index = torch.stack((index2, index1))
+ # shifts = -torch.matmul(unit_cell, data.cell).view(-1, 3)
+
+ cell_repeated = torch.repeat_interleave(cell, num_neighbors_image, dim=0)
+ offsets = -cell_offsets.float().view(-1, 1, 3).bmm(cell_repeated.float()).view(-1, 3)
+ return (
+ edge_index,
+ cell_offsets,
+ num_neighbors_image,
+ offsets,
+ torch.sqrt(atom_distance_squared),
+ )
+
+
+def get_max_neighbors_mask(
+ natoms: torch.Tensor,
+ index: torch.Tensor,
+ atom_distance_squared: torch.Tensor,
+ max_num_neighbors_threshold: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
+ """
+ Give a mask that filters out edges so that each atom has at most
+ `max_num_neighbors_threshold` neighbors.
+ Assumes that `index` is sorted.
+ """
+ device = natoms.device
+ num_atoms = natoms.sum()
+
+ # Get number of neighbors
+ # segment_coo assumes sorted index
+ ones = index.new_ones(1).expand_as(index)
+ # required because PyG does not support MPS for the segment_coo operation yet.
+ pyg_device = get_pyg_device()
+ device_before = ones.device
+ num_neighbors = segment_coo(ones.to(pyg_device), index.to(pyg_device), dim_size=num_atoms).to(
+ device_before
+ )
+ max_num_neighbors = num_neighbors.max()
+ num_neighbors_thresholded = num_neighbors.clamp(max=max_num_neighbors_threshold)
+
+ # Get number of (thresholded) neighbors per image
+ image_indptr = torch.zeros(natoms.shape[0] + 1, device=device, dtype=torch.long)
+ image_indptr[1:] = torch.cumsum(natoms, dim=0)
+ num_neighbors_image = segment_csr(
+ num_neighbors_thresholded.to(pyg_device), image_indptr.to(pyg_device)
+ ).to(device_before)
+
+ # If max_num_neighbors is below the threshold, return early
+ if max_num_neighbors <= max_num_neighbors_threshold or max_num_neighbors_threshold <= 0:
+ mask_num_neighbors = torch.tensor([True], dtype=bool, device=device).expand_as(index)
+ return mask_num_neighbors, num_neighbors_image
+
+ # Create a tensor of size [num_atoms, max_num_neighbors] to sort the distances of the neighbors.
+ # Fill with infinity so we can easily remove unused distances later.
+ distance_sort = torch.full([num_atoms * max_num_neighbors], np.inf, device=device)
+
+ # Create an index map to map distances from atom_distance to distance_sort
+ # index_sort_map assumes index to be sorted
+ index_neighbor_offset = torch.cumsum(num_neighbors, dim=0) - num_neighbors
+ index_neighbor_offset_expand = torch.repeat_interleave(index_neighbor_offset, num_neighbors)
+ index_sort_map = (
+ index * max_num_neighbors
+ + torch.arange(len(index), device=device)
+ - index_neighbor_offset_expand
+ )
+ distance_sort.index_copy_(0, index_sort_map, atom_distance_squared)
+ distance_sort = distance_sort.view(num_atoms, max_num_neighbors)
+
+ # Sort neighboring atoms based on distance
+ distance_sort, index_sort = torch.sort(distance_sort, dim=1)
+ # Select the max_num_neighbors_threshold neighbors that are closest
+ distance_sort = distance_sort[:, :max_num_neighbors_threshold]
+ index_sort = index_sort[:, :max_num_neighbors_threshold]
+
+ # Offset index_sort so that it indexes into index
+ index_sort = index_sort + index_neighbor_offset.view(-1, 1).expand(
+ -1, max_num_neighbors_threshold
+ )
+ # Remove "unused pairs" with infinite distances
+ mask_finite = torch.isfinite(distance_sort)
+ index_sort = torch.masked_select(index_sort, mask_finite)
+
+ # At this point index_sort contains the index into index of the
+ # closest max_num_neighbors_threshold neighbors per atom
+ # Create a mask to remove all pairs not in index_sort
+ mask_num_neighbors = torch.zeros(len(index), device=device, dtype=bool)
+ mask_num_neighbors.index_fill_(0, index_sort, True)
+
+ return mask_num_neighbors, num_neighbors_image
diff --git a/model/denoiser.py b/model/denoiser.py
new file mode 100644
index 0000000000000000000000000000000000000000..de6bf1905e48b5cce90d0858f17990f0b82963be
--- /dev/null
+++ b/model/denoiser.py
@@ -0,0 +1,282 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+from typing import Callable
+
+import torch
+import torch.nn as nn
+
+from onescience.datapipes.materials.mattergen.chemgraph import ChemGraph
+from onescience.datapipes.materials.mattergen.types import PropertySourceId
+from .common.utils.globals import MAX_ATOMIC_NUM, SELECTED_ATOMIC_NUMBERS
+from .diffusion.model_utils import NoiseLevelEncoding
+from .diffusion.score_models.base import ScoreModel
+from onescience.modules.embedding.mattergen_property_embeddings import (
+ ChemicalSystemMultiHotEmbedding,
+ get_property_embeddings,
+ get_use_unconditional_embedding,
+)
+
+BatchTransform = Callable[[ChemGraph], ChemGraph]
+
+
+def atomic_numbers_to_mask(atomic_numbers: torch.LongTensor, max_atomic_num: int) -> torch.Tensor:
+ """Convert atomic numbers to a mask.
+
+ Args:
+ atomic_numbers (torch.LongTensor): One-based atomic numbers of shape (batch_size, )
+
+ Returns:
+ torch.Tensor: Mask of shape (batch_size, num_classes)
+ """
+ k_hot_mask = torch.eye(max_atomic_num, device=atomic_numbers.device)[atomic_numbers - 1]
+ return k_hot_mask
+
+
+def mask_logits(logits: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
+ """Mask logits by setting the logits for masked items to -inf.
+
+ Args:
+ logits (torch.Tensor): Logits of shape (batch_size, num_classes)
+ mask (torch.Tensor): Mask of shape (batch_size, num_classes). Values with zero are masked.
+
+ Returns:
+ torch.Tensor: Masked logits
+ """
+ return logits + (1 - mask) * -1e10
+
+
+def mask_disallowed_elements(
+ logits: torch.FloatTensor,
+ x: ChemGraph | None = None,
+ batch_idx: torch.LongTensor | None = None,
+ predictions_are_zero_based: bool = True,
+):
+ """
+ Mask out atom types that are disallowed in general,
+ as well as potentially all elements not in the chemical system we condition on.
+
+ Args:
+ logits (torch.Tensor): Logits of shape (batch_size, num_classes)
+ x (ChemGraph)
+ batch_idx (torch.LongTensor, optional): Batch indices. Defaults to None. Must be provided if condition is not None.
+ predictions_are_zero_based (bool, optional): Whether the logits are zero-based. Defaults to True. Basically, if we're using D3PM,
+ the logits are zero-based (model predicts atomic number index)
+ """
+ # First, mask out generally undesired elements
+ # (1, num_selected_elements)
+ selected_atomic_numbers = torch.tensor(SELECTED_ATOMIC_NUMBERS, device=logits.device)
+ predictions_are_one_based = not predictions_are_zero_based
+ # (num_atoms, num_classes)
+ one_hot_selected_elements = atomic_numbers_to_mask(
+ atomic_numbers=selected_atomic_numbers + int(predictions_are_one_based),
+ max_atomic_num=logits.shape[1],
+ )
+ # (1, num_classes)
+ k_hot_mask = one_hot_selected_elements.sum(0)[None]
+ # Set the logits for disallowed elements to -inf
+ logits = mask_logits(logits=logits, mask=k_hot_mask)
+
+ # Optionally, also mask out elements that are not in the chemical system we condition on
+ if x is not None and "chemical_system" in x and x["chemical_system"] is not None:
+ try:
+ # torch.BoolTensor, shape (batch_size, 1) -- do not mask logits when we use an unconditional embedding
+ do_not_mask_atom_logits = get_use_unconditional_embedding(
+ batch=x, cond_field="chemical_system"
+ )
+ except KeyError:
+ # if no mask provided to use conditional/unconditional labels then do not mask logits
+ do_not_mask_atom_logits = torch.ones(
+ (len(x["chemical_system"]), 1), dtype=torch.bool, device=x["num_atoms"].device
+ )
+
+ # mypy
+ assert batch_idx is not None, "batch_idx must be provided if condition is not None"
+ # Only mask atom types where the condition is not masked
+ # A 1 means that we do not alter the logit, a 0 means that we change the logit to -inf
+ # keep_logits.shape=(Nbatch, MAX_ATOMIC_NUM+1)
+
+ # 1 = keep logit, 0 = set logit to -inf, shape = (Nbatch, MAX_ATOMIC_NUM+1)
+ keep_all_logits = torch.ones((len(x["chemical_system"]), 1), device=x["num_atoms"].device)
+
+ # torch.Tensor, shape=(Nbatch,MAX_ATOMIC_NUM+1) -- 1s where elements are present in chemical system condition, 0 elsewhere
+ multi_hot_chemical_system = ChemicalSystemMultiHotEmbedding.sequences_to_multi_hot(
+ x=ChemicalSystemMultiHotEmbedding.convert_to_list_of_str(x=x["chemical_system"]),
+ device=x["num_atoms"].device,
+ )
+
+ keep_logits = torch.where(
+ do_not_mask_atom_logits,
+ keep_all_logits,
+ multi_hot_chemical_system,
+ )
+ # This is converting the 1-based chemical system condition to a 0-based
+ # condition -- we're doing it on the multi-hot representation of the
+ # chemical system, so we need to shift the indices by one.
+ if predictions_are_zero_based:
+ keep_logits = keep_logits[:, 1:]
+ # If we use mask diffusion, logits is shape [batch_size, MAX_ATOMIC_NUM + 1]
+ # instead of [batch_size, MAX_ATOMIC_NUM], so we have to add one dummy column
+ if keep_logits.shape[1] == logits.shape[1] - 1:
+ keep_logits = torch.cat([keep_logits, torch.zeros_like(keep_logits[:, :1])], dim=-1)
+ # Mask out all logits outside the chemical system we condition on
+ logits = mask_logits(logits, keep_logits[batch_idx])
+
+ return logits
+
+
+def get_chemgraph_from_denoiser_output(
+ pred_atom_types: torch.Tensor,
+ pred_lattice_eps: torch.Tensor,
+ pred_cart_pos_eps: torch.Tensor,
+ training: bool,
+ element_mask_func: Callable | None,
+ x_input: ChemGraph,
+) -> ChemGraph:
+ """
+ Convert raw denoiser output to ChemGraph and optionally apply masking to element logits.
+
+ Keyword arguments
+ -----------------
+ pred_atom_atoms: predicted logits for atom types
+ pred_lattice_eps: predicted lattice noise
+ pred_cart_pos_eps: predicted cartesian position noise
+ training: whether or not the model is in training mode - logit masking is only applied when sampling
+ element_mask_func: when not training, a function can be applied to mask logits for certain atom types
+ x_input: the nosiy state input to the score model, contains the lattice to convert cartesisan to fractional noise.
+ """
+ if not training and element_mask_func:
+ # when sampling we may want to mask logits for atom types depending on info in x['chemical_system'] and x['chemical_system_MASK']
+ pred_atom_types = element_mask_func(
+ logits=pred_atom_types,
+ x=x_input,
+ batch_idx=x_input.get_batch_idx("pos"),
+ )
+
+ replace_dict = dict(
+ # convert from cartesian to fractional coordinate score
+ pos=(
+ x_input["cell"].inverse().transpose(1, 2)[x_input.get_batch_idx("pos")]
+ @ pred_cart_pos_eps.unsqueeze(-1)
+ ).squeeze(-1),
+ cell=pred_lattice_eps,
+ atomic_numbers=pred_atom_types,
+ )
+ return x_input.replace(
+ **replace_dict,
+ )
+
+
+class GemNetTDenoiser(ScoreModel):
+ """Denoiser"""
+
+ def __init__(
+ self,
+ gemnet: nn.Module,
+ hidden_dim: int = 512,
+ denoise_atom_types: bool = True,
+ atom_type_diffusion: str = [
+ "mask",
+ "uniform",
+ ][0],
+ property_embeddings: torch.nn.ModuleDict | None = None,
+ property_embeddings_adapt: torch.nn.ModuleDict | None = None,
+ element_mask_func: Callable | None = None,
+ **kwargs,
+ ):
+ """Construct a GemNetTDenoiser object.
+
+ Args:
+ gemnet: a GNN module
+ hidden_dim (int, optional): Number of hidden dimensions in the GemNet. Defaults to 128.
+ denoise_atom_types (bool, optional): Whether to denoise the atom types. Defaults to False.
+ atom_type_diffusion (str, optional): Which type of atom type diffusion to use. Defaults to "mask".
+ condition_on (Optional[List[str]], optional): Which aspects of the data to condition on. Strings must be in ["property", "chemical_system"]. If None (default), condition on ["chemical_system"].
+ """
+ super(GemNetTDenoiser, self).__init__()
+
+ self.gemnet = gemnet
+ self.noise_level_encoding = NoiseLevelEncoding(hidden_dim)
+ self.hidden_dim = hidden_dim
+ self.denoise_atom_types = denoise_atom_types
+ self.atom_type_diffusion = atom_type_diffusion
+
+ # torch.nn.ModuleDict: Dict[PropertyName, PropertyEmbedding]
+ self.property_embeddings = torch.nn.ModuleDict(property_embeddings or {})
+
+ with_mask_type = self.denoise_atom_types and "mask" in self.atom_type_diffusion
+ self.fc_atom = nn.Linear(hidden_dim, MAX_ATOMIC_NUM + int(with_mask_type))
+
+ self.element_mask_func = element_mask_func
+
+ def forward(self, x: ChemGraph, t: torch.Tensor) -> ChemGraph:
+ """
+ args:
+ x: tuple containing:
+ frac_coords: (N_atoms, 3)
+ lattice: (N_cryst, 3, 3)
+ atom_types: (N_atoms, ), need to use atomic number e.g. H = 1 or ion state
+ num_atoms: (N_cryst,)
+ batch: (N_atoms,)
+ t: (N_cryst,): timestep per crystal
+ returns:
+ tuple of:
+ predicted epsilon: (N_atoms, 3)
+ lattice update: (N_crystals, 3, 3)
+ predicted atom types: (N_atoms, MAX_ATOMIC_NUM)
+ """
+ (frac_coords, lattice, atom_types, num_atoms, batch) = (
+ x["pos"],
+ x["cell"],
+ x["atomic_numbers"],
+ x["num_atoms"],
+ x.get_batch_idx("pos"),
+ )
+
+ # (num_atoms, hidden_dim) (num_crysts, 3)
+ t_enc = self.noise_level_encoding(t).to(lattice.device)
+ z_per_crystal = t_enc
+
+ # evaluate property embedding values
+ property_embedding_values = get_property_embeddings(
+ batch=x, property_embeddings=self.property_embeddings
+ )
+
+ if len(property_embedding_values) > 0:
+ z_per_crystal = torch.cat([z_per_crystal, property_embedding_values], dim=-1)
+
+ output = self.gemnet(
+ z=z_per_crystal,
+ frac_coords=frac_coords,
+ atom_types=atom_types,
+ num_atoms=num_atoms,
+ batch=batch,
+ lengths=None,
+ angles=None,
+ lattice=lattice,
+ # we construct the graph on the fly, hence pass None for these:
+ edge_index=None,
+ to_jimages=None,
+ num_bonds=None,
+ )
+ pred_atom_types = self.fc_atom(output.node_embeddings)
+
+ return get_chemgraph_from_denoiser_output(
+ pred_atom_types=pred_atom_types,
+ pred_lattice_eps=output.stress,
+ pred_cart_pos_eps=output.forces,
+ training=self.training,
+ element_mask_func=self.element_mask_func,
+ x_input=x,
+ )
+
+ @property
+ def cond_fields_model_was_trained_on(self) -> list[PropertySourceId]:
+ """
+ We adopt the convention that all property embeddings are stored in torch.nn.ModuleDicts of
+ name property_embeddings or property_embeddings_adapt in the case of a fine tuned model.
+
+ This function returns the list of all field names that a given score model was trained to
+ condition on.
+ """
+ return list(self.property_embeddings)
diff --git a/model/diffusion/__init__.py b/model/diffusion/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/model/diffusion/cgmanifest.json b/model/diffusion/cgmanifest.json
new file mode 100644
index 0000000000000000000000000000000000000000..8e4f2373b2e506eaca2d81d823382ad738ac1420
--- /dev/null
+++ b/model/diffusion/cgmanifest.json
@@ -0,0 +1,16 @@
+{
+ "$schema": "https://json.schemastore.org/component-detection-manifest.json",
+ "version": 1,
+ "registrations":[
+ {
+ "component": {
+ "type": "git",
+ "git": {
+ "repositoryUrl": "https://github.com/yang-song/score_sde_pytorch",
+ "commitHash": "cb1f359f4aadf0ff9a5e122fe8fffc9451fd6e44"
+ }
+ },
+ "developmentDependency" : false
+ }
+ ]
+}
\ No newline at end of file
diff --git a/model/diffusion/config.py b/model/diffusion/config.py
new file mode 100644
index 0000000000000000000000000000000000000000..453d88e91bdacfa823c42d28e325c5a1a2ed2f87
--- /dev/null
+++ b/model/diffusion/config.py
@@ -0,0 +1,33 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+from dataclasses import dataclass, field
+from typing import Any
+
+
+@dataclass
+class Config:
+
+ # This is for CLI applications that need to reuse a CLI parameter in multiple places
+ # in the config file. The idea is that you use `my_cli params.output_dir=foobar`
+ # and in other places in the config file `output_dir: ${params.output_dir}`
+ params: dict[str, Any] = field(default_factory=dict)
+
+ checkpoint_path: str | None = None # Required if train == False
+
+ # if load_original is True then we load original weights in validation mode instead of EMA
+ load_original: bool = False
+
+ # When auto_resume is set to `True` the trainer saves a copy of each checkpoint in
+ # {trainer.default_root_dir}/checkpoints. Before starting training, we look in this
+ # directory for a checkpoint from which to resume training.
+ auto_resume: bool = False
+
+ # DiffusionLightningModule
+ lightning_module: dict[str, Any] = field(default_factory=dict)
+
+ # pytorch_lightning.Trainer
+ trainer: dict[str, Any] = field(default_factory=dict)
+
+ # LightningDataModule
+ data_module: dict[str, Any] = field(default_factory=dict)
diff --git a/model/diffusion/corruption/__init__.py b/model/diffusion/corruption/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/model/diffusion/corruption/corruption.py b/model/diffusion/corruption/corruption.py
new file mode 100644
index 0000000000000000000000000000000000000000..04b6e74a2a6b8eecf65faf80c388b98b6d500c84
--- /dev/null
+++ b/model/diffusion/corruption/corruption.py
@@ -0,0 +1,118 @@
+"""
+Copyright 2020 The Google Research Authors.
+Copyright (c) Microsoft Corporation.
+Licensed under the MIT License.
+
+Based on code from https://github.com/yang-song/score_sde_pytorch/blob/main/sde_lib.py
+which is released under Apache licence.
+
+Abstract SDE classes, Reverse SDE, and VE/VP SDEs.
+
+Key changes:
+- Rename SDE => Corruption
+- Remove several methods like .reverse(), .discretize()
+"""
+
+import abc
+import logging
+from typing import Optional, Tuple, Union
+
+import torch
+
+from ...diffusion.data.batched_data import BatchedData
+
+B = Optional[torch.LongTensor]
+
+
+def _broadcast_like(x, like):
+ """
+ add broadcast dimensions to x so that it can be broadcast over ``like``
+ """
+ if like is None:
+ return x
+ return x[(...,) + (None,) * (like.ndim - x.ndim)]
+
+
+def maybe_expand(x: torch.Tensor, batch: B, like: torch.Tensor = None) -> torch.Tensor:
+ """
+
+ Args:
+ x: shape (batch_size, ...)
+ batch: shape (num_thingies,) with integer entries in the range [0, batch_size), indicating which sample each thingy belongs to
+ like: shape x.shape + potential additional dimensions
+ Returns:
+ expanded x with shape (num_thingies,), or if given like.shape, containing value of x for each thingy.
+ If `batch` is None, just returns `x` unmodified, to avoid pointless work if you have exactly one thingy per sample.
+ """
+ x = _broadcast_like(x, like)
+ if batch is None:
+ return x
+ else:
+ if x.shape[0] == batch.shape[0]:
+ logging.warning(
+ "Warning: batch shape is == x shape, are you trying to expand something that is already expanded?"
+ )
+ return x[batch]
+
+
+class Corruption(abc.ABC):
+ """Abstract base class for corruption processes"""
+
+ @property
+ @abc.abstractmethod
+ def T(self) -> float:
+ """End time of the corruption process."""
+ pass
+
+ @abc.abstractmethod
+ def marginal_prob(
+ self,
+ x: torch.Tensor,
+ t: torch.Tensor,
+ batch_idx: B = None,
+ batch: Optional[BatchedData] = None,
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
+ """Parameters to determine the marginal distribution of the SDE, $p_t(x)$."""
+ pass # mean: (num_nodes, num_features), std (num_nodes,)
+
+ @abc.abstractmethod
+ def prior_sampling(
+ self,
+ shape: Union[torch.Size, Tuple],
+ conditioning_data: Optional[BatchedData] = None,
+ batch_idx: B = None, # This is normally unused but is needed for special cases such as sample-wise zero-centering.
+ ) -> torch.Tensor:
+ """Generate one sample from the prior distribution, $p_T(x)$."""
+ pass
+
+ @abc.abstractmethod
+ def prior_logp(
+ self,
+ z: torch.Tensor,
+ batch_idx: B = None,
+ batch: Optional[BatchedData] = None,
+ ) -> torch.Tensor:
+ """Compute log-density of the prior distribution.
+
+ Useful for computing the log-likelihood via probability flow ODE.
+
+ Args:
+ z: latent code
+ Returns:
+ log probability density
+ """
+ pass # prior_logp: (batch_size,)
+
+ @abc.abstractmethod
+ def sample_marginal(
+ self,
+ x: torch.Tensor,
+ t: torch.Tensor,
+ batch_idx: B = None,
+ batch: Optional[BatchedData] = None,
+ ) -> torch.Tensor:
+ """Sample marginal for x(t) given x(0).
+ Returns:
+ sampled x(t) (same shape as input x).
+ """
+ pass
diff --git a/model/diffusion/corruption/d3pm_corruption.py b/model/diffusion/corruption/d3pm_corruption.py
new file mode 100644
index 0000000000000000000000000000000000000000..09a2c8d66c9c7b3ac1887513fb0407d96d4b99fb
--- /dev/null
+++ b/model/diffusion/corruption/d3pm_corruption.py
@@ -0,0 +1,108 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+from typing import Optional, Tuple, Union
+
+import torch
+from torch_scatter import scatter_add
+
+from ...diffusion.corruption.corruption import B, Corruption, maybe_expand
+from ...diffusion.d3pm import d3pm
+from ...diffusion.data.batched_data import BatchedData
+from ...diffusion.discrete_time import to_discrete_time
+
+
+class D3PMCorruption(Corruption):
+ """D3PM discrete corruption process. Has discret time and discrete (categorical) values."""
+
+ def __init__(
+ self,
+ d3pm: d3pm.DiscreteDiffusionBase,
+ offset: int = 0,
+ ):
+ super().__init__()
+ self.d3pm = d3pm
+ # Often, the data is not zero-indexed, so we need to offset the data
+ # E.g., if we are dealing with one-based class labels, we might want to offset by 1 to convert from zero-based indices to actual classes.
+ self.offset = offset
+
+ @property
+ def N(self) -> int:
+ """Number of diffusion timesteps i.e. number of noise levels.
+ Must match number of noise levels used for sampling. To change this, we'd need to implement continuous-time diffusion for discrete things
+ as in e.g. Campbell et al. https://arxiv.org/abs/2205.14987"""
+ return self.d3pm.num_steps
+
+ def _to_zero_based(self, x: torch.Tensor) -> torch.Tensor:
+ """Convert from non-zero-based indices to zero-based indices."""
+ return x - self.offset
+
+ def _to_non_zero_based(self, x: torch.Tensor) -> torch.Tensor:
+ """Convert from zero-based indices to non-zero-based indices."""
+ return x + self.offset
+
+ @property
+ def T(self) -> float:
+ """End time of the Corruption process."""
+ return 1
+
+ def marginal_prob(
+ self,
+ x: torch.Tensor,
+ t: torch.Tensor,
+ batch_idx: B = None,
+ batch: Optional[BatchedData] = None,
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
+ """Parameters to determine the marginal distribution of the corruption process, $p_t(x | x_0)$."""
+ # plus 1 because t=0 is actually no corruption for D3PM and it has N corruption steps, i.e., values go from 0 to N.
+ t_discrete = maybe_expand(to_discrete_time(t, N=self.N, T=self.T), batch_idx) + 1
+ _, logits = d3pm.q_sample(
+ self._to_zero_based(x.long()), t_discrete, diffusion=self.d3pm, return_logits=True
+ )
+ return logits, None # mean: (nodes_per_sample * batch_size, ), std None
+
+ def prior_sampling(
+ self,
+ shape: Union[torch.Size, Tuple],
+ conditioning_data: Optional[BatchedData] = None,
+ batch_idx: B = None,
+ ) -> torch.Tensor:
+ """Generate one sample from the prior distribution, $p_T(x)$."""
+ # sample and then add offset to convert to non-zero-based class labels
+ return self._to_non_zero_based(self.d3pm.sample_stationary(shape))
+
+ def prior_logp(
+ self,
+ z: torch.Tensor,
+ batch_idx: B = None,
+ batch: Optional[BatchedData] = None,
+ ) -> torch.Tensor:
+ """Compute log-density of the prior distribution.
+
+ Args:
+ z: samples, non-zero-based indices, i.e., we first need to subtract the offset
+ Returns:
+ log probability density
+ """
+ probs = self.d3pm.stationary_probs(z.shape).to(z.device)
+ log_probs = (probs + 1e-8).log()
+ log_prob_per_sample = log_probs[:, self._to_zero_based(z.long())]
+ log_prob_per_structure = scatter_add(log_prob_per_sample, batch_idx, dim=0)
+ return log_prob_per_structure
+
+ def sample_marginal(
+ self,
+ x: torch.Tensor,
+ t: torch.Tensor,
+ batch_idx: B = None,
+ batch: Optional[BatchedData] = None,
+ ) -> torch.Tensor:
+ """Sample marginal for x(t) given x(0).
+ Returns:
+ sampled x(t), non-zero-based indices
+ where raw_noise is drawn from standard Gaussian
+ """
+ logits = self.marginal_prob(x=x, t=t, batch_idx=batch_idx, batch=batch)[0]
+ sample = torch.distributions.Categorical(logits=logits).sample()
+ # samples are zero-based, so we need to add the offset to convert to non-zero-based class labels.
+ return self._to_non_zero_based(sample)
diff --git a/model/diffusion/corruption/multi_corruption.py b/model/diffusion/corruption/multi_corruption.py
new file mode 100644
index 0000000000000000000000000000000000000000..811353b972c39219cb8b4fdb7aa5af395e7a8e95
--- /dev/null
+++ b/model/diffusion/corruption/multi_corruption.py
@@ -0,0 +1,170 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from functools import cached_property
+from typing import (
+ Any,
+ Callable,
+ Dict,
+ Generic,
+ Iterable,
+ List,
+ Mapping,
+ Optional,
+ Tuple,
+ TypeVar,
+)
+
+import torch
+
+from ...diffusion.corruption.d3pm_corruption import D3PMCorruption
+from ...diffusion.corruption.sde_lib import SDE, Corruption
+from ...diffusion.data.batched_data import BatchedData
+
+R = TypeVar("R")
+Diffusable = TypeVar("Diffusable", bound=BatchedData)
+
+
+def _first(s: Iterable):
+ return next(iter(s))
+
+
+@dataclass
+class MultiCorruptionConfig:
+ discrete_corruptions: dict[str, Any] = field(default_factory=dict)
+ sdes: dict[str, Any] = field(default_factory=dict)
+
+
+class MultiCorruption(Generic[Diffusable]):
+ """Wraps multiple `Corruption` instances to operate on different fields of a State
+
+ In the forward process, each field of State is corrupted independently.
+
+ In the reverse process, a single score model takes in the entire State and
+ uses it to estimate the score with respect to each field of the State.
+ """
+
+ def _get_batch_indices(self, batch: Diffusable) -> Dict[str, torch.Tensor]:
+ return {k: batch.get_batch_idx(k) for k in self.corrupted_fields}
+
+ def __init__(
+ self,
+ sdes: Optional[Mapping[str, SDE]] = None,
+ discrete_corruptions: Optional[Mapping[str, D3PMCorruption]] = None,
+ ):
+ """
+ Args:
+ sdes: mapping from fields of batch to SDE corruption processes
+ discrete_corruptions: mapping from fields of batch to discrete corruption processes
+ """
+ if sdes is None:
+ sdes = {}
+ if discrete_corruptions is None:
+ discrete_corruptions = {}
+ assert (
+ len(sdes) + len(discrete_corruptions) > 0
+ ), "Must have at least one corruption process."
+
+ self._sdes = sdes
+ self._discrete_corruptions = discrete_corruptions
+ assert (
+ set(self._sdes.keys()).intersection(set(self._discrete_corruptions.keys())) == set()
+ ), "SDEs and corruptions have overlapping keys."
+ self._corruptions: Dict[str, Corruption] = {**self._sdes, **self._discrete_corruptions}
+ # Make the dict sorted by key (to prevent mismatching checkpoints):
+ self._corruptions = {k: self._corruptions[k] for k in sorted(self._corruptions.keys())}
+
+ # All SDEs must have the same T
+ T_vals = [corruption.T for corruption in self.corruptions.values()]
+ assert len(set(T_vals)) == 1
+
+ @property
+ def sdes(self) -> Mapping[str, SDE]:
+ return self._sdes
+
+ @property
+ def has_discrete_corruptions(self) -> bool:
+ return len(self.discrete_corruptions) > 0
+
+ @property
+ def discrete_corruptions(self) -> Mapping[str, Corruption]:
+ return self._discrete_corruptions
+
+ @property
+ def corruptions(self) -> Mapping[str, Corruption]:
+ return self._corruptions
+
+ @property
+ def corrupted_fields(self) -> List[str]:
+ return list(self.corruptions.keys())
+
+ @cached_property
+ def T(self) -> float:
+ return _first(self.corruptions.values()).T
+
+ def sample_marginal(self, batch: Diffusable, t) -> Diffusable:
+ def fn_getter(corruption: Corruption) -> Callable[..., Tuple[torch.Tensor, torch.Tensor]]:
+ return corruption.sample_marginal
+
+ noisy_data = self._apply_corruption_fn(
+ fn_getter,
+ x=batch,
+ batch_idx=self._get_batch_indices(batch),
+ broadcast=dict(t=t),
+ )
+ noisy_batch = batch.replace(**noisy_data)
+ return noisy_batch
+
+ def sde(
+ self, batch: Diffusable, t: torch.Tensor
+ ) -> Dict[str, Tuple[torch.Tensor, torch.Tensor]]:
+ """Get drift and diffusion for each component of the state"""
+ assert (
+ not self.has_discrete_corruptions
+ ), "Cannot call `sde` on a MultiCorruption with non-SDE corruptions"
+
+ fns = {k: sde.sde for k, sde in self.sdes.items()}
+ return apply(
+ fns=fns,
+ broadcast={"batch": batch, "t": t},
+ x=batch,
+ batch_idx=self._get_batch_indices(batch),
+ )
+
+ def _apply_corruption_fn(
+ self,
+ fn_getter: Callable[[Corruption], Callable[..., R]],
+ x: BatchedData,
+ batch_idx: Mapping[str, torch.LongTensor],
+ broadcast: Optional[Dict] = None,
+ apply_to: Optional[Mapping[str, Corruption]] = None,
+ **kwargs,
+ ) -> Dict[str, R]:
+ if apply_to is None:
+ apply_to = self.corruptions
+ fns = {field_name: fn_getter(corruption) for field_name, corruption in apply_to.items()}
+ return apply(
+ fns=fns,
+ broadcast={**(broadcast or dict()), "batch": x},
+ x=x,
+ batch_idx=batch_idx,
+ **kwargs,
+ )
+
+
+def apply(fns: Dict[str, Callable[..., R]], broadcast, **kwargs) -> Dict[str, R]:
+ """Apply different function with different argument values to each field.
+ fns: dict of the form {field_name: function_to_apply}
+ broadcast: arguments that are identical for every field_name
+ kwargs: dict of the form {argument_name: {field_name: argument_value}}
+ """
+ return {
+ field_name: fn(
+ **{k: v[field_name] for k, v in kwargs.items() if field_name in v},
+ **(broadcast or dict()),
+ )
+ for field_name, fn in fns.items()
+ }
diff --git a/model/diffusion/corruption/sde_lib.py b/model/diffusion/corruption/sde_lib.py
new file mode 100644
index 0000000000000000000000000000000000000000..af47d8dadd5d297d623053b259700e798b67ffbe
--- /dev/null
+++ b/model/diffusion/corruption/sde_lib.py
@@ -0,0 +1,270 @@
+"""
+Copyright 2020 The Google Research Authors.
+Copyright (c) Microsoft Corporation.
+Licensed under the MIT License.
+
+Based on code from https://github.com/yang-song/score_sde_pytorch
+which is released under Apache licence.
+
+Abstract SDE classes, Reverse SDE, and VE/VP SDEs.
+
+Key changes:
+- Adapted to work on batched pytorch_geometric style data
+- Added '...given_score' methods so that score for a composite
+state can be calculated in single forward pass of a shared score model,
+and the scores for different fields then forwarded to the different reverse SDEs.
+"""
+
+import abc
+from typing import Callable, Optional, Protocol, Tuple, Union
+
+import numpy as np
+import torch
+from torch_scatter import scatter_add
+
+from ...diffusion.corruption.corruption import B, Corruption, maybe_expand
+from ...diffusion.data.batched_data import BatchedData
+
+
+class ScoreFunction(Protocol):
+ def __call__(
+ self,
+ x: torch.Tensor,
+ t: torch.Tensor,
+ batch_idx: B = None,
+ ) -> torch.Tensor:
+ """Calculate score.
+
+ Args:
+ x: Samples at which the score should be calculated. Shape [num_nodes, ...]
+ t: Timestep for each sample. Shape [num_samples,]
+ batch_idx: Indicates which sample each row of x belongs to. Shape [num_nodes,]
+
+ """
+ pass
+
+
+class SDE(Corruption):
+ """Corruption using a stochastic differential equation."""
+
+ @abc.abstractmethod
+ def sde(
+ self,
+ x: torch.Tensor,
+ t: torch.Tensor,
+ batch_idx: B = None,
+ batch: Optional[BatchedData] = None,
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
+ """Returns drift f and diffusion coefficient g such that dx = f * dt + g * sqrt(dt) * standard Gaussian"""
+ pass # drift: (nodes_per_sample * batch_size, num_features), diffusion (batch_size,)
+
+ @abc.abstractmethod
+ def marginal_prob(
+ self,
+ x: torch.Tensor,
+ t: torch.Tensor,
+ batch_idx: B = None,
+ batch: Optional[BatchedData] = None,
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
+ """Returns mean and standard deviation of the marginal distribution of the SDE, $p_t(x)$."""
+ pass # mean: (nodes_per_sample * batch_size, num_features), std: (nodes_per_sample * batch_size, 1)
+
+ def mean_coeff_and_std(
+ self,
+ x: torch.Tensor,
+ t: torch.Tensor,
+ batch_idx: B = None,
+ batch: Optional[BatchedData] = None,
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
+ """Returns mean coefficient and standard deviation of marginal distribution at time t."""
+ return self.marginal_prob(
+ torch.ones_like(x), t, batch_idx, batch
+ ) # mean_coeff: same shape as x, std: same shape as x
+
+ def sample_marginal(
+ self,
+ x: torch.Tensor,
+ t: torch.Tensor,
+ batch_idx: B = None,
+ batch: Optional[BatchedData] = None,
+ ) -> torch.Tensor:
+ """Sample marginal for x(t) given x(0).
+ Returns:
+ sampled x(t)
+ """
+ mean, std = self.marginal_prob(x=x, t=t, batch_idx=batch_idx, batch=batch)
+ z = torch.randn_like(x)
+
+ return mean + std * z
+
+
+class BaseVPSDE(SDE):
+ """Base class for variance-preserving SDEs of the form
+ dx = - 0.5 * beta_t * x * dt + sqrt(beta_t) * z * sqrt(dt)
+ where z is unit Gaussian noise, or equivalently
+ dx = - 0.5 * beta_t *x * dt + sqrt(beta_t) * dW
+
+ """
+
+ @abc.abstractmethod
+ def beta(self, t: torch.Tensor) -> torch.Tensor: ...
+
+ @abc.abstractmethod
+ def _marginal_mean_coeff(self, t: torch.Tensor) -> torch.Tensor:
+ """This should be implemented to compute exp(-0.5 * int_0^t beta(s) ds). See equation (29) of Song et al."""
+ ...
+
+ @property
+ def T(self) -> float:
+ return 1.0
+
+ def marginal_prob(
+ self,
+ x: torch.Tensor,
+ t: torch.Tensor,
+ batch_idx: B = None,
+ batch: Optional[BatchedData] = None,
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
+ mean_coeff = self._marginal_mean_coeff(t)
+ mean = maybe_expand(mean_coeff, batch_idx, x) * x
+ std = maybe_expand(torch.sqrt(1.0 - mean_coeff**2), batch_idx, x)
+ return mean, std
+
+ def prior_sampling(
+ self,
+ shape: Union[torch.Size, Tuple],
+ conditioning_data: Optional[BatchedData] = None,
+ batch_idx: B = None,
+ ) -> torch.Tensor:
+ return torch.randn(*shape)
+
+ def prior_logp(
+ self,
+ z: torch.Tensor,
+ batch_idx: B = None,
+ batch: Optional[BatchedData] = None,
+ ) -> torch.Tensor:
+ return unit_gaussian_logp(z, batch_idx)
+
+ def sde(
+ self,
+ x: torch.Tensor,
+ t: torch.Tensor,
+ batch_idx: B = None,
+ batch: Optional[BatchedData] = None,
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
+ beta_t = self.beta(t)
+ drift = -0.5 * maybe_expand(beta_t, batch_idx, x) * x
+ diffusion = maybe_expand(torch.sqrt(beta_t), batch_idx, x)
+ return drift, diffusion
+
+
+class VPSDE(BaseVPSDE):
+ def __init__(self, beta_min: float = 0.1, beta_max: float = 20):
+ """Variance-preserving SDE with drift coefficient changing linearly over time."""
+ super().__init__()
+ self.beta_0 = beta_min
+ self.beta_1 = beta_max
+
+ def beta(self, t: torch.Tensor) -> torch.Tensor:
+ return self.beta_0 + t * (self.beta_1 - self.beta_0)
+
+ def _marginal_mean_coeff(self, t: torch.Tensor) -> torch.Tensor:
+ log_mean_coeff = -0.25 * t**2 * (self.beta_1 - self.beta_0) - 0.5 * t * self.beta_0
+ return torch.exp(log_mean_coeff)
+
+
+def unit_gaussian_logp(z: torch.Tensor, batch_idx: B = None) -> torch.Tensor:
+ shape = z.shape
+ N = np.prod(shape[1:])
+ if batch_idx is None:
+ logps = -N / 2.0 * np.log(2 * np.pi) - torch.sum(z**2, dim=tuple(range(1, z.ndim))) / 2.0
+ else:
+ if z.ndim > 2:
+ raise NotImplementedError
+
+ logps = -N / 2.0 * np.log(2 * np.pi) - scatter_add(torch.sum(z**2, dim=1), batch_idx) / 2.0
+
+ return logps
+
+
+class VESDE(SDE):
+ def __init__(self, sigma_min: float = 0.01, sigma_max: float = 50.0):
+ """Construct a Variance Exploding SDE.
+
+ The marginal standard deviation grows exponentially from sigma_min to sigma_max.
+
+ Args:
+ sigma_min: smallest sigma.
+ sigma_max: largest sigma.
+ """
+ super().__init__()
+ self.sigma_min = sigma_min
+ self.sigma_max = sigma_max
+
+ @property
+ def T(self) -> float:
+ return 1.0
+
+ def sde(
+ self,
+ x: torch.Tensor,
+ t: torch.Tensor,
+ batch_idx: B = None,
+ batch: Optional[BatchedData] = None,
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
+ sigma = self.sigma_min * (self.sigma_max / self.sigma_min) ** t
+ drift = torch.zeros_like(x)
+ diffusion = maybe_expand(
+ sigma
+ * torch.sqrt(
+ torch.tensor(2 * (np.log(self.sigma_max) - np.log(self.sigma_min)), device=t.device)
+ ),
+ batch_idx,
+ x,
+ )
+ return drift, diffusion
+
+ def marginal_prob(
+ self,
+ x: torch.Tensor,
+ t: torch.Tensor,
+ batch_idx: B = None,
+ batch: Optional[BatchedData] = None,
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
+ std = maybe_expand(self.sigma_min * (self.sigma_max / self.sigma_min) ** t, batch_idx, x)
+ mean = x
+ return mean, std
+
+ def prior_sampling(
+ self,
+ shape: Union[torch.Size, Tuple],
+ conditioning_data: Optional[BatchedData] = None,
+ batch_idx: B = None,
+ ) -> torch.Tensor:
+ return torch.randn(*shape) * self.sigma_max
+
+ def prior_logp(
+ self,
+ z: torch.Tensor,
+ batch_idx: B = None,
+ batch: Optional[BatchedData] = None,
+ ) -> torch.Tensor:
+ shape = z.shape
+ N = np.prod(shape[1:])
+ if batch_idx is not None:
+ return -N / 2.0 * np.log(2 * np.pi * self.sigma_max**2) - scatter_add(
+ torch.sum(z**2, dim=1), batch_idx
+ ) / (2 * self.sigma_max**2)
+ else:
+ return -N / 2.0 * np.log(2 * np.pi * self.sigma_max**2) - torch.sum(
+ z**2, dim=tuple(range(1, z.ndim))
+ ) / (2 * self.sigma_max**2)
+
+
+def check_score_fn_defined(score_fn: Optional[Callable], fn_name_given_score: str):
+ """Check that a reverse SDE has a score_fn. Give a useful error message if not."""
+ if score_fn is None:
+ raise ValueError(
+ f"This reverse SDE does not know its score_fn. You must either a) pass a score_fn when you construct this reverse SDE or b) call {fn_name_given_score} instead."
+ )
diff --git a/model/diffusion/d3pm/__init__.py b/model/diffusion/d3pm/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/model/diffusion/d3pm/cgmanifest.json b/model/diffusion/d3pm/cgmanifest.json
new file mode 100644
index 0000000000000000000000000000000000000000..3a1e925996d9ba8deb3b75027f1799907edc1269
--- /dev/null
+++ b/model/diffusion/d3pm/cgmanifest.json
@@ -0,0 +1,16 @@
+{
+ "$schema": "https://json.schemastore.org/component-detection-manifest.json",
+ "version": 1,
+ "registrations":[
+ {
+ "component": {
+ "type": "git",
+ "git": {
+ "repositoryUrl": "https://github.com/google-research/google-research",
+ "commitHash": "ad2d81983e4c717f477a232f625d0da2808b15aa"
+ }
+ },
+ "developmentDependency" : false
+ }
+ ]
+}
\ No newline at end of file
diff --git a/model/diffusion/d3pm/d3pm.py b/model/diffusion/d3pm/d3pm.py
new file mode 100644
index 0000000000000000000000000000000000000000..e94e226427791479c8360503144f95b626bd2810
--- /dev/null
+++ b/model/diffusion/d3pm/d3pm.py
@@ -0,0 +1,872 @@
+# Copyright (c) 2022 The Google Research Authors.
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+# Adapted from https://github.com/google-research/google-research/blob/master/d3pm/text/diffusion.py
+# Siginificant changes:
+# * adapt code style/ formatting
+# * Jax -> PyTorch
+# * Remove Diffusion types that are not used by MatterGen
+# ORIGINAL LICENSE NOTICE:
+# Copyright 2022 The Google Research Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Diffusions for training and noise scheduling."""
+
+import abc
+import dataclasses
+from typing import Any, Callable, Dict, Optional, Union
+
+import torch
+import torch.nn.functional as F
+from torch.distributions import Categorical
+
+
+class DiffusionSchedule:
+ """A wrapper around a simple schedule function."""
+
+ def __init__(self, schedule_fn, num_steps, is_constant=False):
+ self._schedule_fn = schedule_fn
+ self.num_steps = num_steps
+ self.is_constant = is_constant
+
+ def __call__(self, step):
+ return self._schedule_fn(step)
+
+ def __repr__(self):
+ return f"DiffusionSchedule(steps: {self.num_steps}, is_constant: {self.is_constant})"
+
+
+class DiscreteDiffusionBase(abc.ABC):
+ """Base class for all matrix-noise schedules."""
+
+ num_steps: int
+ dim: int
+ precision: Any = torch.float32
+
+ @abc.abstractmethod
+ def stationary_probs(self, shape):
+ """Returns probs for the stationary distribution."""
+
+ @abc.abstractmethod
+ def sample_stationary(self, shape):
+ """Draws a sample from the stationary distribution (q(x_T))."""
+
+ @property
+ def has_state(self):
+ """Indicates if the diffusion has state which needs to be set/updated."""
+ return False
+
+ def set_state(self, state):
+ pass
+
+ def reset_state(self):
+ pass
+
+ def update_state(self, state):
+ pass
+
+ def sample_t(self, shape=(1,)):
+ """Samples batches of time steps to use."""
+
+ num_steps = self.num_steps
+ t = torch.randint(shape, minval=0, maxval=num_steps)
+ return t
+
+ @abc.abstractmethod
+ def get_qt_given_q0(self, q0, t, return_logits=False, make_one_hot=False, epsilon=1e-20):
+ """Get q(x_t), the n-step posterior.
+
+ For example, for t = 0, it returns q0 unchanged.
+
+ Args:
+ q0: an array of floats specifying a distribution over p(x_0).
+ t: t in q(x_t | x_0).
+ return_logits: if True, return the output logits
+ make_one_hot: if True, will convert q0 to floats if needed.
+ epsilon: a small number to normalize logits conversion with, if needed.
+
+ Returns:
+ q(x_t | x_0).
+ """
+
+ @abc.abstractmethod
+ def sample_and_compute_posterior_q(
+ self,
+ x_0,
+ t,
+ samples=None,
+ transition_probs=None,
+ return_logits=True,
+ return_transition_probs=False,
+ transition_probs_in_logits=True,
+ make_one_hot=True,
+ epsilon=1e-20,
+ step_size=1,
+ ):
+ """Samples from q(x_{t+1} | x_0), then computes q(x_t | x_{t+1}, x_0).
+
+ Args:
+ x_0: an array containing x_0 samples. These are expected to be integral
+ unless make_one_hot is False (in which case probabilities can be
+ provided).
+ t: the timestep to compute (as an int or integer array with shape that
+ matches x_0.
+ samples: if not None, use these samples to compute the posterior.
+ transition_probs: precomputed transition probabilities.
+ return_logits: if True, returns the (noisy) log of the probabilities.
+ return_transition_probs: if true, returns the transition probs as well.
+ transition_probs_in_logits: include transition probs in logits.
+ make_one_hot: if True, will convert the input to a one_hot vector.
+ epsilon: a small amount of noise to add to logits if needed.
+ step_size: if provided, computes q(x_{t + step_size} | x_0), etc. This is
+ used to sample fewer steps for ELBO evaluation on a longer trained
+ model.
+
+ Returns:
+ a list of samples with the same shape as x_0 and the associated posterior
+ probabilities (or logits).
+ """
+
+
+class DiscreteDiffusionMatrixBase(DiscreteDiffusionBase):
+ """Base class for all matrix-noise schedulers."""
+
+ num_steps: int
+ dim: int
+ precision: Any = torch.float32
+
+ def get(self, t):
+ """Returns the transition matrix q(x_{t+1} | x_t)."""
+ raise NotImplementedError
+
+ def custom_product_fn(self, t):
+ """Returns q(x_t | x_0), the product of the first t matrices."""
+ raise NotImplementedError
+
+ def supports_efficient_get(self):
+ """Returns true if get() is implemented/efficient."""
+ return False
+
+ def supports_efficient_inference(self):
+ """Returns true if custom_product_fn is implemented.
+
+ The ontology of efficient_get and efficient_inference is this:
+ * if efficient_inference is enabled, it is used to return q(x_t | x_0)
+ without computing expensive products.
+ * if efficient_get is enabled, get(...) is used to get the posterior of
+ q(x_{t-1} | x_t, x_0). If not, get_q_given_q0 is called to get
+ q(x_{t+1} | x_0), and qt_reverse is called to get the q(x_{t+1} | x_t).
+ """
+ return False
+
+ def qt_reverse(self, qt_plus_1, t, return_logits=False, make_one_hot=False, epsilon=1e-20):
+ """Get q(x_{t+1} | x_t), for each possible value of x_t. Thus, the rows of the output do not sum to 1.
+
+ Args:
+ qt_plus_1: an array of floats specifying a distribution over q(x_{t+1} | x_0).
+ t: t in q(x_{t+1} | x_t).
+ return_logits: if True, return the output logits
+ make_one_hot: if True, will convert q(x_{t+1}) to floats if needed.
+ epsilon: a small number to normalize logits conversion with, if needed.
+
+ Returns:
+ q(x_{t+1} | x_t), shape [num_samples, num_classes].
+ """
+ raise NotImplementedError
+
+ def get_qt_matrix(self, t):
+ """Returns the matrix Q = q(x_t | x_0) materialized over all x_0."""
+ if self.supports_efficient_inference():
+ return self.custom_product_fn(t)
+
+ # otherwise, multiply by the ith matrix in a for-loop.
+ def product_fn(i, state):
+ return torch.matmul(self.get(torch.tensor(i)), state)
+
+ val = torch.eye(self.dim, device=t.device)
+ for i in range(0, t):
+ val = product_fn(i, val)
+ return val
+
+ def get_qt_given_q0(self, q0, t, return_logits=False, make_one_hot=False, epsilon=1e-20):
+ """Get q(x_t), the n-step posterior.
+
+ For example, for t = 0, it returns q0 unchanged.
+
+ Args:
+ q0: an array of floats specifying a distribution over p(x_0).
+ t: t in q(x_t | x_0).
+ return_logits: if True, return the output logits
+ make_one_hot: if True, will convert q0 to floats if needed.
+ epsilon: a small number to normalize logits conversion with, if needed.
+
+ Returns:
+ q(x_t | x_0).
+ """
+
+ if make_one_hot:
+ assert q0.dtype == torch.long or q0.dtype == torch.int32
+ q0 = torch.eye(self.dim, device=q0.device)[q0]
+
+ assert q0.dtype == torch.float32
+
+ # if efficient inference is supported, just return those matrices.
+ if self.supports_efficient_inference():
+ prob_at_time_t = torch.einsum("bij,bj->bi", self.get_qt_matrix(t).to(q0.dtype), q0)
+
+ if return_logits:
+ return torch.log(prob_at_time_t + epsilon)
+ else:
+ return prob_at_time_t
+
+ @dataclasses.dataclass
+ class ScanState:
+ final_time: int # target time
+ q: Any
+
+ def product_fn(state, current_time):
+ cond = current_time < state.final_time
+ transition = self.get(current_time)
+ q_t_plus_1 = torch.einsum("ij,sj->si", transition, state.q)
+
+ new_q = torch.where(cond[:, None], q_t_plus_1, state.q)
+ return ScanState(final_time=state.final_time, q=new_q), None
+
+ init_val = ScanState(final_time=t, q=q0)
+ carry = init_val
+ idx = torch.arange(self.num_steps, device=q0.device)
+ for i in idx:
+ carry, _ = product_fn(carry, i)
+ final_state = carry
+ prob_at_time_t = final_state.q
+
+ if return_logits:
+ return torch.log(prob_at_time_t + epsilon)
+ else:
+ return prob_at_time_t
+
+ def sample_and_compute_posterior_q(
+ self,
+ x_0,
+ t,
+ samples=None,
+ transition_probs=None,
+ return_logits=True,
+ return_transition_probs=False,
+ transition_probs_in_logits=True,
+ make_one_hot=True,
+ epsilon=1e-20,
+ step_size=1,
+ ):
+ """Samples from q(x_{t+1} | x_0), then computes q(x_t | x_{t+1}, x_0).
+
+ Args:
+ x_0: an array containing x_0 samples. These are expected to be integral
+ unless make_one_hot is False (in which case probabilities can be
+ provided).
+ t: the timestep to compute (as an int or integer array with shape that
+ matches x_0.
+ samples: if not None, use these samples to compute the posterior.
+ transition_probs: precomputed transition probabilities.
+ return_logits: if True, returns the (noisy) log of the probabilities.
+ return_transition_probs: if true, returns the transition probs as well.
+ transition_probs_in_logits: include transition probs in logits.
+ make_one_hot: if True, will convert the input to a one_hot vector.
+ epsilon: a small amount of noise to add to logits if needed.
+ step_size: if provided, computes q(x_{t + step_size} | x_0), etc. This is
+ used to sample fewer steps for ELBO evaluation on a longer trained
+ model.
+
+ Returns:
+ a list of samples with the same shape as x_0 and the associated posterior
+ probabilities (or logits).
+ """
+
+ dim = self.dim
+ device = x_0.device
+ # t = torch.tensor(t, device=x_0.device)
+ if make_one_hot:
+ assert x_0.dtype in [torch.long, torch.int32]
+ x_0 = torch.eye(dim, device=device)[x_0].reshape(x_0.shape + (dim,))
+ assert x_0.dtype == torch.float32
+ assert t.dtype in [torch.long, torch.int32]
+ prob_at_time_t = self.get_qt_given_q0(q0=x_0, t=t)
+ # most methods support efficiently returning the t-th transition matrix
+ # if so, we use that. Otherwise we recompute the t+1th probability.
+ if self.supports_efficient_get():
+ if step_size > 1:
+ transition_matrix = torch.eye(self.dim, device=x_0.device)
+
+ for i in range(step_size):
+ transition_matrix = self.get(t + i) @ transition_matrix
+
+ else:
+ transition_matrix = self.get(t)
+
+ prob_at_time_t_plus_one = torch.einsum(
+ "bij,bj->bi",
+ transition_matrix,
+ prob_at_time_t,
+ )
+
+ else:
+ prob_at_time_t_plus_one = self.get_qt_given_q0(q0=x_0, t=t + step_size)
+
+ if samples is None and transition_probs is not None:
+ raise ValueError("samples were not provided but transition_probs were.")
+
+ # if samples are provided, we use those. otherwise, we sample more.
+ if samples is None:
+ logits = torch.log(prob_at_time_t_plus_one + epsilon)
+ samples = Categorical(logits=logits).sample()
+
+ # we can optionally provide transition probs from another call to this
+ # function. If not, we recompute this. For most methods, we can reuse the
+ # transition matrix. If we didn't compute it, our method must support
+ # qt_reverse which usually computes efficient backwards VJPs.
+
+ if transition_probs is None:
+ if self.supports_efficient_get():
+ transition_probs = transition_matrix[range(samples.shape[0]), samples]
+ else:
+ if step_size > 1:
+ transition_probs = torch.eye(self.dim, device=samples.device)[samples]
+ for i in range(step_size):
+ transition_probs = self.qt_reverse(
+ qt_plus_1=transition_probs, make_one_hot=False, t=t + step_size - 1 - i
+ )
+ else:
+ # Computes q(x_{t+1} | x_t), i.e., for each possible x_t, what is the probability of transitioning to each x_{t+1}.
+ # Thus, these probabilities do not sum to 1 per row.
+ # If we don't return logits, transition_probs will be used to compute q(x_t | x_{t+1}).
+ # Otherwise, we return the logits of q(x_t | x_{t+1}) = q(x_{t+1} | x_t) * q(x_t | x_0), i.e., omit normalization by q(x_{t+1} | x_0).
+ # Shape [batch_size, num_classes]
+ transition_probs = self.qt_reverse(qt_plus_1=samples, make_one_hot=True, t=t)
+
+ if not transition_probs_in_logits and not return_logits:
+ raise ValueError(
+ "Cannot exclude transition probs from logits if return_logits is false."
+ )
+
+ if return_logits:
+ # for numerical stability, we can compute log(a*b) = log(a) + log(b)
+ posterior_logits = torch.log(prob_at_time_t + epsilon)
+
+ if transition_probs_in_logits:
+ posterior_logits += torch.log(transition_probs + epsilon)
+
+ if return_transition_probs:
+ return posterior_logits, samples, transition_probs
+ else:
+ return posterior_logits, samples
+ else:
+ # here we hope this never actually sums to zero. There's a chance
+ # this will produce NaN gradients, but that's OK because they'll be
+ # skipped.
+ posterior = transition_probs * prob_at_time_t
+ denominator = torch.sum(posterior, dim=-1, keepdims=True)
+ posterior = posterior / denominator
+
+ if return_transition_probs:
+ return posterior, samples, transition_probs
+ else:
+ return posterior, samples
+
+
+class MaskDiffusion(DiscreteDiffusionMatrixBase):
+ """A simple schedule that diffuses away from the identity matrix."""
+
+ def __init__(self, dim, schedule, precision=torch.float32, use_fast_inference=True):
+ """A simple scheduler for masking policies.
+
+ Args:
+ dim: int, the dimensionality of the state space.
+ schedule: a DiffusionSchedule object for scheduling rates.
+ precision: matmul precision.
+ use_fast_inference: if False, uses a slower, brute force approach.
+ """
+
+ self.num_steps = schedule.num_steps
+ self.schedule = schedule
+ self.use_fast_inference = use_fast_inference
+ self.precision = precision
+ self.dim = dim # allow mask
+ self.state = self._create_state()
+
+ def _create_state(self):
+ """Initializes values used by the get function."""
+ betas = torch.cat([torch.tensor([0.0]), self.schedule(torch.arange(self.num_steps))]).to(
+ torch.float64
+ )
+ alphas = 1 - betas
+ state = torch.cumprod(alphas, dim=0)
+ state[-1] = 0.0
+
+ return state.float()
+
+ def supports_efficient_inference(self):
+ return self.use_fast_inference
+
+ def stationary_probs(self, shape):
+ """Stationary distribution is one-hot at mask token."""
+ sample = torch.full(shape, self.dim - 1)
+ probs = torch.eye(self.dim, device=sample.device)[sample]
+ return probs
+
+ def sample_stationary(self, shape):
+ """Stationary distribution is one-hot at mask token."""
+ return torch.full(shape, self.dim - 1)
+
+ def custom_product_fn(self, t):
+ """Returns product of first n matrices. Only supported for beta constant."""
+ dim = self.dim
+
+ if self.schedule.is_constant:
+ beta = self.schedule(0)
+ return (1 - beta) ** t * torch.eye(dim) + (1 - (1 - beta) ** t) * self._get_mask()
+
+ else:
+ p = self.state[t]
+ return p * torch.eye(dim) + (1 - p) * self._get_mask()
+
+ def _get_mask(self):
+ dim = self.dim
+ return torch.ones((dim, dim)) * (torch.arange(0, dim)[:, None] == (dim - 1)).to(
+ torch.float32
+ )
+
+ def get(self, t):
+ _t = t if len(t.shape) == 1 else t[None]
+ beta = self.schedule(_t)
+ dim = self.dim
+
+ ret = (1 - beta)[:, None, None] * torch.eye(dim, device=_t.device)[None] + beta[
+ :, None, None
+ ] * self._get_mask().to(_t.device)[None]
+ return ret if len(t.shape) == 1 else ret.squeeze(0)
+
+ def qt_reverse(self, qt_plus_1, t, return_logits=False, make_one_hot=False, epsilon=1e-20):
+ """Get q(x_{t+1} | x_t), for each possible value of x_t. Thus, the rows of the output do not sum to 1.
+
+ Args:
+ qt_plus_1: an array of floats specifying a distribution over q(x_{t+1} | x_0).
+ t: t in q(x_{t+1} | x_t).
+ return_logits: if True, return the output logits
+ make_one_hot: if True, will convert q(x_{t+1}) to floats if needed.
+ epsilon: a small number to normalize logits conversion with, if needed.
+
+ Returns:
+ q(x_{t+1} | x_t), shape [num_samples, num_classes].
+ """
+
+ if make_one_hot:
+ assert qt_plus_1.dtype in [torch.long, torch.int32]
+ qt_plus_1 = torch.eye(self.dim, device=qt_plus_1.device)[qt_plus_1]
+
+ assert qt_plus_1.dtype == torch.float32
+
+ beta = self.schedule(t)
+
+ # q(x_{t+1} | x_t) = (1 - beta) if x_t = x_{t+1} != mask type
+ # else: beta if x_t != mask type else 1. (beta is the probability of transitioning to the absorbing state at t).
+ # I.e., if x_{t+1} is in some non-masked state S, then the probability of transitioning from S in t to S in t+1 is (1 - beta).
+ # Else, if x_{t+1} is in the masked state, then the probability of transitioning from a non-masked state S in t to the masked state in t+1 is beta,
+ # and the probability of transitioning from the masked state to itself is 1.
+ non_mask_prob = (1 - beta)[:, None] * qt_plus_1[:, :-1] + beta[:, None] * qt_plus_1[:, -1:]
+ prob_at_time_t = (
+ torch.eye(self.dim, device=qt_plus_1.device)[self.dim - 1][None] * qt_plus_1[:, -1:]
+ )
+ prob_at_time_t[:, :-1] = non_mask_prob
+
+ if return_logits:
+ return torch.log(prob_at_time_t + epsilon)
+ else:
+ return prob_at_time_t
+
+ def get_qt_given_q0(self, q0, t, return_logits=False, make_one_hot=False, epsilon=1e-20):
+ """Get q(x_t), the n-step posterior.
+
+ Can do efficiently for masks.
+
+ For example, for t = 0, it returns q0 unchanged.
+
+ Args:
+ q0: an array of floats specifying a distribution over p(x_0).
+ t: t in q(x_t | x_0).
+ return_logits: if True, return the output logits
+ make_one_hot: if True, will convert q0 to floats if needed.
+ epsilon: a small number to normalize logits conversion with, if needed.
+
+ Returns:
+ q(x_t | x_0).
+ """
+ if not self.supports_efficient_inference():
+ return super().get_qt_given_q0(
+ q0, t, return_logits=return_logits, make_one_hot=make_one_hot, epsilon=epsilon
+ )
+
+ if make_one_hot:
+ assert q0.dtype in [torch.int32, torch.long]
+ q0 = torch.eye(self.dim, device=q0.device)[q0]
+
+ assert q0.dtype == torch.float32
+ assert len(q0.shape) == 2
+
+ # p is probability of staying the same. (1 - p) is prob of masking.
+ p = self.state.to(q0.device)[t]
+
+ non_mask_prob = p[:, None] * q0[:, :-1]
+ mask_prob = 1 - non_mask_prob.sum(-1)
+
+ prob_at_time_t = (
+ mask_prob[:, None] * torch.eye(self.dim, device=q0.device)[self.dim - 1][None]
+ )
+ prob_at_time_t[:, :-1] = non_mask_prob
+
+ prob_at_time_t = torch.where(t[:, None] == 0, q0, prob_at_time_t)
+
+ if return_logits:
+ return torch.log(prob_at_time_t + epsilon)
+ else:
+ return prob_at_time_t
+
+ def supports_efficient_get(self):
+ return not self.use_fast_inference
+
+
+def create_discrete_diffusion_schedule(
+ kind="linear",
+ beta_min=1e-3,
+ beta_max=1e-1,
+ num_steps=100,
+ scale=1.0,
+):
+ """Creates a callable schedule object to use for diffusion rates.
+
+ Args:
+ kind: str, one of 'standard', 'linear', 'cosine', 'mutual_information'. If
+ standard, performs standard binomial diffusion taken from Sohl-Dicksteein
+ et al, ignoring betas. Otherwise, linear schedule between beta_min and
+ beta_max.
+ beta_min: the minimum beta. Ignored if kind == standard.
+ beta_max: the maximum beta.
+ num_steps: int, the number of steps to take.
+ scale: for standard schedule, rescales num_steps by this amount.
+
+ Returns:
+ a DiffusionSchedule object.
+ """
+
+ assert beta_min <= beta_max
+ assert num_steps > 0
+ assert scale >= 1
+
+ if kind == "standard":
+
+ def schedule_fn(step: Union[int, torch.Tensor]):
+ return 1 / (scale * num_steps - step)
+
+ return DiffusionSchedule(schedule_fn, num_steps, is_constant=False)
+
+ elif kind == "linear":
+ is_constant = beta_min == beta_max
+
+ linspace = torch.linspace(beta_min, beta_max, num_steps)
+
+ def schedule_fn(step: Union[int, torch.Tensor]):
+ return linspace[step]
+
+ return DiffusionSchedule(schedule_fn, num_steps, is_constant=is_constant)
+ elif kind == "cosine":
+ s = 0.008
+
+ def cosine_fn(step: torch.Tensor):
+ return torch.cos((step / num_steps + s) / (1 + s) * torch.pi / 2)
+
+ def schedule_fn(step: Union[int, torch.Tensor]):
+ if isinstance(step, int):
+ step = torch.tensor(step)
+ return torch.clamp(1 - (cosine_fn(step + 1) / cosine_fn(step)), 0, 0.999)
+
+ return DiffusionSchedule(schedule_fn, num_steps, is_constant=False)
+ else:
+ raise ValueError(f"kind {kind} is not supported.")
+
+
+def p_forward(
+ denoise_fn,
+ x_t,
+ t,
+ diffusion,
+ predict_x0=True,
+ return_x0=False,
+ return_logits=False,
+ special_case_x0=False,
+ transition_probs=None,
+ transition_probs_in_logits=True,
+ maximum_likelihood=False,
+ epsilon=1e-20,
+ step_size=1,
+):
+ """Returns probabilities from the reverse process p(x_{t-1} | x_t).
+
+ Args:
+ denoise_fn: the reverse process. Must support embed, call, and attend.
+ x_t: the current value of x_t to condition on.
+ t: the timestep t.
+ diffusion: the Diffusion object to use for noise.
+ predict_x0: if True, assumes the model output corresponds to its prediction
+ for p(x_0 | x_t). Otherwise assumes model predicts p(x_{t-1} | x_t).
+ return_x0: if True, will return probs for x_0 as well as x_{t-1}.
+ return_logits: if True, will return logits instead of probabilities.
+ special_case_x0: if True, will directly predict x0 instead of using the
+ forward process probabilities.
+ transition_probs: if provided, q(x_{t+1} | x_t) probs to reuse.
+ transition_probs_in_logits: if False, will ignore transition probs in logits
+ (only allowed if return_logits is True). This is because this term is
+ independent of theta.
+ maximum_likelihood: if true, will draw the most likely x0 before applying
+ the forward process.
+ epsilon: a small number.
+ step_size: step size to compute posterior from.
+
+ Returns:
+ probabilities for q(x_{t-1} | x_t) (and probabilities for x0 if predict_x0
+ is True)
+ """
+ assert not (step_size > 1 and not predict_x0)
+
+ logits = denoise_fn(targets=x_t, timestep=t)
+ probs = logits.softmax(dim=-1)
+
+ if not predict_x0:
+ retval = logits if return_logits else probs
+ if return_x0:
+ return retval, None
+ else:
+ return retval
+
+ if maximum_likelihood:
+ probs = probs.argmax(-1)
+
+ # we use this to compute p(x_{t-1} | x_t) = sum_x0 q(x_{t-1} | x_t, x_0)
+ # p(x_0 | x_t).
+ qt_probs, _ = diffusion.sample_and_compute_posterior_q(
+ x_0=probs,
+ t=t - step_size,
+ make_one_hot=maximum_likelihood,
+ return_logits=return_logits,
+ transition_probs_in_logits=transition_probs_in_logits,
+ transition_probs=transition_probs,
+ samples=x_t,
+ epsilon=epsilon,
+ step_size=step_size,
+ )
+
+ retval_x0 = logits if return_logits else probs
+ retval = qt_probs
+
+ # we can special case t = 1 to just use the raw logits outputs.
+ mask = (t == step_size) & special_case_x0
+ retval = mask[:, None] * retval_x0 + (mask.logical_not())[:, None] * retval
+
+ if return_x0:
+ return retval, retval_x0
+ else:
+ return retval
+
+
+def q_sample(x_start, t, diffusion, return_logits=False):
+ """Draws a sample from the posterior q(x_t | x_start)."""
+
+ assert x_start.dtype in [torch.int32, torch.long]
+
+ dim = diffusion.dim
+ x_start = torch.eye(dim, device=x_start.device)[x_start]
+
+ logits = diffusion.get_qt_given_q0(q0=x_start, t=t, return_logits=True)
+ sample = Categorical(logits=logits).sample()
+ if return_logits:
+ return sample, logits
+ return sample
+
+
+def compute_prior_kl(x_start, diffusion, target_mask=None):
+ """Computes KL divergence between q(x_T) and the true distribution."""
+ assert x_start.dtype in [torch.long, torch.int32]
+
+ num_steps = diffusion.num_steps
+
+ q_probs = diffusion.get_qt_given_q0(
+ q0=x_start,
+ t=torch.tensor(
+ [
+ num_steps,
+ ],
+ device=x_start.device,
+ ),
+ return_logits=False,
+ make_one_hot=True,
+ ) # get end step
+ p_probs = diffusion.stationary_probs(q_probs.shape[:-1]).to(q_probs.device)
+
+ d1 = Categorical(probs=q_probs)
+ d2 = Categorical(probs=p_probs)
+ loss = torch.distributions.kl_divergence(d1, d2)
+
+ if target_mask is not None:
+ loss = (loss * target_mask).sum()
+ else:
+ loss = loss.sum()
+
+ return loss
+
+
+def compute_kl_reverse_process(
+ x_start: torch.Tensor,
+ t: torch.Tensor,
+ *,
+ x_t_plus_1: Optional[torch.Tensor] = None,
+ diffusion: DiscreteDiffusionBase,
+ denoise_fn: Callable[[torch.Tensor, torch.Tensor], torch.Tensor],
+ predict_x0: bool = True,
+ log_space: bool = False,
+ label_smoothing: float = 0.0,
+ hybrid_lambda: float = 0.0,
+ use_cached_transition: bool = True,
+ target_mask: Optional[torch.Tensor] = None,
+ step_size: int = 1,
+) -> Dict[str, torch.Tensor]:
+ """Returns the KL for one term in the ELBO (time t) (loss L_t).
+
+ This assumes x_start is a sample from x_0, from which we draw samples from
+ q(x_t | x_0) and then compute q(x_{t-1} | x_t, x_0) following the LaTeX. This
+ is the KL divergence for terms L_1 through L_{T-1}.
+
+ Args:
+ x_start: a sample from p(data) (or q(x_0)).
+ t: the loss term to compute.
+ diffusion: the diffusion object to use.
+ denoise_fn: a functool.partial-ed version of the model_apply function which
+ takes a set of targets (x_t) and noise level and returns q(x_{t-1} | x_t,
+ x_0).
+ predict_x0: if True, will predict a distribution over x0 instead of x_{t-1}.
+ log_space: if True, will perform the loss calculations in log space.
+ label_smoothing: label smoothing for cross entropy.
+ hybrid_lambda: coefficient for hybrid cross-entropy loss.
+ use_cached_transition: if True, will reuse q(x_{t+1} | x_t) computation.
+ target_mask: mask for target sequence.
+ step_size: the step size over which the ELBO is computed.
+
+ Returns:
+ the KL divergence and denominator.
+ """
+ assert x_start.dtype in [torch.int32, torch.long]
+
+ if step_size > 1 and not predict_x0:
+ raise ValueError("cannot skip steps when not predicting x0.")
+
+ # If x_t_plus_1 is None, sample from q(x_{t+1} | x_start). Otherwise use the provided samples for x_{t+1}.
+ # Then compute q(x_t | x_{t+1}, x_start)
+ # q_t and p_t can be logits or probs depending on log_space.
+ q_t, x_t_plus_1, transition_probs = diffusion.sample_and_compute_posterior_q(
+ x_0=x_start,
+ t=t,
+ return_logits=log_space,
+ return_transition_probs=True,
+ step_size=step_size,
+ samples=x_t_plus_1,
+ )
+
+ transition_probs = transition_probs if use_cached_transition else None
+
+ p_t = p_forward(
+ denoise_fn=denoise_fn,
+ x_t=x_t_plus_1,
+ t=t + step_size,
+ diffusion=diffusion,
+ predict_x0=predict_x0,
+ return_x0=predict_x0 and hybrid_lambda > 0.0,
+ return_logits=log_space,
+ transition_probs=transition_probs,
+ step_size=step_size,
+ )
+
+ hybrid_loss = torch.tensor(0.0, device=x_start.device)
+ if predict_x0 and hybrid_lambda > 0.0:
+ # p_t, p_0 are shape [num_atoms, ].
+ p_t, p_0 = p_t
+ if log_space:
+ # [num_atoms, ]
+ cross_entropy = F.cross_entropy(
+ input=p_0, target=x_start, label_smoothing=label_smoothing, reduction="none"
+ )
+ else:
+ # [num_atoms, ]
+ cross_entropy = F.cross_entropy(
+ input=(p_0 + 1e-7).log(),
+ target=x_start,
+ label_smoothing=label_smoothing,
+ reduction="none",
+ )
+
+ hybrid_loss = hybrid_lambda * cross_entropy
+
+ assert not q_t.isnan().any() and not p_t.isnan().any()
+
+ if log_space:
+ d1 = Categorical(logits=q_t)
+ d2 = Categorical(logits=p_t)
+ # [num_atoms, ]
+ kl = torch.distributions.kl_divergence(p=d1, q=d2)
+ # [num_atoms, ]
+ cross_entropy = F.cross_entropy(
+ input=p_t, target=x_start, label_smoothing=label_smoothing, reduction="none"
+ )
+ else:
+ d1 = Categorical(logits=(q_t + 1e-7).log())
+ d2 = Categorical(logits=(p_t + 1e-7).log())
+ # [num_atoms, ]
+ kl = torch.distributions.kl_divergence(p=d1, q=d2)
+ # [num_atoms, ]
+ cross_entropy = F.cross_entropy(
+ input=(p_t + 1e-7).log(),
+ target=x_start,
+ label_smoothing=label_smoothing,
+ reduction="none",
+ )
+
+ if target_mask is not None: # can be used for inpainting
+ kl = kl * target_mask
+ cross_entropy = cross_entropy * target_mask
+ hybrid_loss = hybrid_loss * target_mask
+
+ # [num_atoms, ]
+ mask = t == 0
+ base_loss = mask * cross_entropy + (mask.logical_not()) * kl
+ loss = base_loss + hybrid_loss
+ denominator = torch.tensor(1, device=x_start.device)
+
+ metrics_dict = {
+ "loss": loss,
+ "denominator": denominator,
+ "kl/hybrid_loss": hybrid_loss,
+ "kl/base_loss": base_loss,
+ "kl/cross_entropy_loss": cross_entropy,
+ "kl/t0_loss": mask * cross_entropy,
+ "kl/kl_loss": kl,
+ }
+
+ return metrics_dict
diff --git a/model/diffusion/d3pm/d3pm_predictors_correctors.py b/model/diffusion/d3pm/d3pm_predictors_correctors.py
new file mode 100644
index 0000000000000000000000000000000000000000..0cc2a1e8cf5702870f9e95568938e9bf9b0e50b9
--- /dev/null
+++ b/model/diffusion/d3pm/d3pm_predictors_correctors.py
@@ -0,0 +1,106 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+from typing import Optional, cast
+
+import torch
+
+from ...diffusion.corruption.corruption import Corruption
+from ...diffusion.corruption.d3pm_corruption import D3PMCorruption
+from ...diffusion.corruption.sde_lib import ScoreFunction
+from ...diffusion.data.batched_data import BatchedData
+from ...diffusion.discrete_time import to_discrete_time
+from ...diffusion.sampling.predictors import Predictor
+from ...diffusion.sampling.predictors_correctors import SampleAndMean
+
+
+class D3PMAncestralSamplingPredictor(Predictor):
+ """
+ Ancestral sampling predictor for D3PM.
+ """
+
+ def __init__(
+ self,
+ *,
+ corruption: D3PMCorruption,
+ score_fn: ScoreFunction,
+ predict_x0: bool = True,
+ ):
+ super().__init__(corruption=corruption, score_fn=score_fn)
+ # if True, self.denoiser returns p(x_0|x_t), otherwise p(x_{t-1}|x_t)
+ self.predict_x0 = predict_x0
+
+ @classmethod
+ def is_compatible(cls, corruption: Corruption) -> bool:
+ return isinstance(corruption, D3PMCorruption)
+
+ @property
+ def N(self) -> int:
+ self.corruption = cast(D3PMCorruption, self.corruption) # mypy
+ return self.corruption.N
+
+ def update_given_score(
+ self,
+ *,
+ x: torch.Tensor,
+ t: torch.Tensor,
+ dt: torch.Tensor,
+ batch_idx: torch.LongTensor,
+ score: torch.Tensor,
+ batch: Optional[BatchedData],
+ ) -> SampleAndMean:
+ """
+ Takes the atom coordinates, cell vectors and atom types at time t and
+ returns the atom types at time t-1, sampled using the learned reverse
+ atom diffusion model.
+
+ Look at https://github.com/google-research/google-research/blob/master/d3pm/text/diffusion.py
+
+ lines 3201-3229. NOTE: we do implement the taking the softmax of the initial
+ sample as per 3226-3227. This could be to avoid weird behaving for picking
+ initial states that happened to have very low probability in latent space.
+ Try adding if there proves to be a problem generating samples.
+ """
+ # t is continuous, needs to be integer
+ t = to_discrete_time(t=t, N=self.N, T=self.corruption.T)
+
+ class_logits = score
+
+ assert isinstance(self.corruption, D3PMCorruption)
+
+ # sample from categorical distribution
+ x_sample = self.corruption._to_non_zero_based(
+ torch.distributions.Categorical(logits=class_logits).sample()
+ )
+
+ # convert logit output to normalized probabilities
+ class_probs = torch.softmax(class_logits, dim=-1)
+
+ # get expected atom type from categorical distribution
+ class_expected = self.corruption._to_non_zero_based(torch.argmax(class_probs, dim=-1))
+
+ if self.predict_x0:
+ # if self.predict_x0, the model predicts p(x_0|x_t), not p(x_{t-1}|x_t)
+ # We need to evaluate p(x_{t-1}|x_t) by Eq 4. in https://arxiv.org/pdf/2107.03006v1.pdf
+ assert isinstance(self.corruption, D3PMCorruption)
+ class_logits, _ = self.corruption.d3pm.sample_and_compute_posterior_q(
+ x_0=class_probs,
+ t=t[batch_idx].to(torch.long), # requires torch.long or torch.int32
+ make_one_hot=False,
+ samples=self.corruption._to_zero_based(
+ x
+ ), # d3pm expects 0 offset atom type integers
+ return_logits=True,
+ )
+
+ x_sample = self.corruption._to_non_zero_based(
+ torch.distributions.Categorical(logits=class_logits).sample()
+ )
+
+ # get expected atom type
+ class_expected = self.corruption._to_non_zero_based(
+ torch.argmax(torch.softmax(class_logits.to(class_probs.dtype), dim=-1), dim=-1)
+ )
+
+ # (sampled states), (expected states)
+ return x_sample, class_expected
diff --git a/model/diffusion/data/__init__.py b/model/diffusion/data/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/model/diffusion/data/batched_data.py b/model/diffusion/data/batched_data.py
new file mode 100644
index 0000000000000000000000000000000000000000..5dad923174b210ac42d48d3ca039dded23cbe912
--- /dev/null
+++ b/model/diffusion/data/batched_data.py
@@ -0,0 +1,188 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+from __future__ import annotations
+
+import logging
+from copy import deepcopy
+from dataclasses import dataclass, replace
+from typing import Any, Mapping, Protocol, Sequence, TypeVar, runtime_checkable
+
+import torch
+from torch_scatter import scatter
+
+T = TypeVar("T")
+
+logger = logging.getLogger(__name__)
+
+
+@runtime_checkable
+class BatchedData(Protocol):
+ def replace(self: T, **vals: torch.Tensor) -> T:
+ """Return a copy of self with some fields replaced with new values."""
+
+ def get_batch_idx(self, field_name: str) -> torch.LongTensor | None:
+ """Get the batch index (i.e., which row belongs to which sample) for a given field.
+ For 'dense' type data, where every sample has the same shape and the first dimension is the
+ batch dimension, this method should return None. Mathematically,
+ returning None will be treated the same as returning a tensor [0, 1, 2, ..., batch_size - 1]
+ but I expect memory access in other functions to be more efficient if you return None.
+ """
+
+ def get_batch_size(self) -> int:
+ """Get the batch size."""
+
+ def device(self) -> torch.device:
+ """Get the device of the batch."""
+
+ def __getitem__(self, field_name: str) -> torch.Tensor:
+ """Get a field from the batch."""
+
+ def to(self: T, device: torch.device) -> T:
+ """Move the batch to a given device."""
+
+ def clone(self: T) -> T:
+ """Return a copy with all the tensors cloned."""
+
+
+@dataclass
+class SimpleBatchedData(BatchedData):
+ """Implements BatchedData as a pair of mappings from field names to tensors."""
+
+ data: Mapping[str, Any]
+ batch_idx: Mapping[str, torch.LongTensor]
+
+ def replace(self, **vals: torch.Tensor) -> SimpleBatchedData:
+ """Return a copy of self with some fields of self.data replaced with new values."""
+ return replace(self, data=self._updated_data(**vals))
+
+ def get_batch_idx(self, field_name: str) -> torch.LongTensor | None:
+ return self.batch_idx[field_name]
+
+ def _updated_data(self, **vals):
+ return dict(self.data, **vals)
+
+ def __getitem__(self, key):
+ return self.data[key]
+
+ def __contains__(self, key):
+ return key in self.data
+
+ def get_batch_size(self) -> int:
+ L = []
+ for k, v in self.batch_idx.items():
+ if v is None:
+ d = self.data[k]
+ L.append(d.shape[0] if isinstance(d, torch.Tensor) else len(d))
+ else:
+ if len(v) == 0:
+ logger.warning(f"Empty batch index for field {k}")
+ L.append(0)
+ else:
+ L.append(int(torch.max(v).item()) + 1)
+ return max(L)
+
+ @property
+ def device(self) -> torch.device:
+ # Beware, there are no checks that all values are on the same device
+ return next(v.device for v in self.data.values())
+
+ def to(self, device) -> "SimpleBatchedData":
+ """Modify self in-place to move all tensors to the given device, and return self"""
+ if isinstance(self.data, dict):
+ for k in self.data.keys():
+ if isinstance(self.data[k], torch.Tensor):
+ self.data[k] = self.data[k].to(device)
+ if isinstance(self.batch_idx, dict):
+ for key in self.batch_idx.keys():
+ if self.batch_idx[key] is None:
+ continue
+ self.batch_idx[key] = self.batch_idx[key].to(device)
+
+ return self
+
+ def clone(self) -> SimpleBatchedData:
+ return SimpleBatchedData(
+ data={
+ k: v.clone() if isinstance(v, torch.Tensor) else deepcopy(v)
+ for k, v in self.data.items()
+ },
+ batch_idx={k: v.clone() if v is not None else None for k, v in self.batch_idx.items()},
+ )
+
+ def to_data_list(self) -> list[dict[str, torch.Tensor]]:
+ """Converts this instance to a list of dictionaries, each of which corresponds to a single datapoint in
+ `batched_data`. The keys of the dictionaries match the keys of `batched_data`.
+ """
+
+ batch_size = self.get_batch_size()
+ if batch_size == 0:
+ return []
+
+ def _unpack(k, i):
+ if self.batch_idx[k] is not None:
+ return self.data[k][self.batch_idx[k] == i]
+ elif isinstance(self.data[k], torch.Tensor):
+ return self.data[k][i : i + 1]
+ else:
+ return self.data[k][i]
+
+ return [{k: _unpack(k, i) for k in self.data.keys()} for i in range(batch_size)]
+
+
+def collate_fn(
+ states: list[dict[str, Any]], dense_field_names: Sequence[str] = ()
+) -> SimpleBatchedData:
+ """
+ Combine a list of samples into a SimpleBatchedData object.
+
+ The association between the index in `states[i][k]` and a row in the `batched_data[k]` is
+ stored in `batched_data.batch_idx[k]`. If the `k` appears in
+ `dense_field_names`, `batched_data.batch_idx[k]` is `None` and the data is
+ simply stacked along the first dimension.
+
+ Non-tensor values are put into lists.
+ """
+ assert states, "Cannot collate empty list"
+ concatenated_data = {}
+ batch_idx: dict[str, torch.Tensor | None] = {}
+ for k, v in states[0].items():
+ if isinstance(v, torch.Tensor):
+ concatenated_data[k] = torch.cat([x[k] for x in states], dim=0)
+ if k in dense_field_names:
+ if any(x[k].shape[0] != 1 for x in states):
+ raise ValueError(
+ f"First dimension should be batch dimension. Instead key {k} has shape {states[0][k].shape}"
+ )
+ batch_idx[k] = None
+ else:
+ batch_idx[k] = _construct_batch_idx(states, k)
+ else:
+ concatenated_data[k] = [x[k] for x in states]
+ batch_idx[k] = None
+
+ batch = SimpleBatchedData(data=concatenated_data, batch_idx=batch_idx)
+ if "edge_index" in batch.data:
+ batch = batch.replace(
+ edge_index=_batch_edge_index(
+ batch["edge_index"],
+ batch.batch_idx["atomic_numbers"],
+ batch.batch_idx["edge_index"],
+ ),
+ )
+ return batch
+
+
+def _batch_edge_index(edge_index, atom_batch_idx, edge_batch_idx):
+ num_atoms = scatter(torch.ones_like(atom_batch_idx), atom_batch_idx)
+ num_atoms_acc = torch.nn.functional.pad(torch.cumsum(num_atoms, 0)[:-1], [1, 0], "constant", 0)
+ return edge_index + num_atoms_acc[edge_batch_idx].unsqueeze(1)
+
+
+def _construct_batch_idx(data_list: list[Any], field_name: str) -> torch.LongTensor:
+ """Construct batch index tensor for one field."""
+ batch_size = len(data_list)
+ return torch.repeat_interleave(
+ torch.arange(0, batch_size),
+ torch.tensor([x[field_name].shape[0] for x in data_list]),
+ )
diff --git a/model/diffusion/diffusion_module.py b/model/diffusion/diffusion_module.py
new file mode 100644
index 0000000000000000000000000000000000000000..ce61b0647cbfbd930d40555da7c28faab3ec3c79
--- /dev/null
+++ b/model/diffusion/diffusion_module.py
@@ -0,0 +1,158 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+from typing import Callable, Generic, TypeVar
+
+import torch
+
+from ..diffusion.corruption.multi_corruption import MultiCorruption, apply
+from ..diffusion.data.batched_data import BatchedData
+from ..diffusion.losses import Loss
+from ..diffusion.model_target import ModelTarget
+from ..diffusion.model_utils import convert_model_out_to_score
+from ..diffusion.score_models.base import ScoreModel
+from ..diffusion.timestep_samplers import TimestepSampler, UniformTimestepSampler
+
+T = TypeVar("T", bound=BatchedData)
+BatchTransform = Callable[[T], T]
+
+
+def identity(x: T) -> T:
+ return x
+
+
+class DiffusionModule(torch.nn.Module, Generic[T]):
+ """Denoising diffusion model for a multi-part state"""
+
+ def __init__(
+ self,
+ model: ScoreModel[T],
+ corruption: MultiCorruption[T],
+ loss_fn: Loss,
+ pre_corruption_fn: BatchTransform | None = None,
+ timestep_sampler: TimestepSampler | None = None,
+ ) -> None:
+ super().__init__()
+ self.model = model
+ self.corruption = corruption
+ self.loss_fn = loss_fn
+ self.pre_corruption_fn = pre_corruption_fn or identity
+ self.model_targets = {k: ModelTarget(v) for k, v in loss_fn.model_targets.items()}
+
+ self.timestep_sampler = timestep_sampler or UniformTimestepSampler(
+ min_t=1e-5,
+ max_t=corruption.T,
+ )
+
+ # Check corruption for nn.Modules and register them here.
+ self._register_corruption_modules()
+
+ def _register_corruption_modules(self):
+ """
+ Register corruptions that are instances of `torch.nn.Module`s for proper device, parameter,
+ etc handling.
+ """
+ assert isinstance(self.corruption, MultiCorruption)
+ for idx, (key, _corruption) in enumerate(self.corruption._corruptions.items()):
+ if isinstance(_corruption, torch.nn.Module):
+ self.register_module(f"MultiCorruption:{idx}:{key}", _corruption)
+
+ def calc_loss(
+ self, batch: T, node_is_unmasked: torch.LongTensor | None = None
+ ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
+ """
+ Calculate loss and metrics given a batch of clean data which may include
+ context/conditioning fields. Add noise, predict score using score model, then calculate
+ loss.
+
+ Args:
+ batch: batch of training data
+ node_is_unmasked: mask that has a value 1 for nodes that are included in the loss, and
+ a value of 0 for nodes that should be ignored. If None, all nodes are included.
+
+ Returns:
+ loss: the loss for the batch
+ metrics: a dictionary of metrics for the batch
+ """
+ batch = self.pre_corruption_fn(batch)
+
+ noisy_batch, t = self._corrupt_batch(batch)
+
+ score_model_output = self.model(noisy_batch, t)
+ loss, metrics = self.loss_fn(
+ multi_corruption=self.corruption,
+ batch=batch,
+ noisy_batch=noisy_batch,
+ score_model_output=score_model_output,
+ t=t,
+ node_is_unmasked=node_is_unmasked,
+ )
+ assert loss.numel() == 1
+
+ return loss, metrics
+
+ def _corrupt_batch(
+ self,
+ batch: T,
+ ) -> tuple[T, torch.Tensor]:
+ """
+ Corrupt a batch of data for use in a training step:
+ - sample a different timestep for each sample in the batch
+ - add noise according to the corruption process
+
+ Args:
+ batch: Batch of clean states
+
+ Returns:
+ noisy_batch: batch of noisy samples
+ t: the timestep used for each sample in the batch
+
+ """
+ # Sample timesteps
+ t = self.sample_timesteps(batch)
+
+ # Add noise to data
+ noisy_batch = self.corruption.sample_marginal(batch, t)
+
+ return noisy_batch, t
+
+ def score_fn(self, x: T, t: torch.Tensor) -> T:
+ """Calculate the score of a batch of data at a given timestep
+
+ Args:
+ x: batch of data
+ t: timestep
+
+ Returns:
+ score: score of the batch of data at the given timestep
+ """
+ model_out: T = self.model(x, t)
+ fns = {k: convert_model_out_to_score for k in self.corruption.sdes.keys()}
+
+ scores = apply(
+ fns=fns,
+ model_out=model_out,
+ broadcast=dict(t=t, batch=x),
+ sde=self.corruption.sdes,
+ model_target=self.model_targets,
+ batch_idx=self.corruption._get_batch_indices(x),
+ )
+
+ return model_out.replace(**scores)
+
+ def sample_timesteps(self, batch: T) -> torch.Tensor:
+ """Sample the timesteps, which will be used to determine how much noise
+ to add to data.
+
+ Args:
+ batch: batch of data to be corrupted
+
+ Returns: sampled timesteps
+ """
+ return self.timestep_sampler(
+ batch_size=batch.get_batch_size(),
+ device=self._get_device(batch),
+ )
+
+ def _get_device(self, batch: T) -> torch.device:
+ return next(batch[k].device for k in self.corruption.sdes.keys())
diff --git a/model/diffusion/discrete_time.py b/model/diffusion/discrete_time.py
new file mode 100644
index 0000000000000000000000000000000000000000..6829d3a0103c615ef2e0d9623879bfbb4238f2f2
--- /dev/null
+++ b/model/diffusion/discrete_time.py
@@ -0,0 +1,17 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+import torch
+
+
+def to_discrete_time(t: torch.Tensor, N: int, T: float) -> torch.LongTensor:
+ """Convert continuous time to integer timestep.
+
+ Args:
+ t: continuous time between 0 and T
+ N: number of timesteps
+ T: max time
+ Returns:
+ Integer timesteps between 0 and N-1
+ """
+ return ((t * (N - 1)) / T).long()
diff --git a/model/diffusion/exceptions.py b/model/diffusion/exceptions.py
new file mode 100644
index 0000000000000000000000000000000000000000..fb36a0beeadec5c2c7da571fb6d2f044a7da9876
--- /dev/null
+++ b/model/diffusion/exceptions.py
@@ -0,0 +1,11 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+class IncompatibleSampler(ValueError):
+ # Raised when sampler type and SDE are incompatible.
+ pass
+
+
+class AmbiguousConfig(ValueError):
+ # Raised when the config is ambiguous
+ pass
diff --git a/model/diffusion/lightning_module.py b/model/diffusion/lightning_module.py
new file mode 100644
index 0000000000000000000000000000000000000000..5771ede76d34e33c774fbd3118513ee2f195fa6b
--- /dev/null
+++ b/model/diffusion/lightning_module.py
@@ -0,0 +1,172 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+from __future__ import annotations
+
+from collections import deque
+from typing import Any, Dict, Generic, Optional, Protocol, Sequence, TypeVar, Union
+
+import numpy as np
+import pytorch_lightning as pl
+import torch
+from hydra.errors import InstantiationException
+from hydra.utils import instantiate
+from omegaconf import DictConfig
+from pytorch_lightning.utilities.types import STEP_OUTPUT
+from torch.optim import AdamW, Optimizer
+from tqdm import tqdm
+
+from ..diffusion.config import Config
+from ..diffusion.data.batched_data import BatchedData
+from ..diffusion.diffusion_module import DiffusionModule
+
+T = TypeVar("T", bound=BatchedData)
+
+
+class OptimizerPartial(Protocol):
+ """Callable to instantiate an optimizer."""
+
+ def __call__(self, params: Any) -> Optimizer:
+ raise NotImplementedError
+
+
+class SchedulerPartial(Protocol):
+ """Callable to instantiate a learning rate scheduler."""
+
+ def __call__(self, optimizer: Optimizer) -> Any:
+ raise NotImplementedError
+
+
+def get_default_optimizer(params):
+ return AdamW(params=params, lr=1e-4, weight_decay=0, amsgrad=True)
+
+
+class DiffusionLightningModule(pl.LightningModule, Generic[T]):
+ """LightningModule for instantiating and training a DiffusionModule."""
+
+ def __init__(
+ self,
+ diffusion_module: DiffusionModule[T],
+ optimizer_partial: Optional[OptimizerPartial] = None,
+ scheduler_partials: Optional[Sequence[Dict[str, Union[Any, SchedulerPartial]]]] = None,
+ ):
+ """_summary_
+
+ Args:
+ diffusion_module: The diffusion module to use.
+ optimizer_partial: Used to instantiate optimizer.
+ scheduler_partials: used to instantiate learning rate schedulers
+ """
+ super().__init__()
+ scheduler_partials = scheduler_partials or []
+ optimizer_partial = optimizer_partial or get_default_optimizer
+ self.save_hyperparameters(
+ ignore=("optimizer_partial", "scheduler_partials", "diffusion_module")
+ )
+
+ self.diffusion_module = diffusion_module
+ self._optimizer_partial = optimizer_partial
+ self._scheduler_partials = scheduler_partials
+
+ @classmethod
+ def load_from_checkpoint(
+ cls,
+ checkpoint_path: str,
+ map_location: Optional[str] = None,
+ **kwargs,
+ ) -> DiffusionLightningModule:
+ """Load model from checkpoint. kwargs are passed to hydra's instantiate and can override
+ arguments from the checkpoint config."""
+ checkpoint = torch.load(checkpoint_path, map_location=map_location)
+
+ # The config should have been saved in the checkpoint by AddConfigCallback in run.py
+ config = Config(**checkpoint["config"])
+ try:
+ lightning_module = instantiate(config.lightning_module, **kwargs)
+ except InstantiationException as e:
+ print("Could not instantiate model from the checkpoint.")
+ print(
+ "If the error is due to an unexpected argument because the checkpoint and the code have diverged, try using load_from_checkpoint_and_config instead."
+ )
+ raise e
+ assert isinstance(lightning_module, cls)
+
+ # Restore state of the DiffusionLightningModule.
+ lightning_module.load_state_dict(checkpoint["state_dict"])
+ return lightning_module
+
+ @classmethod
+ def load_from_checkpoint_and_config(
+ cls,
+ checkpoint_path: str,
+ config: DictConfig,
+ map_location: Optional[str] = None,
+ strict: bool = True,
+ ) -> tuple[DiffusionLightningModule, torch.nn.modules.module._IncompatibleKeys]:
+ """Load model from checkpoint, but instead of using the config stored in the checkpoint,
+ use the config passed in as an argument. This is useful when, e.g., an unused argument was
+ removed in the code but is still present in the checkpoint config."""
+ checkpoint = torch.load(checkpoint_path, map_location=map_location)
+
+ lightning_module = instantiate(config)
+ assert isinstance(lightning_module, cls)
+
+ # Restore state of the DiffusionLightningModule.
+ result = lightning_module.load_state_dict(checkpoint["state_dict"], strict=strict)
+
+ return lightning_module, result
+
+ def configure_optimizers(self) -> Any:
+ optimizer = self._optimizer_partial(params=self.diffusion_module.parameters())
+ if self._scheduler_partials:
+ lr_schedulers = [
+ {
+ **scheduler_dict,
+ "scheduler": scheduler_dict["scheduler"](
+ optimizer=optimizer,
+ ),
+ }
+ for scheduler_dict in self._scheduler_partials
+ ]
+
+ return [
+ optimizer,
+ ], lr_schedulers
+ else:
+ return optimizer
+
+ def training_step(self, train_batch: T, batch_idx: int) -> STEP_OUTPUT:
+ return self._calc_loss(train_batch, True)
+
+ def validation_step(self, val_batch: T, batch_idx: int) -> Optional[STEP_OUTPUT]:
+ return self._calc_loss(val_batch, False)
+
+ def test_step(self, test_batch: T, batch_idx: int) -> Optional[STEP_OUTPUT]:
+ return self._calc_loss(test_batch, False)
+
+ def _calc_loss(self, batch: T, train: bool) -> Optional[STEP_OUTPUT]:
+ """Calculate loss and metrics given a batch of clean data."""
+ loss, metrics = self.diffusion_module.calc_loss(batch)
+ # Log the results
+ step_type = "train" if train else "val"
+ batch_size = batch.get_batch_size()
+ self.log(
+ f"loss_{step_type}",
+ loss,
+ on_step=train,
+ on_epoch=True,
+ prog_bar=train,
+ batch_size=batch_size,
+ sync_dist=True,
+ )
+ for k, v in metrics.items():
+ self.log(
+ f"{k}_{step_type}",
+ v,
+ on_step=train,
+ on_epoch=True,
+ prog_bar=train,
+ batch_size=batch_size,
+ sync_dist=True,
+ )
+ return loss
diff --git a/model/diffusion/losses.py b/model/diffusion/losses.py
new file mode 100644
index 0000000000000000000000000000000000000000..e0ed59089e0a4ca1ba1664c8069ae0e1a177b7cc
--- /dev/null
+++ b/model/diffusion/losses.py
@@ -0,0 +1,125 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+from functools import partial
+from typing import Dict, Literal, Optional, Protocol, Tuple, TypeVar
+
+import torch
+
+from ..diffusion.corruption.multi_corruption import MultiCorruption, apply
+from ..diffusion.data.batched_data import BatchedData
+from ..diffusion.model_target import ModelTargets
+from ..diffusion.training.field_loss import FieldLoss, denoising_score_matching
+
+T = TypeVar("T", bound=BatchedData)
+
+
+class Loss(Protocol[T]):
+ """Loss function for training a score model on multi-field data."""
+
+ def __call__(
+ self,
+ *,
+ multi_corruption: MultiCorruption[T],
+ batch: T,
+ noisy_batch: T,
+ score_model_output: T,
+ t: torch.Tensor,
+ node_is_unmasked: Optional[torch.LongTensor] = None,
+ ) -> Tuple[torch.Tensor, Dict[str, float]]:
+ pass
+
+ """model_targets tells us what this loss function trains the score model to predict.
+ We need this information in order to convert the model output to a score during sampling.
+ """
+ model_targets: ModelTargets
+
+
+class SummedFieldLoss(Loss[T]):
+ """(Weighted) sum of different loss functions applied on each field."""
+
+ def __init__(
+ self,
+ loss_fns: Dict[str, FieldLoss],
+ model_targets: ModelTargets,
+ weights: Optional[Dict[str, float]] = None,
+ ) -> None:
+ self.model_targets = model_targets
+ self.loss_fns = loss_fns
+ # weights are optional, if not provided, all fields are weighted equally with weight 1.
+ if weights is None:
+ self.loss_weights = {k: 1.0 for k in self.loss_fns.keys()}
+ else:
+ assert set(weights.keys()) == set(
+ self.loss_fns.keys()
+ ), f"weight keys {set(weights.keys())} do not match loss_fns keys {set(self.loss_fns.keys())}"
+ self.loss_weights = weights
+
+ def __call__(
+ self,
+ *,
+ multi_corruption: MultiCorruption[T],
+ batch: T,
+ noisy_batch: T,
+ score_model_output: T,
+ t: torch.Tensor,
+ node_is_unmasked: Optional[torch.LongTensor] = None,
+ ) -> Tuple[torch.Tensor, Dict[str, float]]:
+ batch_idx = {k: batch.get_batch_idx(k) for k in self.loss_fns.keys()}
+ node_is_unmasked = {k: node_is_unmasked for k in self.loss_fns.keys()}
+
+ # Dict[str, torch.Tensor]
+ # Keys are field names and values are loss per sample, with shape (batch_size,).
+ loss_per_sample_per_field = apply(
+ fns=self.loss_fns,
+ corruption=multi_corruption.corruptions,
+ x=batch,
+ noisy_x=noisy_batch,
+ score_model_output=score_model_output,
+ batch_idx=batch_idx,
+ broadcast=dict(t=t, batch_size=batch.get_batch_size(), batch=batch),
+ node_is_unmasked=node_is_unmasked,
+ )
+ assert set([v.shape for v in loss_per_sample_per_field.values()]) == {
+ (batch.get_batch_size(),)
+ }, "All losses should have shape (batch_size,)."
+ # Aggregate losses per field over samples.
+ scalar_loss_per_field = {k: v.mean() for k, v in loss_per_sample_per_field.items()}
+
+ # Dict[str, torch.Tensor], dictionary containing metrics to be logged,
+ metrics_dict = scalar_loss_per_field
+ # This is the loss that is used for backpropagation (after mean aggregation over samples).
+ # Shape: (batch_size,)
+ agg_loss = torch.stack(
+ [self.loss_weights[k] * v for k, v in loss_per_sample_per_field.items()], dim=0
+ ).sum(0)
+
+ return (
+ agg_loss.mean(),
+ metrics_dict,
+ )
+
+
+class DenoisingScoreMatchingLoss(SummedFieldLoss):
+ def __init__(
+ self,
+ model_targets: ModelTargets,
+ reduce: Literal["sum", "mean"] = "mean",
+ weights: Optional[Dict[str, float]] = None,
+ field_center_zero: Optional[Dict[str, bool]] = None, # Whether to zero center each field.
+ ):
+ if field_center_zero is not None:
+ assert set(field_center_zero.keys()) == set(model_targets.keys())
+
+ super().__init__(
+ loss_fns={
+ k: partial(
+ denoising_score_matching,
+ reduce=reduce,
+ model_target=v,
+ )
+ for k, v in model_targets.items()
+ },
+ model_targets=model_targets,
+ weights=weights,
+ )
diff --git a/model/diffusion/model_target.py b/model/diffusion/model_target.py
new file mode 100644
index 0000000000000000000000000000000000000000..1a473762fcf9e2c32fefada351de9a9309cb0004
--- /dev/null
+++ b/model/diffusion/model_target.py
@@ -0,0 +1,16 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+from enum import Enum
+from typing import Mapping, Union
+
+
+class ModelTarget(Enum):
+ """Specifies what the score model is trained to predict.
+ Only relevant for fields that are corrupted with an SDE."""
+
+ score_times_std = "score_times_std" # Predict -z where z is gaussian noise with unit variance used to corrupt the data
+ logits = "logits" # Predict logits for a categorical variable
+
+
+ModelTargets = Mapping[str, Union[ModelTarget, str]]
diff --git a/model/diffusion/model_utils.py b/model/diffusion/model_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..d8565c19592a1cca35705affb7a518fbdb497c1d
--- /dev/null
+++ b/model/diffusion/model_utils.py
@@ -0,0 +1,74 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+import math
+from typing import Any, TypeVar
+
+import torch
+
+from ..diffusion.corruption.sde_lib import SDE
+from ..diffusion.data.batched_data import BatchedData
+from ..diffusion.model_target import ModelTarget
+
+T = TypeVar("T", bound=BatchedData)
+
+
+def convert_model_out_to_score(
+ *,
+ model_target: ModelTarget,
+ sde: SDE,
+ model_out: torch.Tensor,
+ batch_idx: torch.LongTensor,
+ t: torch.Tensor,
+ batch: Any
+) -> torch.Tensor:
+ """
+ Convert a model output to a score, according to the specified model_target.
+
+ model_target: says what the model predicts.
+ For example, in RFDiffusion the model predicts clean coordinates;
+ in EDM the model predicts the raw noise.
+ sde: corruption process
+ model_out: model output
+ batch_idx: indicates which sample each row of model_out belongs to
+ noisy_x: noisy data
+ t: diffusion timestep
+ batch: noisy batch, ignored except by strange SDEs
+ """
+ _, std = sde.marginal_prob(
+ x=torch.ones_like(model_out),
+ t=t,
+ batch_idx=batch_idx,
+ batch=batch,
+ )
+ # Note the slack tolerances in test_model_utils.py: the choice of ModelTarget does make a difference.
+ if model_target == ModelTarget.score_times_std:
+ return model_out / std
+ elif model_target == ModelTarget.logits:
+ # Not really a score, but logits will be handled downstream.
+ return model_out
+ else:
+ raise NotImplementedError
+
+
+class NoiseLevelEncoding(torch.nn.Module):
+ """
+ From: https://pytorch.org/tutorials/beginner/transformer_tutorial.html
+ """
+
+ def __init__(self, d_model: int, dropout: float = 0.0):
+ super().__init__()
+ self.dropout = torch.nn.Dropout(p=dropout)
+ self.d_model = d_model
+ div_term = torch.exp(torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model))
+ self.register_buffer("div_term", div_term)
+
+ def forward(self, t: torch.Tensor) -> torch.Tensor:
+ """
+ Args:
+ t: Tensor, shape [batch_size]
+ """
+ x = torch.zeros((t.shape[0], self.d_model), device=self.div_term.device)
+ x[:, 0::2] = torch.sin(t[:, None] * self.div_term[None])
+ x[:, 1::2] = torch.cos(t[:, None] * self.div_term[None])
+ return self.dropout(x)
diff --git a/model/diffusion/run.py b/model/diffusion/run.py
new file mode 100644
index 0000000000000000000000000000000000000000..951e3daefe53f734e8e3a5652362482dfb2803d3
--- /dev/null
+++ b/model/diffusion/run.py
@@ -0,0 +1,179 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+import argparse
+import logging
+import os
+import random
+import re
+from glob import glob
+from typing import Any, Mapping, TypeVar
+
+import numpy as np
+import pytorch_lightning as pl
+import torch
+import yaml
+from hydra.utils import instantiate
+from omegaconf import DictConfig, OmegaConf
+from pytorch_lightning import Callback
+from pytorch_lightning.callbacks import ModelCheckpoint
+from pytorch_lightning.cli import SaveConfigCallback
+from pytorch_lightning.utilities import rank_zero_only
+
+from ..common.utils.config_utils import get_config
+from ..diffusion.config import Config
+from ..diffusion.exceptions import AmbiguousConfig
+from ..diffusion.lightning_module import DiffusionLightningModule
+
+T = TypeVar("T")
+
+# logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+def maybe_instantiate(instance_or_config: T | Mapping, expected_type=None, **kwargs) -> T:
+ """
+ If instance_or_config is a mapping with a _target_ field, instantiate it.
+ Otherwise, return it as is.
+ """
+ if isinstance(instance_or_config, Mapping) and "_target_" in instance_or_config:
+ instance = instantiate(instance_or_config, **kwargs)
+ else:
+ instance = instance_or_config
+ assert expected_type is None or isinstance(
+ instance, expected_type
+ ), f"Expected {expected_type}, got {type(instance)}"
+ return instance
+
+
+def _find_latest_checkpoint(dirpath: str) -> str | None:
+ """Finds the most recent checkpoint inside `dirpath`."""
+
+ # checkpoint names are like "epoch=0-step=0.ckpt."
+ # Find the checkpoint with highest epoch:
+ def extract_epoch(ckpt):
+ match = re.search(r"epoch=(\d+)", ckpt)
+ if match:
+ return int(match.group(1))
+ return -1
+
+ ckpts = glob(f"{dirpath}/*.ckpt")
+ epochs = np.array([extract_epoch(ckpt) for ckpt in ckpts])
+ if len(epochs) == 0 or epochs.max() < 0:
+ # No checkpoints found.
+ return None
+ latest_checkpoint = ckpts[epochs.argmax()]
+ return latest_checkpoint
+
+
+class SimpleParser:
+ def save(self, config, path, **_):
+ with open(path, "w") as f:
+ yaml.dump(config, f)
+
+
+class AddConfigCallback(Callback):
+ """Adds a copy of the config to the checkpoint, so that `load_from_checkpoint` can use it to instantiate everything."""
+
+ def __init__(self, config: dict[str, Any]):
+ self._config_dict = config
+
+ def on_save_checkpoint(
+ self, trainer: "pl.Trainer", pl_module: "pl.LightningModule", checkpoint: dict[str, Any]
+ ) -> None:
+ checkpoint["config"] = self._config_dict
+
+
+def main(
+ config: Config | DictConfig, save_config: bool = True, seed: int | None = None
+) -> tuple[pl.Trainer, pl.LightningModule]:
+ """
+ Main entry point to train and evaluate a diffusion model.
+
+ save_config: if True, the config will be saved both as a YAML file and in each checkpoint. This doesn't work if the config contains things that can't be `yaml.dump`-ed, so
+ if you don't care about saving and loading checkpoints and want to use a config that contains things like `torch.nn.Module`s already instantiated, set this to False.
+ """
+ if config.checkpoint_path and config.auto_resume:
+ raise AmbiguousConfig(
+ f"Ambiguous config: you set both a checkpoint path {config.checkpoint_path} and `auto_resume` which means automatically select a checkpoint path to resume from."
+ )
+
+ if seed is not None:
+ torch.manual_seed(seed)
+ torch.cuda.manual_seed(seed)
+ torch.cuda.manual_seed_all(seed)
+ np.random.seed(seed)
+ random.seed(seed)
+
+ trainer: pl.Trainer = maybe_instantiate(config.trainer, pl.Trainer)
+
+ if save_config:
+ if isinstance(config, DictConfig):
+ config_as_dict = OmegaConf.to_container(config, resolve=True)
+ # This callback will save a config.yaml file.
+ trainer.callbacks.append(
+ SaveConfigCallback(
+ parser=SimpleParser(),
+ config=config_as_dict,
+ overwrite=True if config.auto_resume else False,
+ )
+ )
+
+ # This callback will add a copy of the config to each checkpoint.
+ trainer.callbacks.append(AddConfigCallback(config_as_dict))
+ else:
+ raise NotImplementedError
+ datamodule: pl.LightningDataModule = maybe_instantiate(
+ config.data_module, pl.LightningDataModule
+ )
+
+ # If checkpoint_path is provided training will be resumed from this point.
+ # Beware: the old checkpoint will be deleted when a new one is saved.
+
+ ckpt_path = config.checkpoint_path
+ if config.auto_resume:
+ # Add an additional checkpointer with a fixed directory path to restore from.
+ dirpath = os.path.join(trainer.default_root_dir, "checkpoints")
+ trainer.callbacks.append(ModelCheckpoint(dirpath=dirpath))
+ ckpt_path = _find_latest_checkpoint(dirpath)
+ pl_module: DiffusionLightningModule = maybe_instantiate(
+ config.lightning_module, DiffusionLightningModule
+ )
+ if rank_zero_only.rank == 0 and isinstance(trainer.logger, pl.loggers.WandbLogger):
+ # Log the config to wandb so that it shows up in the portal.
+ trainer.logger.experiment.config.update(
+ {**OmegaConf.to_container(config, resolve=True)},
+ allow_val_change=True,
+ )
+ trainer.fit(
+ pl_module,
+ datamodule=datamodule,
+ ckpt_path=ckpt_path,
+ )
+
+ return trainer, pl_module
+
+
+def cli(argv: list[str] | None) -> None:
+ """
+ Args:
+ argv: list of command-line arguments as strings, or None. If None,
+ command-line arguments will be got from sys.argv
+ """
+
+ parser = argparse.ArgumentParser(allow_abbrev=False) # prevent prefix matching issues
+ parser.add_argument(
+ "--seed",
+ type=int,
+ help="Random seed to use. If not provided, a random seed will be used.",
+ )
+ args, argv = parser.parse_known_args(argv)
+
+ # Create config from command-line arguments.
+ config = get_config(argv, Config)
+ main(config, seed=args.seed)
+
+
+if __name__ == "__main__":
+ cli(argv=None)
diff --git a/model/diffusion/sampling/__init__.py b/model/diffusion/sampling/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/model/diffusion/sampling/classifier_free_guidance.py b/model/diffusion/sampling/classifier_free_guidance.py
new file mode 100644
index 0000000000000000000000000000000000000000..d72c29d85089644da581c814ab03e556cfdeec6e
--- /dev/null
+++ b/model/diffusion/sampling/classifier_free_guidance.py
@@ -0,0 +1,100 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+from typing import Callable
+
+import torch
+
+from ...diffusion.sampling.pc_sampler import Diffusable, PredictorCorrector
+from onescience.datapipes.materials.mattergen.collate import collate
+
+BatchTransform = Callable[[Diffusable], Diffusable]
+
+
+def identity(x: Diffusable) -> Diffusable:
+ """
+ Default function that transforms data to its conditional state
+ """
+ return x
+
+
+class GuidedPredictorCorrector(PredictorCorrector):
+ """
+ Sampler for classifier-free guidance.
+ """
+
+ def __init__(
+ self,
+ *,
+ guidance_scale: float,
+ remove_conditioning_fn: BatchTransform,
+ keep_conditioning_fn: BatchTransform | None = None,
+ **kwargs,
+ ):
+ """
+ guidance_scale: gamma in p_gamma(x|y)=p(x)p(y|x)**gamma for classifier-free guidance
+ remove_conditioning_fn: function that removes conditioning from the data
+ keep_conditioning_fn: function that will be applied to the data before evaluating the conditional score. For example, this function might drop some fields that you never want to condition on or add fields that indicate which conditions should be respected.
+ **kwargs: passed on to parent class constructor.
+ """
+
+ super().__init__(**kwargs)
+ self._remove_conditioning_fn = remove_conditioning_fn
+ self._keep_conditioning_fn = keep_conditioning_fn or identity
+ self._guidance_scale = guidance_scale
+
+ def _score_fn(
+ self,
+ x: Diffusable,
+ t: torch.Tensor,
+ ) -> Diffusable:
+ """For each field, regardless of whether the corruption process is SDE or D3PM, we guide the score in the same way here,
+ by taking a linear combination of the conditional and unconditional score model output.
+
+ For discrete fields, the score model outputs are interpreted as logits, so the linear combination here means we compute logits for
+ p_\gamma(x|y)=p(x)^(1-\gamma) p(x|y)^\gamma
+
+ """
+
+ def get_unconditional_score():
+ return super(GuidedPredictorCorrector, self)._score_fn(
+ x=self._remove_conditioning_fn(x), t=t
+ )
+
+ def get_conditional_score():
+ return super(GuidedPredictorCorrector, self)._score_fn(
+ x=self._keep_conditioning_fn(x), t=t
+ )
+
+ if abs(self._guidance_scale - 1) < 1e-15:
+ return get_conditional_score()
+ elif abs(self._guidance_scale) < 1e-15:
+ return get_unconditional_score()
+ else:
+ # guided_score = guidance_factor * conditional_score + (1-guidance_factor) * unconditional_score
+ batch_no_condition = self._remove_conditioning_fn(x)
+ batch_with_condition = self._keep_conditioning_fn(x)
+ joint_batch = collate([batch_no_condition, batch_with_condition])
+
+ for attr,value in batch_no_condition.items():
+ if isinstance(value, list):
+ joint_batch[attr] = batch_no_condition[attr]+batch_with_condition[attr]
+
+
+ combined_score = super(GuidedPredictorCorrector, self)._score_fn(
+ x=joint_batch, t=torch.cat([t, t], dim=0),
+ )
+ # Split the combined score back into unconditional and conditional parts.
+ # Any batch.attr: list fields will be wrong here because of the manual concatenation above
+ # this should be ok as self._multi_corruption.corrupted_fields are always torch.Tensor
+ unconditional_score = combined_score[0]
+ conditional_score = combined_score[1]
+
+ return unconditional_score.replace(
+ **{
+ k: torch.lerp(
+ unconditional_score[k], conditional_score[k], self._guidance_scale
+ )
+ for k in self._multi_corruption.corrupted_fields
+ }
+ )
diff --git a/model/diffusion/sampling/pc_partials.py b/model/diffusion/sampling/pc_partials.py
new file mode 100644
index 0000000000000000000000000000000000000000..eac9ead898c1061d925f1ce6288078a3984a43e5
--- /dev/null
+++ b/model/diffusion/sampling/pc_partials.py
@@ -0,0 +1,21 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+from typing import Protocol
+
+from ...diffusion.corruption.corruption import Corruption
+from ...diffusion.corruption.sde_lib import ScoreFunction
+from ...diffusion.sampling.predictors import Predictor
+from ...diffusion.sampling.predictors_correctors import LangevinCorrector
+
+
+class PredictorPartial(Protocol):
+ def __call__(self, *, corruption: Corruption, score_fn: ScoreFunction | None) -> Predictor:
+ raise NotImplementedError
+
+
+class CorrectorPartial(Protocol):
+ def __call__(
+ self, *, corruption: Corruption, n_steps: int, score_fn: ScoreFunction | None
+ ) -> LangevinCorrector:
+ raise NotImplementedError
diff --git a/model/diffusion/sampling/pc_sampler.py b/model/diffusion/sampling/pc_sampler.py
new file mode 100644
index 0000000000000000000000000000000000000000..4e0af18985885570790a928e56b42f3de646c92a
--- /dev/null
+++ b/model/diffusion/sampling/pc_sampler.py
@@ -0,0 +1,279 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+from __future__ import annotations
+
+from typing import Generic, Mapping, Tuple, TypeVar
+
+import torch
+from tqdm.auto import tqdm
+
+from ...diffusion.corruption.multi_corruption import MultiCorruption, apply
+from ...diffusion.data.batched_data import BatchedData
+from ...diffusion.diffusion_module import DiffusionModule
+from ...diffusion.lightning_module import DiffusionLightningModule
+from ...diffusion.sampling.pc_partials import CorrectorPartial, PredictorPartial
+
+Diffusable = TypeVar(
+ "Diffusable", bound=BatchedData
+) # Don't use 'T' because it clashes with the 'T' for time
+SampleAndMean = Tuple[Diffusable, Diffusable]
+SampleAndMeanAndMaybeRecords = Tuple[Diffusable, Diffusable, list[Diffusable] | None]
+SampleAndMeanAndRecords = Tuple[Diffusable, Diffusable, list[Diffusable]]
+
+
+class PredictorCorrector(Generic[Diffusable]):
+ """Generates samples using predictor-corrector sampling."""
+
+ def __init__(
+ self,
+ *,
+ diffusion_module: DiffusionModule,
+ predictor_partials: dict[str, PredictorPartial] | None = None,
+ corrector_partials: dict[str, CorrectorPartial] | None = None,
+ device: torch.device,
+ n_steps_corrector: int,
+ N: int,
+ eps_t: float = 1e-3,
+ max_t: float | None = None,
+ ):
+ """
+ Args:
+ diffusion_module: diffusion module
+ predictor_partials: partials for constructing predictors. Keys are the names of the corruptions.
+ corrector_partials: partials for constructing correctors. Keys are the names of the corruptions.
+ device: device to run on
+ n_steps_corrector: number of corrector steps
+ N: number of noise levels
+ eps_t: diffusion time to stop denoising at
+ max_t: diffusion time to start denoising at. If None, defaults to the maximum diffusion time. You may want to start at T-0.01, say, for numerical stability.
+ """
+ self._diffusion_module = diffusion_module
+ self.N = N
+
+ if max_t is None:
+ max_t = self._multi_corruption.T
+ assert max_t <= self._multi_corruption.T, "Denoising cannot start from beyond T"
+
+ self._max_t = max_t
+ assert (
+ corrector_partials or predictor_partials
+ ), "Must specify at least one predictor or corrector"
+ corrector_partials = corrector_partials or {}
+ predictor_partials = predictor_partials or {}
+ if self._multi_corruption.discrete_corruptions:
+ # These all have property 'N' because they are D3PM type
+ assert set(c.N for c in self._multi_corruption.discrete_corruptions.values()) == {N} # type: ignore
+
+ self._predictors = {
+ k: v(corruption=self._multi_corruption.corruptions[k], score_fn=None)
+ for k, v in predictor_partials.items()
+ }
+
+ self._correctors = {
+ k: v(
+ corruption=self._multi_corruption.corruptions[k],
+ n_steps=n_steps_corrector,
+ score_fn=None,
+ )
+ for k, v in corrector_partials.items()
+ }
+ self._eps_t = eps_t
+ self._n_steps_corrector = n_steps_corrector
+ self._device = device
+
+ @property
+ def diffusion_module(self) -> DiffusionModule:
+ return self._diffusion_module
+
+ @property
+ def _multi_corruption(self) -> MultiCorruption:
+ return self._diffusion_module.corruption
+
+ def _score_fn(self, x: Diffusable, t: torch.Tensor) -> Diffusable:
+ return self._diffusion_module.score_fn(x, t)
+
+ @classmethod
+ def from_pl_module(cls, pl_module: DiffusionLightningModule, **kwargs) -> PredictorCorrector:
+ return cls(diffusion_module=pl_module.diffusion_module, device=pl_module.device, **kwargs)
+
+ @torch.no_grad()
+ def sample(
+ self, conditioning_data: BatchedData, mask: Mapping[str, torch.Tensor] | None = None
+ ) -> SampleAndMean:
+ """Create one sample for each of a batch of conditions.
+ Args:
+ conditioning_data: batched conditioning data. Even if you think you don't want conditioning, you still need to pass a batch of conditions
+ because the sampler uses these to determine the shapes of things to generate.
+ mask: for inpainting. Keys should be a subset of the keys in `data`. 1 indicates data that should be fixed, 0 indicates data that should be replaced with sampled values.
+ Shapes of values in `mask` must match the shapes of values in `conditioning_data`.
+ Returns:
+ (batch, mean_batch). The difference between these is that `mean_batch` has no noise added at the final denoising step.
+
+ """
+ return self._sample_maybe_record(conditioning_data, mask=mask, record=False)[:2]
+
+ @torch.no_grad()
+ def sample_with_record(
+ self, conditioning_data: BatchedData, mask: Mapping[str, torch.Tensor] | None = None
+ ) -> SampleAndMeanAndRecords:
+ """Create one sample for each of a batch of conditions.
+ Args:
+ conditioning_data: batched conditioning data. Even if you think you don't want conditioning, you still need to pass a batch of conditions
+ because the sampler uses these to determine the shapes of things to generate.
+ mask: for inpainting. Keys should be a subset of the keys in `data`. 1 indicates data that should be fixed, 0 indicates data that should be replaced with sampled values.
+ Shapes of values in `mask` must match the shapes of values in `conditioning_data`.
+ Returns:
+ (batch, mean_batch). The difference between these is that `mean_batch` has no noise added at the final denoising step.
+
+ """
+ return self._sample_maybe_record(conditioning_data, mask=mask, record=True)
+
+ @torch.no_grad()
+ def _sample_maybe_record(
+ self,
+ conditioning_data: BatchedData,
+ mask: Mapping[str, torch.Tensor] | None = None,
+ record: bool = False,
+ ) -> SampleAndMeanAndMaybeRecords:
+ """Create one sample for each of a batch of conditions.
+ Args:
+ conditioning_data: batched conditioning data. Even if you think you don't want conditioning, you still need to pass a batch of conditions
+ because the sampler uses these to determine the shapes of things to generate.
+ mask: for inpainting. Keys should be a subset of the keys in `data`. 1 indicates data that should be fixed, 0 indicates data that should be replaced with sampled values.
+ Shapes of values in `mask` must match the shapes of values in `conditioning_data`.
+ Returns:
+ (batch, mean_batch, recorded_samples, recorded_predictions).
+ The difference between the former two is that `mean_batch` has no noise added at the final denoising step.
+ The latter two are only returned if `record` is True, and contain the samples and predictions from each step of the diffusion process.
+
+ """
+ if isinstance(self._diffusion_module, torch.nn.Module):
+ self._diffusion_module.eval()
+ mask = mask or {}
+ conditioning_data = conditioning_data.to(self._device)
+ mask = {k: v.to(self._device) for k, v in mask.items()}
+ batch = _sample_prior(self._multi_corruption, conditioning_data, mask=mask)
+ return self._denoise(batch=batch, mask=mask, record=record)
+
+ @torch.no_grad()
+ def _denoise(
+ self,
+ batch: Diffusable,
+ mask: dict[str, torch.Tensor],
+ record: bool = False,
+ ) -> SampleAndMeanAndMaybeRecords:
+ """Denoise from a prior sample to a t=eps_t sample."""
+ recorded_samples = None
+ if record:
+ recorded_samples = []
+ for k in self._predictors:
+ mask.setdefault(k, None)
+ for k in self._correctors:
+ mask.setdefault(k, None)
+ mean_batch = batch.clone()
+
+ # Decreasing timesteps from T to eps_t
+ timesteps = torch.linspace(self._max_t, self._eps_t, self.N, device=self._device)
+ dt = -torch.tensor((self._max_t - self._eps_t) / (self.N - 1)).to(self._device)
+
+ for i in tqdm(range(self.N), miniters=50, mininterval=5):
+ # Set the timestep
+ t = torch.full((batch.get_batch_size(),), timesteps[i], device=self._device)
+
+ # Corrector updates.
+ if self._correctors:
+ for _ in range(self._n_steps_corrector):
+ score = self._score_fn(batch, t)
+ fns = {
+ k: corrector.step_given_score for k, corrector in self._correctors.items()
+ }
+ samples_means: dict[str, Tuple[torch.Tensor, torch.Tensor]] = apply(
+ fns=fns,
+ broadcast={"t": t, "dt": dt},
+ x=batch,
+ score=score,
+ batch_idx=self._multi_corruption._get_batch_indices(batch),
+ )
+ if record:
+ recorded_samples.append(batch.clone().to("cpu"))
+ batch, mean_batch = _mask_replace(
+ samples_means=samples_means, batch=batch, mean_batch=mean_batch, mask=mask
+ )
+
+ # Predictor updates
+ score = self._score_fn(batch, t)
+ predictor_fns = {
+ k: predictor.update_given_score for k, predictor in self._predictors.items()
+ }
+ samples_means = apply(
+ fns=predictor_fns,
+ x=batch,
+ score=score,
+ broadcast=dict(t=t, batch=batch, dt=dt),
+ batch_idx=self._multi_corruption._get_batch_indices(batch),
+ )
+ if record:
+ recorded_samples.append(batch.clone().to("cpu"))
+ batch, mean_batch = _mask_replace(
+ samples_means=samples_means, batch=batch, mean_batch=mean_batch, mask=mask
+ )
+
+ return batch, mean_batch, recorded_samples
+
+
+def _mask_replace(
+ samples_means: dict[str, Tuple[torch.Tensor, torch.Tensor]],
+ batch: BatchedData,
+ mean_batch: BatchedData,
+ mask: dict[str, torch.Tensor | None],
+) -> SampleAndMean:
+ # Apply masks
+ samples_means = apply(
+ fns={k: _mask_both for k in samples_means},
+ broadcast={},
+ sample_and_mean=samples_means,
+ mask=mask,
+ old_x=batch,
+ )
+
+ # Put the updated values in `batch` and `mean_batch`
+ batch = batch.replace(**{k: v[0] for k, v in samples_means.items()})
+ mean_batch = mean_batch.replace(**{k: v[1] for k, v in samples_means.items()})
+ return batch, mean_batch
+
+
+def _mask_both(
+ *, sample_and_mean: Tuple[torch.Tensor, torch.Tensor], old_x: torch.Tensor, mask: torch.Tensor
+) -> Tuple[torch.Tensor, torch.Tensor]:
+ return tuple(_mask(old_x=old_x, new_x=x, mask=mask) for x in sample_and_mean) # type: ignore
+
+
+def _mask(*, old_x: torch.Tensor, new_x: torch.Tensor, mask: torch.Tensor | None) -> torch.Tensor:
+ """Replace new_x with old_x where mask is 1."""
+ if mask is None:
+ return new_x
+ else:
+ return new_x.lerp(old_x, mask)
+
+
+def _sample_prior(
+ multi_corruption: MultiCorruption,
+ conditioning_data: BatchedData,
+ mask: Mapping[str, torch.Tensor] | None,
+) -> BatchedData:
+ samples = {
+ k: multi_corruption.corruptions[k]
+ .prior_sampling(
+ shape=conditioning_data[k].shape,
+ conditioning_data=conditioning_data,
+ batch_idx=conditioning_data.get_batch_idx(field_name=k),
+ )
+ .to(conditioning_data[k].device)
+ for k in multi_corruption.corruptions
+ }
+ mask = mask or {}
+ for k, msk in mask.items():
+ if k in multi_corruption.corrupted_fields:
+ samples[k].lerp_(conditioning_data[k], msk)
+ return conditioning_data.replace(**samples)
diff --git a/model/diffusion/sampling/predictors.py b/model/diffusion/sampling/predictors.py
new file mode 100644
index 0000000000000000000000000000000000000000..f2db7e63a123c678f0b3702ad3581a3c9c811aea
--- /dev/null
+++ b/model/diffusion/sampling/predictors.py
@@ -0,0 +1,169 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+"""Adapted from https://github.com/yang-song/score_sde_pytorch which is released under Apache license.
+
+Key changes:
+- Introduced batch_idx argument to work with graph-like data (e.g. molecules)
+- Introduced `..._given_score` methods so that multiple fields can be sampled at once using a shared score model. See PredictorCorrector for how this is used.
+"""
+
+import abc
+import logging
+
+import torch
+
+from ...diffusion.corruption.corruption import Corruption
+from ...diffusion.corruption.sde_lib import SDE, ScoreFunction, check_score_fn_defined
+from ...diffusion.data.batched_data import BatchedData
+from ...diffusion.sampling.predictors_correctors import SampleAndMean, Sampler
+from ...diffusion.wrapped.wrapped_sde import WrappedSDEMixin
+
+logger = logging.getLogger(__name__)
+
+
+class Predictor(Sampler):
+ """The abstract class for something that takes x_t and predicts x_{t-dt},
+ where t is diffusion timestep."""
+
+ def __init__(
+ self,
+ corruption: Corruption,
+ score_fn: ScoreFunction | None,
+ ):
+ super().__init__(corruption, score_fn=score_fn)
+
+ def update_fn(
+ self,
+ *,
+ x: torch.Tensor,
+ t: torch.Tensor,
+ dt: torch.Tensor,
+ batch_idx: torch.LongTensor,
+ batch: BatchedData | None,
+ ) -> SampleAndMean:
+ """One update of the predictor.
+
+ Args:
+ x: current state
+ t: timesteps
+ batch_idx: indicates which sample each row of x belongs to
+
+ Returns:
+ (sampled next state, mean next state)
+ """
+ check_score_fn_defined(self.score_fn, "update_given_score")
+ assert self.score_fn is not None
+ score = self.score_fn(x=x, t=t, batch_idx=batch_idx)
+ return self.update_given_score(
+ x=x, t=t, dt=dt, batch_idx=batch_idx, score=score, batch=batch
+ )
+
+ @abc.abstractmethod
+ def update_given_score(
+ self,
+ *,
+ x: torch.Tensor,
+ t: torch.Tensor,
+ dt: torch.Tensor,
+ batch_idx: torch.LongTensor,
+ score: torch.Tensor,
+ batch: BatchedData | None,
+ ) -> SampleAndMean:
+ pass
+
+
+class AncestralSamplingPredictor(Predictor):
+ """Suitable for all linear SDEs.
+
+ This predictor is derived by converting the score prediction to a prediction of x_0 given x_t, and then
+ sampling from the conditional distribution of x_{t-dt} given x_0 and x_t according to the corruption process.
+ It corresponds to equation (47) in Song et al. for VESDE (https://openreview.net/forum?id=PxTIG12RRHS)
+ and equation (7) in Ho et al. for VPSDE (https://arxiv.org/abs/2006.11239)
+
+ In more detail: suppose the SDE has marginals x_t ~ N(alpha_t *x_0, sigma_t**2)
+
+ We estimate x_0 as follows:
+ x_0 \approx (x_t + sigma_t^2 * score) / alpha_t
+
+ For any s < t, the forward corruption process implies that
+ x_t| x_s ~ N(alpha_t/alpha_s * x_s, sigma_t^2 - sigma_s^2 * alpha_t^2 / alpha_s^2)
+
+ Now go away and do some algebra to get the mean and variance of x_s given x_t
+ and x_0, and you will get the coefficients in the `update_given_score` method below.
+
+ """
+
+ def update_given_score(
+ self,
+ *,
+ x: torch.Tensor,
+ t: torch.Tensor,
+ dt: torch.Tensor,
+ batch_idx: torch.LongTensor,
+ score: torch.Tensor,
+ batch: BatchedData | None,
+ ) -> SampleAndMean:
+ x_coeff, score_coeff, std = self._get_coeffs(
+ x=x,
+ t=t,
+ dt=dt,
+ batch_idx=batch_idx,
+ batch=batch,
+ )
+ # Sample random noise.
+ z = torch.randn_like(x_coeff)
+
+ mean = x_coeff * x + score_coeff * score
+ sample = mean + std * z
+
+ return sample, mean
+
+ def _get_coeffs(self, x, t, dt, batch_idx, batch):
+ """
+ Compute coefficients for ancestral sampling.
+ This is in a separate method to make it easier to test."""
+ sde = self.corruption
+ assert isinstance(sde, SDE)
+
+ # Previous timestep
+ s = t + dt
+
+ alpha_t, sigma_t = sde.mean_coeff_and_std(x=x, t=t, batch_idx=batch_idx, batch=batch)
+ if batch_idx is None:
+ is_time_zero = s <= 0
+ else:
+ is_time_zero = s[batch_idx] <= 0
+ alpha_s, sigma_s = sde.mean_coeff_and_std(x=x, t=s, batch_idx=batch_idx, batch=batch)
+ sigma_s[is_time_zero] = 0
+
+ # If you are trying to match this up with algebra in papers, it may help to
+ # notice that for VPSDE, sigma2_t_given_s == 1 - alpha_t_given_s**2, except
+ # that alpha_t_given_s**2 is clipped.
+ sigma2_t_given_s = sigma_t**2 - sigma_s**2 * alpha_t**2 / alpha_s**2
+ sigma_t_given_s = torch.sqrt(sigma2_t_given_s)
+ std = sigma_t_given_s * sigma_s / sigma_t
+
+ # Clip alpha_t_given_s so that we do not divide by zero.
+ min_alpha_t_given_s = 0.001
+ alpha_t_given_s = alpha_t / alpha_s
+ if torch.any(alpha_t_given_s < min_alpha_t_given_s):
+ # If this warning is raised, you probably should change something: either modify your noise schedule
+ # so that the diffusion coefficient does not blow up near sde.T, or only denoise from sde.T - eps,
+ # rather than sde.T.
+ logger.warning(
+ f"Clipping alpha_t_given_s to {min_alpha_t_given_s} to avoid divide-by-zero. You should probably change something else to avoid this."
+ )
+ alpha_t_given_s = torch.clip(alpha_t_given_s, min_alpha_t_given_s, 1)
+
+ score_coeff = sigma2_t_given_s / alpha_t_given_s
+
+ x_coeff = 1.0 / alpha_t_given_s
+
+ std[is_time_zero] = 0
+
+ return x_coeff, score_coeff, std
+
+ @classmethod
+ def is_compatible(cls, corruption: Corruption) -> bool:
+ return super().is_compatible(corruption) and not isinstance(corruption, WrappedSDEMixin)
diff --git a/model/diffusion/sampling/predictors_correctors.py b/model/diffusion/sampling/predictors_correctors.py
new file mode 100644
index 0000000000000000000000000000000000000000..4cc8c54d1ee504a872d0d3aa1323b94c9df893f3
--- /dev/null
+++ b/model/diffusion/sampling/predictors_correctors.py
@@ -0,0 +1,124 @@
+# Copyright 2020 The Google Research Authors.
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+# Adapted from https://github.com/yang-song/score_sde_pytorch which is released under Apache license.
+
+# Key changes:
+# - Introduced batch_idx argument to work with graph-like data (e.g. molecules)
+# - Introduced `..._given_score` methods so that multiple fields can be sampled at once using a shared score model. See PredictorCorrector for how this is used.
+
+import abc
+
+import torch
+from torch_scatter import scatter_add
+
+from ...diffusion.corruption.corruption import maybe_expand
+from ...diffusion.corruption.sde_lib import (
+ VESDE,
+ VPSDE,
+ BaseVPSDE,
+ Corruption,
+ ScoreFunction,
+)
+from ...diffusion.exceptions import IncompatibleSampler
+from ...diffusion.wrapped.wrapped_sde import WrappedSDEMixin
+
+SampleAndMean = tuple[torch.Tensor, torch.Tensor]
+
+
+class Sampler(abc.ABC):
+ def __init__(self, corruption: Corruption, score_fn: ScoreFunction | None):
+ if not self.is_compatible(corruption):
+ raise IncompatibleSampler(
+ f"{self.__class__.__name__} is not compatible with {corruption}"
+ )
+ self.corruption = corruption
+ self.score_fn = score_fn
+
+ @classmethod
+ def is_compatible(cls, corruption: Corruption) -> bool:
+ return True
+
+
+class LangevinCorrector(Sampler):
+ def __init__(
+ self,
+ corruption: Corruption,
+ score_fn: ScoreFunction | None,
+ n_steps: int,
+ snr: float = 0.2,
+ max_step_size: float = 1.0,
+ ):
+ """The Langevin corrector.
+
+ Args:
+ corruption: corruption process
+ score_fn: score function
+ n_steps: number of Langevin steps at each noise level
+ snr: signal-to-noise ratio
+ max_step_size: largest coefficient that the score can be multiplied by for each Langevin step.
+ """
+ super().__init__(corruption=corruption, score_fn=score_fn)
+ self.n_steps = n_steps
+ self.snr = snr
+ self.max_step_size = torch.tensor(max_step_size)
+
+ @classmethod
+ def is_compatible(cls, corruption: Corruption):
+ return (
+ isinstance(corruption, (VESDE, BaseVPSDE))
+ and super().is_compatible(corruption)
+ and not isinstance(corruption, WrappedSDEMixin)
+ )
+
+ def update_fn(self, *, x, t, batch_idx, dt: torch.Tensor) -> SampleAndMean:
+ assert self.score_fn is not None, "Did you mean to use step_given_score?"
+ for _ in range(self.n_steps):
+ score = self.score_fn(x, t, batch_idx)
+ x, x_mean = self.step_given_score(x=x, batch_idx=batch_idx, score=score, t=t, dt=dt)
+
+ return x, x_mean
+
+ def get_alpha(self, t: torch.FloatTensor, dt: torch.FloatTensor) -> torch.Tensor:
+ sde = self.corruption
+
+ if isinstance(sde, VPSDE):
+ alpha_bar = sde._marginal_mean_coeff(t) ** 2
+ alpha_bar_before = sde._marginal_mean_coeff(t + dt) ** 2
+ alpha = alpha_bar / alpha_bar_before
+ else:
+ alpha = torch.ones_like(t)
+ return alpha
+
+ def step_given_score(
+ self, *, x, batch_idx: torch.LongTensor | None, score, t: torch.Tensor, dt: torch.Tensor
+ ) -> SampleAndMean:
+ alpha = self.get_alpha(t, dt=dt)
+ snr = self.snr
+ noise = torch.randn_like(score)
+ grad_norm_square = torch.square(score).reshape(score.shape[0], -1).sum(dim=1)
+ noise_norm_square = torch.square(noise).reshape(noise.shape[0], -1).sum(dim=1)
+ if batch_idx is None:
+ grad_norm = grad_norm_square.sqrt().mean()
+ noise_norm = noise_norm_square.sqrt().mean()
+ else:
+ grad_norm = torch.sqrt(scatter_add(grad_norm_square, dim=-1, index=batch_idx)).mean()
+
+ noise_norm = torch.sqrt(scatter_add(noise_norm_square, dim=-1, index=batch_idx)).mean()
+
+ # If gradient is zero (i.e., we are sampling from an improper distribution that's flat over the whole of R^n)
+ # the step_size blows up. Clip step_size to avoid this.
+ # The EGNN reports zero scores when there are no edges between nodes.
+ step_size = (snr * noise_norm / grad_norm) ** 2 * 2 * alpha
+ step_size = torch.minimum(step_size, self.max_step_size)
+ step_size[grad_norm == 0, :] = self.max_step_size
+
+ # Expand step size to batch structure (score and noise have the same shape).
+ step_size = maybe_expand(step_size, batch_idx, score)
+
+ # Perform update, using custom update for SO(3) diffusion on frames.
+ mean = x + step_size * score
+ x = mean + torch.sqrt(step_size * 2) * noise
+
+ return x, mean
diff --git a/model/diffusion/score_models/__init__.py b/model/diffusion/score_models/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/model/diffusion/score_models/base.py b/model/diffusion/score_models/base.py
new file mode 100644
index 0000000000000000000000000000000000000000..2bfc1a38ccdbeb8fac16513c7404b3ab8851558f
--- /dev/null
+++ b/model/diffusion/score_models/base.py
@@ -0,0 +1,23 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+import abc
+from typing import Generic, TypeVar
+
+import torch
+
+from ...diffusion.data.batched_data import BatchedData
+
+Diffusable = TypeVar("Diffusable", bound=BatchedData)
+
+
+class ScoreModel(torch.nn.Module, Generic[Diffusable], abc.ABC):
+ """Abstract base class for score models."""
+
+ @abc.abstractmethod
+ def forward(self, x: Diffusable, t: torch.Tensor) -> Diffusable:
+ """Args:
+ x: batch of noisy data
+ t: timestep. Shape (batch_size, 1)
+ """
+ ...
diff --git a/model/diffusion/timestep_samplers.py b/model/diffusion/timestep_samplers.py
new file mode 100644
index 0000000000000000000000000000000000000000..f0ae868d1e76f096fb3820ec571145ce25a7d1dc
--- /dev/null
+++ b/model/diffusion/timestep_samplers.py
@@ -0,0 +1,39 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+from typing import Protocol
+
+import torch
+
+from ..diffusion.corruption.sde_lib import SDE
+
+
+class TimestepSampler(Protocol):
+ min_t: float
+ max_t: float
+
+ def __call__(self, batch_size: int, device: torch.device) -> torch.FloatTensor:
+ raise NotImplementedError
+
+
+class UniformTimestepSampler:
+ """Samples diffusion timesteps uniformly over the training time."""
+
+ def __init__(
+ self,
+ *,
+ min_t: float,
+ max_t: float,
+ ):
+ """Initializes the sampler.
+
+ Args:
+ min_t (float): Smallest timestep that will be seen during training.
+ max_t (float): Largest timestep that will be seen during training.
+ """
+ super().__init__()
+ self.min_t = min_t
+ self.max_t = max_t
+
+ def __call__(self, batch_size: int, device: torch.device) -> torch.FloatTensor:
+ return torch.rand(batch_size, device=device) * (self.max_t - self.min_t) + self.min_t
diff --git a/model/diffusion/training/__init__.py b/model/diffusion/training/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/model/diffusion/training/field_loss.py b/model/diffusion/training/field_loss.py
new file mode 100644
index 0000000000000000000000000000000000000000..dff301881f87d8a1cfbb4684579562fb5cbd09ee
--- /dev/null
+++ b/model/diffusion/training/field_loss.py
@@ -0,0 +1,198 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+from typing import Literal, Protocol
+
+import torch
+from torch_scatter import scatter
+
+from ...diffusion.corruption.corruption import Corruption
+from ...diffusion.corruption.sde_lib import maybe_expand
+from ...diffusion.d3pm.d3pm import compute_kl_reverse_process
+from ...diffusion.data.batched_data import BatchedData
+from ...diffusion.discrete_time import to_discrete_time
+from ...diffusion.model_target import ModelTarget
+
+
+def compute_noise_given_sample_and_corruption(
+ x: torch.Tensor,
+ x_noisy: torch.Tensor,
+ corruption: Corruption,
+ t: torch.Tensor,
+ batch_idx: torch.LongTensor | None,
+ batch: BatchedData,
+) -> torch.Tensor:
+ """
+ Recover the (unit-Gaussian-distributed) raw noise that was used to corrupt a batch of samples.
+ We first obtain the mean and std of the noisy samples from the corruption via `t` and the clean batch.
+ Then we solve:
+ x_noisy = x_mean + noise * std w.r.t. `noise`:
+ noise = (x_noisy - x_mean) / std
+ """
+ x_mean, std = corruption.marginal_prob(
+ x,
+ t=t,
+ batch_idx=batch_idx,
+ batch=batch,
+ )
+ return (x_noisy - x_mean) / std
+
+
+class FieldLoss(Protocol):
+ """Loss function for a single field. Because loss functions are defined different ways in different papers,
+ we pass loads of keyword arguments. Each loss function will only use a subset of these arguments.
+ """
+
+ def __call__(
+ self,
+ *,
+ corruption: Corruption,
+ score_model_output: torch.Tensor,
+ t: torch.Tensor,
+ batch_idx: torch.LongTensor | None,
+ batch_size: int,
+ x: torch.Tensor,
+ noisy_x: torch.Tensor,
+ reduce: Literal["sum", "mean"],
+ batch: BatchedData,
+ ) -> torch.Tensor:
+ """Calculate loss per sample for a single field. Returns a loss tensor of shape (batch_size,)."""
+ pass
+
+
+def denoising_score_matching(
+ *,
+ corruption: Corruption,
+ score_model_output: torch.Tensor,
+ t: torch.Tensor,
+ batch_idx: torch.LongTensor | None,
+ batch_size: int,
+ x: torch.Tensor,
+ noisy_x: torch.Tensor,
+ reduce: Literal["sum", "mean"],
+ batch: BatchedData,
+ model_target: ModelTarget,
+ node_is_unmasked: torch.LongTensor | None = None,
+ **_,
+) -> torch.Tensor:
+ """Mean square error in predicting raw noise, optionally reweighted."""
+ assert score_model_output.ndim >= 2
+ model_target = ModelTarget(model_target) # in case str was passed
+
+ losses = get_losses(
+ corruption=corruption,
+ score_model_output=score_model_output,
+ t=t,
+ batch_idx=batch_idx,
+ x=x,
+ noisy_x=noisy_x,
+ batch=batch,
+ model_target=model_target,
+ )
+
+ if node_is_unmasked is not None:
+ losses = node_is_unmasked.unsqueeze(-1) * losses # Apply masking.
+ original_reduce = reduce
+ reduce = "sum" # We sum first and handle the division by nodes_per_sample for the mean manually later.
+
+ loss_per_sample = aggregate_per_sample(losses, batch_idx, reduce=reduce, batch_size=batch_size)
+
+ if (node_is_unmasked is not None) and (original_reduce == "mean"):
+ nodes_per_sample = scatter(node_is_unmasked, batch_idx, dim=0, reduce="sum")
+ loss_per_sample /= nodes_per_sample
+
+ return loss_per_sample
+
+
+def get_losses(
+ corruption: Corruption,
+ score_model_output: torch.Tensor,
+ t: torch.Tensor,
+ batch_idx: torch.LongTensor | None,
+ x: torch.Tensor,
+ noisy_x: torch.Tensor,
+ batch: BatchedData,
+ model_target: ModelTarget,
+) -> torch.Tensor:
+ if model_target == ModelTarget.score_times_std:
+ raw_noise = compute_noise_given_sample_and_corruption(
+ x=x, x_noisy=noisy_x, corruption=corruption, t=t, batch_idx=batch_idx, batch=batch
+ )
+ target = -raw_noise
+ losses = (score_model_output - target).square()
+ else:
+ raise ValueError(f"Unknown model_target {model_target}")
+ return losses
+
+
+def aggregate_per_sample(
+ loss_per_row: torch.Tensor,
+ batch_idx: torch.Tensor | None,
+ reduce: Literal["sum", "mean"],
+ batch_size: int,
+):
+ """
+ Aggregate (potentially) batched input tensor to get a scalar for each sample in the batch.
+ E.g., (num_atoms, d1, d2, ..., dn) -> (batch_size, d1, d2, ..., dn) -> (batch_size,),
+ where the first aggregation only happens when batch_idx is provided.
+
+ Args:
+ loss_per_row: shape (num_nodes, any_more_dims). May contain multiple nodes per sample.
+ batch_idx: shape (num_nodes,). Indicates which sample each row belongs to. If not provided,
+ then we assume the first dimension is the batch dimension.
+ reduce: determines how to aggregate over nodes within each sample. (Aggregation over samples
+ and within dims for one node is always mean.)
+ batch_size: number of samples in the batch.
+
+ Returns:
+ Scalar for each sample, shape (batch_size,).
+
+ """
+ # Sum over all but 0th dimension.
+ loss_per_row = torch.mean(loss_per_row.reshape(loss_per_row.shape[0], -1), dim=1)
+
+ if batch_idx is None:
+ # First dimension is batch dimension. In this case 'reduce' is ignored.
+ loss_per_sample = loss_per_row
+ else:
+ # Aggregate over nodes within each sample.
+ loss_per_sample = scatter(
+ src=loss_per_row,
+ index=batch_idx,
+ dim_size=batch_size,
+ reduce=reduce,
+ )
+ return loss_per_sample
+
+
+def d3pm_loss(
+ *,
+ corruption: Corruption,
+ score_model_output: torch.Tensor,
+ t: torch.Tensor,
+ batch_idx: torch.LongTensor | None,
+ batch_size: int,
+ x: torch.Tensor,
+ noisy_x: torch.Tensor,
+ reduce: Literal["sum", "mean"],
+ d3pm_hybrid_lambda: float = 0.0,
+ **_,
+) -> torch.Tensor:
+ assert hasattr(corruption, "N") # mypy
+ assert hasattr(corruption, "_to_zero_based") # mypy
+ assert hasattr(corruption, "d3pm") # mypy
+ t = maybe_expand(to_discrete_time(t, N=corruption.N, T=corruption.T), batch_idx)
+ metrics_dict = compute_kl_reverse_process(
+ corruption._to_zero_based(x.long()),
+ t,
+ diffusion=corruption.d3pm,
+ log_space=True,
+ denoise_fn=lambda targets, timestep: score_model_output,
+ hybrid_lambda=d3pm_hybrid_lambda,
+ x_t_plus_1=corruption._to_zero_based(noisy_x.long()),
+ )
+ loss = metrics_dict.pop("loss")
+ loss_per_structure = aggregate_per_sample(
+ loss, batch_idx=batch_idx, reduce=reduce, batch_size=batch_size
+ )
+ return loss_per_structure
diff --git a/model/diffusion/training/metrics.py b/model/diffusion/training/metrics.py
new file mode 100644
index 0000000000000000000000000000000000000000..83bf975c552b3f7c36808603533cc32e7e70017e
--- /dev/null
+++ b/model/diffusion/training/metrics.py
@@ -0,0 +1,136 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+from typing import Dict, Iterable, Protocol
+
+import torch
+from torch_scatter import scatter
+
+from ...diffusion.corruption.multi_corruption import MultiCorruption
+from ...diffusion.data.batched_data import BatchedData
+from ...diffusion.score_models.base import Diffusable
+
+
+class Metric(Protocol):
+ """
+ Computes a metric to be logged during training.
+ Each metric must have a name which is used as a prefix for the metric in the log.
+ """
+
+ name: str
+
+ def __call__(
+ self,
+ *,
+ loss_per_sample_per_field: Dict[str, torch.Tensor],
+ multi_corruption: MultiCorruption,
+ score_model_output: Diffusable,
+ t: torch.Tensor,
+ batch_idx: Dict[str, torch.LongTensor],
+ batch: BatchedData,
+ noisy_batch: BatchedData,
+ ) -> Dict[str, torch.Tensor]:
+ """
+ Computes a metric to be logged during training. Useful, e.g., for plotting loss over time.
+
+ Args:
+ loss_per_sample_per_field: Dict[str, torch.Tensor], where each tensor has shape (batch_size,).
+ multi_corruption: MultiCorruption
+ score_model_output: the output produced by the model per field.
+ t: shape (batch_size,). Time for each element in the loss.
+ batch_idx: Dict[str, torch.LongTensor]: batch indices per field
+ batch: BatchedData: the clean (un-perturbed) batched data
+ noisy_batch: BatchedData: the corrupted batched data
+ """
+ pass
+
+
+def loss_per_time_bin(
+ loss_per_sample: torch.Tensor, t: torch.Tensor, bins: torch.Tensor
+) -> torch.Tensor:
+ """
+ Aggregate loss per bin. Useful for plotting loss over time.
+
+ Args:
+ loss_per_sample: shape (batch_size,). Loss for each sample.
+ t: shape (batch_size,). Time for each element in the loss.
+ bins: shape (num_bins,). Upper boundaries of the time bins.
+ Returns:
+ avg_loss_per_bin: shape (num_bins,). Average loss per time bin.
+ """
+ bin_per_element = torch.bucketize(t, bins)
+ avg_loss_per_bin = scatter(
+ src=loss_per_sample, index=bin_per_element, dim_size=bins.shape[0], reduce="mean"
+ )
+ return avg_loss_per_bin
+
+
+class LossPerTimeBin(Metric):
+ name = "loss_per_time_bin"
+
+ def __init__(self, t_min: float = 0.0, t_max: float = 1.0, num_bins: int = 10):
+ self.bins = torch.linspace(t_min, t_max, num_bins + 1)
+
+ def __call__(
+ self,
+ *,
+ loss_per_sample_per_field: Dict[str, torch.Tensor],
+ t: torch.Tensor,
+ **_,
+ ) -> Dict[str, torch.Tensor]:
+ """
+ Compute loss bins per diffusion time bin. Useful for plotting loss over diffusion time.
+ """
+ metrics_dict = {}
+ for k, v in loss_per_sample_per_field.items():
+ assert v.shape == t.shape
+
+ # first bin is always empty because no time is less than t_min, so we skip it
+ avg_loss_per_bin = loss_per_time_bin(
+ loss_per_sample_per_field[k],
+ t,
+ bins=self.bins.to(loss_per_sample_per_field[k].device)[1:],
+ )
+ metrics_dict.update(
+ {
+ f"{k}_{self.bins[ix]:.2f}-{self.bins[ix + 1]:.2f}": avg_loss_per_bin[ix]
+ for ix in range(len(avg_loss_per_bin))
+ if avg_loss_per_bin[ix] > 0.0
+ }
+ )
+ return metrics_dict
+
+
+class MetricsCalculator:
+ """
+ Computes a set of metrics to be logged during training.
+ """
+
+ def __init__(self, metric_fns: Iterable[Metric]):
+ self.metric_fns = metric_fns
+
+ def __call__(
+ self,
+ *,
+ loss_per_sample_per_field: Dict[str, torch.Tensor],
+ multi_corruption: MultiCorruption,
+ score_model_output: torch.Tensor,
+ t: torch.Tensor,
+ batch_idx: Dict[str, torch.LongTensor],
+ batch: BatchedData,
+ noisy_batch: BatchedData,
+ ) -> Dict[str, torch.Tensor]:
+ metrics_dict = {}
+ for metric_fn in self.metric_fns:
+ _metrics_dict = metric_fn(
+ loss_per_sample_per_field=loss_per_sample_per_field,
+ multi_corruption=multi_corruption,
+ score_model_output=score_model_output,
+ t=t,
+ batch_idx=batch_idx,
+ batch=batch,
+ noisy_batch=noisy_batch,
+ )
+ # prepend metric name to each metric
+ metrics_dict.update({f"{metric_fn.name}_{k}": v for k, v in _metrics_dict.items()})
+ return metrics_dict
diff --git a/model/diffusion/training/utils.py b/model/diffusion/training/utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..d1f7647361af4153b278453f29b29d715a44370b
--- /dev/null
+++ b/model/diffusion/training/utils.py
@@ -0,0 +1,32 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+from typing import Iterable, Union
+
+import torch
+
+
+def get_grad_norm(
+ parameters: Union[torch.Tensor, Iterable[torch.Tensor]], norm_type: float = 2.0
+) -> torch.Tensor:
+ """
+ Adapted from: https://pytorch.org/docs/stable/_modules/torch/nn/utils/clip_grad.html#clip_grad_norm_
+ """
+
+ if isinstance(parameters, torch.Tensor):
+ parameters = [parameters]
+ parameters = [p for p in parameters if p.grad is not None]
+
+ norm_type = float(norm_type)
+
+ if len(parameters) == 0:
+ return torch.tensor(0.0)
+
+ device = parameters[0].grad.device
+
+ total_norm = torch.norm(
+ torch.stack([torch.norm(p.grad.detach(), norm_type).to(device) for p in parameters]),
+ norm_type,
+ )
+
+ return total_norm
diff --git a/model/diffusion/wrapped/__init__.py b/model/diffusion/wrapped/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/model/diffusion/wrapped/wrapped_normal_loss.py b/model/diffusion/wrapped/wrapped_normal_loss.py
new file mode 100644
index 0000000000000000000000000000000000000000..1d62a133b725f303621ef6324a5311324a899059
--- /dev/null
+++ b/model/diffusion/wrapped/wrapped_normal_loss.py
@@ -0,0 +1,119 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+from typing import Literal, Optional
+
+import torch
+
+from ...diffusion.corruption.sde_lib import SDE, maybe_expand
+from ...diffusion.data.batched_data import BatchedData
+from ...diffusion.training.field_loss import aggregate_per_sample
+
+
+def get_pbc_offsets(pbc: torch.Tensor, max_offset_integer: int = 3) -> torch.Tensor:
+ """Build the Cartesian product of integer offsets of the periodic boundary. That is, if dim=3 and max_offset_integer=1 we build the (2*1 + 1)^3 = 27
+ possible combinations of the Cartesian product of (i,j,k) for i,j,k in -max_offset_integer, ..., max_offset_integer. Then, we construct
+ the tensor of integer offsets of the pbc vectors, i.e., L_{ijk} = row_stack([i * l_1, j * l_2, k * l_3]).
+
+ Args:
+ pbc (torch.Tensor, [batch_size, dim, dim]): The input pbc matrix.
+ max_offset_integer (int): The maximum integer offset per dimension to consider for the Cartesian product. Defaults to 3.
+
+ Returns:
+ torch.Tensor, [batch_size, (2 * max_offset_integer + 1)^dim, dim]: The tensor containing the integer offsets of the pbc vectors.
+ """
+ offset_range = torch.arange(-max_offset_integer, max_offset_integer + 1, device=pbc.device)
+ meshgrid = torch.stack(
+ torch.meshgrid(offset_range, offset_range, offset_range, indexing="xy"), dim=-1
+ )
+ offset = (pbc[:, None, None, None] * meshgrid[None, :, :, :, :, None]).sum(-2)
+ pbc_offset_per_molecule = offset.reshape(pbc.shape[0], -1, 3)
+ return pbc_offset_per_molecule
+
+
+def wrapped_normal_score(
+ x: torch.Tensor,
+ mean: torch.Tensor,
+ wrapping_boundary: torch.Tensor,
+ variance_diag: torch.Tensor,
+ batch: torch.Tensor,
+ max_offset_integer: int = 3,
+) -> torch.Tensor:
+ """Approximate the the score of a 3D wrapped normal distribution with diagonal covariance matrix w.r.t. x via a truncated sum.
+ See docstring of `wrapped_normal_score` for details about the arguments
+
+ Args:
+ x (torch.Tensor, [num_atoms, dim])
+ mean (torch.Tensor, [num_atoms, dim])
+ wrapping_boundary (torch.Tensor, [num_molecules, dim, dim])
+ variance_diag (torch.Tensor, [num_atoms,])
+ batch (torch.Tensor, [num_atoms, ])
+ max_offset_integer (int), Defaults to 3.
+
+ Returns:
+ torch.Tensor, [num_atoms, dim]: The approximated score of the wrapped normal distribution.
+ """
+ offset_add = get_pbc_offsets(
+ wrapping_boundary,
+ max_offset_integer,
+ )
+ diffs_k = (x - mean)[:, None] + offset_add[batch]
+ dists_sqr_k = diffs_k.pow(2).sum(-1)
+ score_softmax = torch.softmax(-dists_sqr_k / (2 * variance_diag[:, None]), dim=-1)
+ score = -(score_softmax[:, :, None] * diffs_k).sum((-2)) / (variance_diag[:, None])
+ return score
+
+
+def wrapped_normal_loss(
+ *,
+ corruption: SDE,
+ score_model_output: torch.Tensor,
+ t: torch.Tensor,
+ batch_idx: Optional[torch.LongTensor],
+ batch_size: int,
+ x: torch.Tensor,
+ noisy_x: torch.Tensor,
+ reduce: Literal["sum", "mean"],
+ batch: BatchedData,
+ **_
+) -> torch.Tensor:
+ """Compute the loss for a wrapped normal distribution.
+ Compares the score of the wrapped normal distribution to the score of the score model.
+ """
+ assert len(t) == batch_size
+ _, std = corruption.marginal_prob(
+ x=torch.zeros((x.shape[0], 1), device=t.device),
+ t=t,
+ batch_idx=batch_idx,
+ batch=batch,
+ ) # std does not depend on x
+
+ pred: torch.Tensor = score_model_output
+ if pred.ndim != 2:
+ raise NotImplementedError
+
+ assert hasattr(
+ corruption, "wrapping_boundary"
+ ), "SDE must be a WrappedSDE, i.e., must have a wrapping boundary."
+ wrapping_boundary = corruption.wrapping_boundary
+ # Scaled identity matrix, i.e., in each dimension we wrap at `wrapping_boundary`.
+ wrapping_boundary = wrapping_boundary * torch.eye(x.shape[-1], device=t.device)[None].expand(
+ batch_size, -1, -1
+ )
+
+ # We multiply the score by the standard deviation because we don't use raw_noise here; raw_noise is -score * std, i.e., we multiply the score by std.
+ target = (
+ wrapped_normal_score(
+ x=noisy_x,
+ mean=x,
+ wrapping_boundary=wrapping_boundary,
+ variance_diag=std.squeeze() ** 2,
+ batch=batch_idx,
+ )
+ * std
+ )
+ delta = target - pred
+
+ losses = delta.square()
+
+ return aggregate_per_sample(losses, batch_idx, reduce=reduce, batch_size=batch_size)
diff --git a/model/diffusion/wrapped/wrapped_predictors_correctors.py b/model/diffusion/wrapped/wrapped_predictors_correctors.py
new file mode 100644
index 0000000000000000000000000000000000000000..53aa57a023b254613fd6ecdeefa799230bbec23c
--- /dev/null
+++ b/model/diffusion/wrapped/wrapped_predictors_correctors.py
@@ -0,0 +1,89 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+from typing import Optional, Tuple
+
+import torch
+
+from ..sampling import predictors_correctors as pc
+from ...diffusion.corruption import sde_lib
+from ...diffusion.corruption.corruption import Corruption
+from ...diffusion.data.batched_data import BatchedData
+from ...diffusion.exceptions import IncompatibleSampler
+from ...diffusion.sampling import predictors
+from ...diffusion.wrapped.wrapped_sde import WrappedSDEMixin
+
+# importing SampleAndMean does not work because of circular imports, so we have to redefine it here.
+SampleAndMean = Tuple[torch.Tensor, torch.Tensor]
+
+
+class WrappedPredictorMixin:
+ """A mixin for wrapping the predictor in a WrappedSDE."""
+
+ def update_given_score(
+ self,
+ *,
+ x: torch.Tensor,
+ t: torch.Tensor,
+ dt: torch.Tensor,
+ batch_idx: torch.LongTensor,
+ score: torch.Tensor,
+ batch: Optional[BatchedData],
+ ) -> SampleAndMean:
+ # mypy
+ assert isinstance(self, predictors.Predictor)
+ _super = super()
+ assert hasattr(_super, "update_given_score")
+ assert hasattr(self, "corruption")
+ if not hasattr(self.corruption, "wrap"):
+ raise IncompatibleSampler(
+ f"{self.__class__.__name__} is not compatible with {self.corruption}."
+ )
+
+ sample, mean = _super.update_given_score(
+ x=x, t=t, dt=dt, batch_idx=batch_idx, score=score, batch=batch
+ )
+ return self.corruption.wrap(sample), self.corruption.wrap(mean)
+
+
+class WrappedCorrectorMixin:
+ """A mixin for wrapping the corrector in a WrappedSDE."""
+
+ def step_given_score(
+ self,
+ *,
+ x: torch.Tensor,
+ batch_idx: torch.LongTensor,
+ score: torch.Tensor,
+ t: torch.Tensor,
+ dt: torch.Tensor,
+ ) -> SampleAndMean:
+ # mypy
+ assert isinstance(self, pc.LangevinCorrector)
+ _super = super()
+ assert hasattr(_super, "step_given_score")
+ assert hasattr(self, "corruption") and hasattr(self.corruption, "wrap")
+ if not hasattr(self.corruption, "wrap"):
+ raise IncompatibleSampler(
+ f"{self.__class__.__name__} is not compatible with {self.corruption}."
+ )
+ sample, mean = _super.step_given_score(x=x, score=score, t=t, batch_idx=batch_idx, dt=dt)
+ return self.corruption.wrap(sample), self.corruption.wrap(mean)
+
+
+class WrappedAncestralSamplingPredictor(
+ WrappedPredictorMixin, predictors.AncestralSamplingPredictor
+):
+ @classmethod
+ def is_compatible(cls, corruption: Corruption):
+ return isinstance(corruption, (sde_lib.VPSDE, sde_lib.VESDE)) and isinstance(
+ corruption, WrappedSDEMixin
+ )
+
+
+class WrappedLangevinCorrector(WrappedCorrectorMixin, pc.LangevinCorrector):
+ @classmethod
+ def is_compatible(cls, corruption: Corruption):
+ return isinstance(corruption, (sde_lib.VPSDE, sde_lib.VESDE)) and isinstance(
+ corruption, WrappedSDEMixin
+ )
diff --git a/model/diffusion/wrapped/wrapped_sde.py b/model/diffusion/wrapped/wrapped_sde.py
new file mode 100644
index 0000000000000000000000000000000000000000..3745129dadc981c90248f2cdf64957080f1f92eb
--- /dev/null
+++ b/model/diffusion/wrapped/wrapped_sde.py
@@ -0,0 +1,82 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+from typing import Optional, Tuple, Union
+
+import torch
+
+from ...diffusion.corruption.sde_lib import SDE, VESDE, VPSDE
+from ...diffusion.data.batched_data import BatchedData
+
+B = Optional[torch.LongTensor]
+
+
+def wrap_at_boundary(x: torch.Tensor, wrapping_boundary: float) -> torch.Tensor:
+ """Wrap x at the boundary given by wrapping_boundary.
+ Args:
+ x: tensor of shape (batch_size, dim)
+ wrapping_boundary: float): wrap at [0, wrapping_boundary] in all dimensions.
+ Returns:
+ wrapped_x: tensor of shape (batch_size, dim)
+ """
+ return torch.remainder(
+ x, wrapping_boundary
+ ) # remainder is the same as mod, but works with negative numbers.
+
+
+class WrappedSDEMixin:
+ def sample_marginal(
+ self,
+ x: torch.Tensor,
+ t: torch.Tensor,
+ batch_idx: torch.LongTensor = None,
+ batch: Optional[BatchedData] = None,
+ ) -> torch.Tensor:
+ _super = super()
+ assert (
+ isinstance(self, SDE)
+ and hasattr(_super, "sample_marginal")
+ and hasattr(self, "wrapping_boundary")
+ )
+ if (x > self.wrapping_boundary).any() or (x < 0).any():
+ # Values outside the wrapping boundary are valid in principle, but could point to an issue in the data preprocessing,
+ # as typically we assume that the input data is inside the wrapping boundary (e.g., angles between 0 and 2*pi).
+ print("Warning: Wrapped SDE has received input outside of the wrapping boundary.")
+ noisy_x = _super.sample_marginal(x=x, t=t, batch_idx=batch_idx, batch=batch)
+ return self.wrap(noisy_x)
+
+ def prior_sampling(
+ self,
+ shape: Union[torch.Size, Tuple],
+ conditioning_data: Optional[BatchedData] = None,
+ batch_idx: B = None,
+ ) -> torch.Tensor:
+ _super = super()
+ assert isinstance(self, SDE) and hasattr(_super, "prior_sampling")
+ return self.wrap(_super.prior_sampling(shape=shape, conditioning_data=conditioning_data))
+
+ def wrap(self, x):
+ assert isinstance(self, SDE) and hasattr(self, "wrapping_boundary")
+ return wrap_at_boundary(x, self.wrapping_boundary)
+
+
+class WrappedVESDE(WrappedSDEMixin, VESDE):
+ def __init__(
+ self,
+ wrapping_boundary: float = 1.0,
+ sigma_min: float = 0.01,
+ sigma_max: float = 50.0,
+ ):
+ super().__init__(sigma_min=sigma_min, sigma_max=sigma_max)
+ self.wrapping_boundary = wrapping_boundary
+
+
+class WrappedVPSDE(WrappedSDEMixin, VPSDE):
+ def __init__(
+ self,
+ wrapping_boundary: float = 1.0,
+ beta_min: float = 0.1,
+ beta_max: float = 20,
+ ):
+ super().__init__(beta_min=beta_min, beta_max=beta_max)
+ self.wrapping_boundary = wrapping_boundary
diff --git a/model/finetune.py b/model/finetune.py
new file mode 100644
index 0000000000000000000000000000000000000000..c7ce0c7c77ddb71c1116e7bedaeaf16eae4c2931
--- /dev/null
+++ b/model/finetune.py
@@ -0,0 +1,93 @@
+"""Fine-tuning helpers backed by the example's local MatterGen model package."""
+
+import logging
+from collections import OrderedDict
+from copy import deepcopy
+from pathlib import Path
+from typing import Tuple
+
+import hydra
+import pytorch_lightning as pl
+import torch
+from omegaconf import DictConfig, open_dict
+
+from .common.utils.data_classes import MatterGenCheckpointInfo, _rewrite_vendored_targets
+from .common.utils.globals import get_device
+
+logger = logging.getLogger(__name__)
+
+
+def init_adapter_lightningmodule_from_pretrained(
+ adapter_cfg: DictConfig, lightning_module_cfg: DictConfig
+) -> Tuple[pl.LightningModule, DictConfig]:
+ """Initialize the local adapter model from a MatterGen checkpoint."""
+ if adapter_cfg.model_path is not None:
+ if adapter_cfg.pretrained_name is not None:
+ logger.warning(
+ "pretrained_name is provided, but will be ignored since model_path is also provided."
+ )
+ model_path = Path(hydra.utils.to_absolute_path(adapter_cfg.model_path))
+ checkpoint_info = MatterGenCheckpointInfo(model_path, adapter_cfg.load_epoch)
+ elif adapter_cfg.pretrained_name is not None:
+ checkpoint_info = MatterGenCheckpointInfo.from_hf_hub(
+ adapter_cfg.pretrained_name
+ )
+ else:
+ raise ValueError("Either adapter.model_path or adapter.pretrained_name is required.")
+
+ checkpoint_path = checkpoint_info.checkpoint_path
+ version_root_path = Path(checkpoint_path).relative_to(
+ checkpoint_info.model_path
+ ).parents[1]
+ config_path = Path(checkpoint_info.model_path) / version_root_path
+ pretrained_cfg_path = (
+ config_path if (config_path / "config.yaml").exists() else config_path.parent.parent
+ )
+
+ hydra.core.global_hydra.GlobalHydra.instance().clear()
+ with hydra.initialize_config_dir(
+ str(pretrained_cfg_path.absolute()), version_base="1.1"
+ ):
+ pretrained_cfg = hydra.compose(config_name="config")
+ pretrained_cfg = _rewrite_vendored_targets(pretrained_cfg)
+ diffusion_module_cfg = deepcopy(pretrained_cfg.lightning_module.diffusion_module)
+ denoiser_cfg = diffusion_module_cfg.model
+
+ with open_dict(adapter_cfg.adapter):
+ for key, value in denoiser_cfg.items():
+ if key not in {"_target_", "property_embeddings_adapt"}:
+ adapter_cfg.adapter[key] = value
+ if key == "property_embeddings":
+ for field in value:
+ if field in adapter_cfg.adapter.property_embeddings_adapt:
+ adapter_cfg.adapter.property_embeddings_adapt.remove(field)
+
+ adapter_cfg.adapter.gemnet["_target_"] = (
+ "model.common.gemnet.gemnet_ctrl.GemNetTCtrl"
+ )
+ adapter_cfg.adapter.gemnet.condition_on_adapt = list(
+ adapter_cfg.adapter.property_embeddings_adapt
+ )
+
+ with open_dict(diffusion_module_cfg):
+ diffusion_module_cfg.model = adapter_cfg.adapter
+ with open_dict(lightning_module_cfg):
+ lightning_module_cfg.diffusion_module = diffusion_module_cfg
+
+ lightning_module = hydra.utils.instantiate(lightning_module_cfg)
+ checkpoint = torch.load(checkpoint_path, map_location=get_device())
+ pretrained_state: OrderedDict = checkpoint["state_dict"]
+ local_state: OrderedDict = lightning_module.state_dict()
+ local_state.update(
+ (key, pretrained_state[key])
+ for key in local_state.keys() & pretrained_state.keys()
+ )
+ lightning_module.load_state_dict(local_state, strict=True)
+
+ if not adapter_cfg.full_finetuning:
+ pretrained_keys = set(pretrained_state.keys())
+ for name, parameter in lightning_module.named_parameters():
+ if name in pretrained_keys:
+ parameter.requires_grad_(False)
+
+ return lightning_module, lightning_module_cfg
diff --git a/model/generator.py b/model/generator.py
new file mode 100644
index 0000000000000000000000000000000000000000..e4150266bdf4d94808e9b322e6a5e4e67b5cdfd1
--- /dev/null
+++ b/model/generator.py
@@ -0,0 +1,394 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+import io
+import os
+from dataclasses import dataclass
+from pathlib import Path
+from zipfile import ZipFile
+
+import ase.io
+import hydra
+import torch
+from hydra.utils import instantiate
+from omegaconf import DictConfig, OmegaConf
+from pymatgen.core.structure import Structure
+from pymatgen.io.ase import AseAtomsAdaptor
+from tqdm import tqdm
+
+from onescience.datapipes.materials.mattergen.chemgraph import ChemGraph
+from onescience.datapipes.materials.mattergen.collate import collate
+from onescience.datapipes.materials.mattergen.condition_factory import ConditionLoader
+from onescience.datapipes.materials.mattergen.num_atoms_distribution import NUM_ATOMS_DISTRIBUTIONS
+from onescience.datapipes.materials.mattergen.types import TargetProperty
+from .common.utils.data_utils import lattice_matrix_to_params_torch
+from .common.utils.eval_utils import (
+ MatterGenCheckpointInfo,
+ get_crystals_list,
+ load_model_diffusion,
+ make_structure,
+ save_structures,
+)
+from .common.utils.globals import DEFAULT_SAMPLING_CONFIG_PATH, get_device
+from .diffusion.lightning_module import DiffusionLightningModule
+from .diffusion.sampling.pc_sampler import PredictorCorrector
+from .common.utils.data_classes import ProgressCallback
+
+
+def draw_samples_from_sampler(
+ sampler: PredictorCorrector,
+ condition_loader: ConditionLoader,
+ properties_to_condition_on: TargetProperty | None = None,
+ output_path: Path | None = None,
+ cfg: DictConfig | None = None,
+ record_trajectories: bool = True,
+ progress_callback: ProgressCallback | None = None,
+) -> list[Structure]:
+
+ # Dict
+ properties_to_condition_on = properties_to_condition_on or {}
+
+ # we cannot conditional sample on something on which the model was not trained to condition on
+ assert all([key in sampler.diffusion_module.model.cond_fields_model_was_trained_on for key in properties_to_condition_on.keys()]) # type: ignore
+
+ all_samples_list = []
+ all_trajs_list = []
+ for batch_idx, (conditioning_data, mask) in enumerate(tqdm(condition_loader, desc="Generating samples")):
+ if progress_callback is not None:
+ progress_callback(progress=batch_idx / len(condition_loader))
+
+ # generate samples
+ if record_trajectories:
+ sample, mean, intermediate_samples = sampler.sample_with_record(conditioning_data, mask)
+ all_trajs_list.extend(list_of_time_steps_to_list_of_trajectories(intermediate_samples))
+ else:
+ sample, mean = sampler.sample(conditioning_data, mask)
+ all_samples_list.extend(mean.to_data_list())
+
+ if progress_callback is not None:
+ # log 100% progress
+ progress_callback(progress=1.0)
+
+ all_samples = collate(all_samples_list)
+ assert isinstance(all_samples, ChemGraph)
+ lengths, angles = lattice_matrix_to_params_torch(all_samples.cell)
+ all_samples = all_samples.replace(lengths=lengths, angles=angles)
+
+ generated_strucs = structure_from_model_output(
+ all_samples["pos"].reshape(-1, 3),
+ all_samples["atomic_numbers"].reshape(-1),
+ all_samples["lengths"].reshape(-1, 3),
+ all_samples["angles"].reshape(-1, 3),
+ all_samples["num_atoms"].reshape(-1),
+ )
+
+ if output_path is not None:
+ assert cfg is not None
+ # Save structures to disk in both a extxyz file and a compressed zip file.
+ # do this before uploading to mongo in case there is an authentication error
+ save_structures(output_path, generated_strucs)
+
+ if record_trajectories:
+ dump_trajectories(
+ output_path=output_path,
+ all_trajs_list=all_trajs_list,
+ )
+
+ return generated_strucs
+
+
+def list_of_time_steps_to_list_of_trajectories(
+ list_of_time_steps: list[ChemGraph],
+) -> list[list[ChemGraph]]:
+ # Rearrange the shapes of the recorded intermediate samples and predictions
+ # We get a list of many ChemGraphBatches, each containing
+ # many ChemGraphs. Instead, we group all the ChemGraphs of the same trajectory together,
+ # i.e., we construct lists of many lists of
+ # many ChemGraphs.
+
+ # many lists of many ChemGraphs
+ data_lists_per_timesteps = [x.to_data_list() for x in list_of_time_steps]
+
+ # many lists of many ChemGraphs.
+ data_lists_per_sample = [
+ [data_lists_per_timesteps[ix_t][ix_traj] for ix_t in range(len(data_lists_per_timesteps))]
+ for ix_traj in range(len(data_lists_per_timesteps[0]))
+ ]
+ return data_lists_per_sample
+
+
+def dump_trajectories(
+ output_path: Path,
+ all_trajs_list: list[list[ChemGraph]],
+) -> None:
+ try:
+ # We gather all trajectories in a single zip file as .extxyz files.
+ # This way we can view them easily after downloading.
+ with ZipFile(output_path / "generated_trajectories.zip", "w") as zip_obj:
+ for ix, traj in enumerate(all_trajs_list):
+ strucs = structures_from_trajectory(traj)
+ ase_atoms = [AseAtomsAdaptor.get_atoms(crystal) for crystal in strucs]
+ str_io = io.StringIO()
+ ase.io.write(str_io, ase_atoms, format="extxyz")
+ str_io.flush()
+ zip_obj.writestr(f"gen_{ix}.extxyz", str_io.getvalue())
+ except IOError as e:
+ print(f"Got error {e} writing the trajectory to disk.")
+ except ValueError as e:
+ print(f"Got error ValueError '{e}' writing the trajectory to disk.")
+
+
+def structure_from_model_output(
+ frac_coords, atom_types, lengths, angles, num_atoms
+) -> list[Structure]:
+ structures = [
+ make_structure(
+ lengths=d["lengths"],
+ angles=d["angles"],
+ atom_types=d["atom_types"],
+ frac_coords=d["frac_coords"],
+ )
+ for d in get_crystals_list(
+ frac_coords.cpu(),
+ atom_types.cpu(),
+ lengths.cpu(),
+ angles.cpu(),
+ num_atoms.cpu(),
+ )
+ ]
+ return structures
+
+
+def structures_from_trajectory(traj: list[ChemGraph]) -> list[Structure]:
+ all_strucs = []
+ for batch in traj:
+ cell = batch.cell
+ lengths, angles = lattice_matrix_to_params_torch(cell)
+ all_strucs.extend(
+ structure_from_model_output(
+ frac_coords=batch.pos,
+ atom_types=batch.atomic_numbers,
+ lengths=lengths,
+ angles=angles,
+ num_atoms=batch.num_atoms,
+ )
+ )
+
+ return all_strucs
+
+
+@dataclass
+class CrystalGenerator:
+ checkpoint_info: MatterGenCheckpointInfo
+
+ # These may be set at runtime
+ batch_size: int | None = None
+ num_batches: int | None = None
+ target_compositions_dict: list[dict[str, float]] | None = None
+ num_atoms_distribution: str = "ALEX_MP_20"
+
+ # Conditional generation
+ diffusion_guidance_factor: float = 0.0
+ properties_to_condition_on: TargetProperty | None = None
+
+ # Additional overrides, only has an effect when using a diffusion-codebase model
+ sampling_config_overrides: list[str] | None = None
+
+ # These only have an effect when using a legacy model
+ num_samples_per_batch: int = 1
+ niggli_reduction: bool = False
+
+ # Config path, if None will default to DEFAULT_SAMPLING_CONFIG_PATH
+ sampling_config_path: Path | None = None
+ sampling_config_name: str = "default"
+
+ record_trajectories: bool = True # store all intermediate samples by default
+
+ # These attributes are set when prepare() method is called.
+ _model: DiffusionLightningModule | None = None
+ _cfg: DictConfig | None = None
+
+ # can be used to monitor progress of generation
+ progress_callback: ProgressCallback | None = None
+
+ def __post_init__(self) -> None:
+ assert self.num_atoms_distribution in NUM_ATOMS_DISTRIBUTIONS, (
+ f"num_atoms_distribution must be one of {list(NUM_ATOMS_DISTRIBUTIONS.keys())}, "
+ f"but got {self.num_atoms_distribution}. To add your own distribution, "
+ "please add it to onescience.datapipes.materials.mattergen.num_atoms_distribution.NUM_ATOMS_DISTRIBUTIONS."
+ )
+ if self.target_compositions_dict:
+ assert self.cfg.lightning_module.diffusion_module.loss_fn.weights.get(
+ "atomic_numbers", 0.0
+ ) == 0.0 and "atomic_numbers" not in self.cfg.lightning_module.diffusion_module.corruption.get(
+ "discrete_corruptions", {}
+ ), "Input model appears to have been trained for crystal generation (i.e., with atom type denoising), not crystal structure prediction. Please use a model trained for crystal structure prediction instead."
+ sampling_cfg = self._load_sampling_config(
+ sampling_config_name=self.sampling_config_name,
+ sampling_config_overrides=self.sampling_config_overrides,
+ sampling_config_path=self.sampling_config_path,
+ )
+ if (
+ "atomic_numbers" in sampling_cfg.sampler_partial.predictor_partials
+ or "atomic_numbers" in sampling_cfg.sampler_partial.corrector_partials
+ ):
+ raise ValueError(
+ "Incompatible sampling config for crystal structure prediction: found atomic_numbers in predictor_partials or corrector_partials. Use the 'csp' sampling config instead, e.g., via --sampling-config-name=csp."
+ )
+
+ @property
+ def model(self) -> DiffusionLightningModule:
+ self.prepare()
+ assert self._model is not None
+ return self._model
+
+ @property
+ def cfg(self) -> DictConfig:
+ self._cfg = self.checkpoint_info.config
+ assert self._cfg is not None
+ return self._cfg
+
+ @property
+ def num_structures_to_generate(self) -> int:
+ """Returns the total number of structures to generate if `batch_size` and `num_batches` are specified at construction time;
+ otherwise, raises an AssertionError.
+ """
+ assert self.batch_size is not None
+ assert self.num_batches is not None
+ return self.batch_size * self.num_batches
+
+ @property
+ def sampling_config(self) -> DictConfig:
+ """Returns the sampling config if `batch_size` and `num_batches` are specified at construction time;
+ otherwise, raises an AssertionError.
+ """
+ assert self.batch_size is not None
+ assert self.num_batches is not None
+ return self.load_sampling_config(
+ batch_size=self.batch_size,
+ num_batches=self.num_batches,
+ target_compositions_dict=self.target_compositions_dict,
+ )
+
+ def get_condition_loader(
+ self,
+ sampling_config: DictConfig,
+ target_compositions_dict: list[dict[str, float]] | None = None,
+ ) -> ConditionLoader:
+ condition_loader_partial = instantiate(sampling_config.condition_loader_partial)
+ if not target_compositions_dict:
+ return condition_loader_partial(properties=self.properties_to_condition_on)
+
+ return condition_loader_partial(target_compositions_dict=target_compositions_dict)
+
+ def load_sampling_config(
+ self,
+ batch_size: int,
+ num_batches: int,
+ target_compositions_dict: list[dict[str, float]] | None = None,
+ ) -> DictConfig:
+ """
+ Create a sampling config from the given parameters.
+ We specify certain sampling hyperparameters via the sampling config that is loaded via hydra.
+ """
+ if self.sampling_config_overrides is None:
+ sampling_config_overrides = []
+ else:
+ # avoid modifying the original list
+ sampling_config_overrides = self.sampling_config_overrides.copy()
+ if not target_compositions_dict:
+ # Default `condition_loader_partial` is
+ # onescience.datapipes.materials.mattergen.condition_factory.get_number_of_atoms_condition_loader
+ sampling_config_overrides += [
+ f"+condition_loader_partial.num_atoms_distribution={self.num_atoms_distribution}",
+ f"+condition_loader_partial.batch_size={batch_size}",
+ f"+condition_loader_partial.num_samples={num_batches * batch_size}",
+ f"sampler_partial.guidance_scale={self.diffusion_guidance_factor}",
+ ]
+ else:
+ # `condition_loader_partial` for fixed atom type (crystal structure prediction)
+ num_structures_to_generate_per_composition = (
+ num_batches * batch_size // len(target_compositions_dict)
+ )
+ sampling_config_overrides += [
+ "condition_loader_partial._target_=onescience.datapipes.materials.mattergen.condition_factory.get_composition_data_loader",
+ f"+condition_loader_partial.num_structures_to_generate_per_composition={num_structures_to_generate_per_composition}",
+ f"+condition_loader_partial.batch_size={batch_size}",
+ ]
+ return self._load_sampling_config(
+ sampling_config_overrides=sampling_config_overrides,
+ sampling_config_path=self.sampling_config_path,
+ sampling_config_name=self.sampling_config_name,
+ )
+
+ def _load_sampling_config(
+ self,
+ sampling_config_path: Path | None = None,
+ sampling_config_name: str = "default",
+ sampling_config_overrides: list[str] | None = None,
+ ) -> DictConfig:
+ if sampling_config_path is None:
+ sampling_config_path = DEFAULT_SAMPLING_CONFIG_PATH
+
+ if sampling_config_overrides is None:
+ sampling_config_overrides = []
+
+ with hydra.initialize_config_dir(os.path.abspath(str(sampling_config_path))):
+ sampling_config = hydra.compose(
+ config_name=sampling_config_name, overrides=sampling_config_overrides
+ )
+ return sampling_config
+
+ def prepare(self) -> None:
+ """Loads the model from checkpoint and prepares for generation."""
+ if self._model is not None:
+ return
+ model = load_model_diffusion(self.checkpoint_info)
+ model = model.to(get_device())
+ self._model = model
+ self._cfg = self.checkpoint_info.config
+
+ def generate(
+ self,
+ batch_size: int | None = None,
+ num_batches: int | None = None,
+ target_compositions_dict: list[dict[str, float]] | None = None,
+ output_dir: str = "outputs",
+ ) -> list[Structure]:
+ # Prioritize the runtime provided batch_size, num_batches and target_compositions_dict
+ batch_size = batch_size or self.batch_size
+ num_batches = num_batches or self.num_batches
+ target_compositions_dict = target_compositions_dict or self.target_compositions_dict
+ assert batch_size is not None
+ assert num_batches is not None
+
+ # print config for debugging and reproducibility
+ print("\nModel config:")
+ print(OmegaConf.to_yaml(self.cfg, resolve=True))
+
+ sampling_config = self.load_sampling_config(
+ batch_size=batch_size,
+ num_batches=num_batches,
+ target_compositions_dict=target_compositions_dict,
+ )
+
+ print("\nSampling config:")
+ print(OmegaConf.to_yaml(sampling_config, resolve=True))
+ condition_loader = self.get_condition_loader(sampling_config, target_compositions_dict)
+
+ sampler_partial = instantiate(sampling_config.sampler_partial)
+ sampler = sampler_partial(pl_module=self.model)
+
+ output_path = Path(output_dir).expanduser().resolve()
+ output_path.mkdir(parents=True, exist_ok=True)
+ generated_structures = draw_samples_from_sampler(
+ sampler=sampler,
+ condition_loader=condition_loader,
+ cfg=self.cfg,
+ output_path=output_path,
+ properties_to_condition_on=self.properties_to_condition_on,
+ record_trajectories=self.record_trajectories,
+ progress_callback=self.progress_callback,
+ )
+
+ return generated_structures
diff --git a/model/package_config.py b/model/package_config.py
new file mode 100644
index 0000000000000000000000000000000000000000..80c5df9da281c046dac7d15950e363219117d34a
--- /dev/null
+++ b/model/package_config.py
@@ -0,0 +1,9 @@
+from pathlib import Path
+
+
+def get_package_data():
+ package = "model"
+ data = {
+ package: ["LICENSE", "NOTICE", "SOURCE.md", "**/*.yaml", "**/*.json"],
+ }
+ return data
diff --git a/model/sampling_conf/csp.yaml b/model/sampling_conf/csp.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..b88ff9e959f4f864d7f1adf35cfaa35c36850cb1
--- /dev/null
+++ b/model/sampling_conf/csp.yaml
@@ -0,0 +1,36 @@
+sampler_partial:
+ _target_: model.diffusion.sampling.classifier_free_guidance.GuidedPredictorCorrector.from_pl_module
+ N: 1000
+ eps_t: ${eval:'1/${.N}'}
+
+ _partial_: true
+ guidance_scale: 0.0
+ remove_conditioning_fn:
+ _target_: onescience.modules.embedding.mattergen_property_embeddings.SetUnconditionalEmbeddingType
+ keep_conditioning_fn:
+ _target_: onescience.modules.embedding.mattergen_property_embeddings.SetConditionalEmbeddingType
+ predictor_partials:
+ pos:
+ _target_: model.diffusion.wrapped.wrapped_predictors_correctors.WrappedAncestralSamplingPredictor
+ _partial_: true
+ cell:
+ _target_: model.common.diffusion.predictors_correctors.LatticeAncestralSamplingPredictor
+ _partial_: true
+
+ corrector_partials:
+ pos:
+ _target_: model.diffusion.wrapped.wrapped_predictors_correctors.WrappedLangevinCorrector
+ _partial_: true
+ max_step_size: 1e6
+ snr: 0.4
+ cell:
+ _target_: model.common.diffusion.predictors_correctors.LatticeLangevinDiffCorrector
+ _partial_: true
+ max_step_size: 1e6
+ snr: 0.2
+
+ n_steps_corrector: 1
+
+condition_loader_partial:
+ _partial_: true
+ _target_: onescience.datapipes.materials.mattergen.condition_factory.get_composition_data_loader
diff --git a/model/sampling_conf/default.yaml b/model/sampling_conf/default.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..4b84005d362e985106e70b5aaca5325545e3b1f4
--- /dev/null
+++ b/model/sampling_conf/default.yaml
@@ -0,0 +1,40 @@
+sampler_partial:
+ _target_: model.diffusion.sampling.classifier_free_guidance.GuidedPredictorCorrector.from_pl_module
+ N: 1000
+ eps_t: ${eval:'1/${.N}'}
+
+ _partial_: true
+ guidance_scale: 0.0
+ remove_conditioning_fn:
+ _target_: onescience.modules.embedding.mattergen_property_embeddings.SetUnconditionalEmbeddingType
+ keep_conditioning_fn:
+ _target_: onescience.modules.embedding.mattergen_property_embeddings.SetConditionalEmbeddingType
+ predictor_partials:
+ pos:
+ _target_: model.diffusion.wrapped.wrapped_predictors_correctors.WrappedAncestralSamplingPredictor
+ _partial_: true
+ cell:
+ _target_: model.common.diffusion.predictors_correctors.LatticeAncestralSamplingPredictor
+ _partial_: true
+ atomic_numbers:
+ _target_: model.diffusion.d3pm.d3pm_predictors_correctors.D3PMAncestralSamplingPredictor
+ predict_x0: True
+ _partial_: true
+
+ corrector_partials:
+ pos:
+ _target_: model.diffusion.wrapped.wrapped_predictors_correctors.WrappedLangevinCorrector
+ _partial_: true
+ max_step_size: 1e6
+ snr: 0.4
+ cell:
+ _target_: model.common.diffusion.predictors_correctors.LatticeLangevinDiffCorrector
+ _partial_: true
+ max_step_size: 1e6
+ snr: 0.2
+
+ n_steps_corrector: 1
+
+condition_loader_partial:
+ _partial_: true
+ _target_: onescience.datapipes.materials.mattergen.condition_factory.get_number_of_atoms_condition_loader
diff --git a/submit_train.sh b/submit_train.sh
new file mode 100644
index 0000000000000000000000000000000000000000..a74c7996c14b8e8df48a7a601659f6fe639c2fa5
--- /dev/null
+++ b/submit_train.sh
@@ -0,0 +1,91 @@
+#!/bin/bash
+# MatterGen Slurm training entry. Users may edit the defaults below directly.
+# The recommended interface is: cd demo && bash run.sh --config configs/train_8dcu.yaml --submit
+# Direct interface: cd examples/matchem/mattergen && sbatch submit_train.sh
+#SBATCH --job-name=mattergen_train
+#SBATCH --partition=hx1hdexclu12
+#SBATCH --nodes=1
+#SBATCH --ntasks-per-node=1
+#SBATCH --gres=dcu:8
+#SBATCH --cpus-per-task=64
+#SBATCH --time=48:00:00
+#SBATCH --output=slurm_%j.out
+#SBATCH --error=slurm_%j.err
+
+set -euo pipefail
+
+SCRIPT_DIR="${SCRIPT_DIR:-${SLURM_SUBMIT_DIR:-$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)}}"
+if [[ ! -f "$SCRIPT_DIR/train.py" ]]; then
+ echo "ERROR: MatterGen root not found: $SCRIPT_DIR" >&2
+ echo "Run this script from examples/matchem/mattergen or use demo/run.sh." >&2
+ exit 1
+fi
+
+
+NUM_DEVICES="${NUM_DEVICES:-8}"
+TRAIN_BATCH_SIZE="${TRAIN_BATCH_SIZE:-4}"
+VAL_BATCH_SIZE="${VAL_BATCH_SIZE:-4}"
+ACCUMULATE_GRAD_BATCHES="${ACCUMULATE_GRAD_BATCHES:-16}"
+SMOKE_TEST="${SMOKE_TEST:-0}"
+
+TRAIN_ARGS=(
+ data_module=mp_20
+ "trainer.devices=$NUM_DEVICES"
+ trainer.num_nodes=1
+ "trainer.accumulate_grad_batches=$ACCUMULATE_GRAD_BATCHES"
+ "data_module.batch_size.train=$TRAIN_BATCH_SIZE"
+ "data_module.batch_size.val=$VAL_BATCH_SIZE"
+ "data_module.batch_size.test=$VAL_BATCH_SIZE"
+ data_module.num_workers.train=2
+ data_module.num_workers.val=2
+ data_module.num_workers.test=2
+ '~trainer.logger'
+)
+
+if [[ -n "${CONFIG_PATH:-}" ]]; then
+ mapfile -t TRAIN_ARGS < <(python3 - "$CONFIG_PATH" <<'PY'
+import os
+import sys
+import yaml
+
+config = yaml.safe_load(open(sys.argv[1]))
+if config.get("task", "train") != "train":
+ raise SystemExit("submit_train.sh only accepts task: train")
+for key, value in config.get("args", {}).items():
+ if isinstance(value, bool):
+ if value:
+ print(key)
+ elif value is not None:
+ print(f"{key}={os.path.expandvars(str(value))}")
+PY
+ )
+fi
+
+if [[ "$SMOKE_TEST" == "1" ]]; then
+ TRAIN_ARGS+=(
+ trainer.max_epochs=1
+ trainer.check_val_every_n_epoch=1
+ +trainer.limit_train_batches=1
+ +trainer.limit_val_batches=1
+ +trainer.num_sanity_val_steps=0
+ )
+fi
+
+for arg in "${TRAIN_ARGS[@]}"; do
+ case "$arg" in
+ trainer.devices=*) NUM_DEVICES="${arg#*=}" ;;
+ data_module.batch_size.train=*) TRAIN_BATCH_SIZE="${arg#*=}" ;;
+ data_module.batch_size.val=*) VAL_BATCH_SIZE="${arg#*=}" ;;
+ trainer.accumulate_grad_batches=*) ACCUMULATE_GRAD_BATCHES="${arg#*=}" ;;
+ esac
+done
+
+echo "MatterGen DDP training"
+echo " devices: $NUM_DEVICES"
+echo " train batch per device: $TRAIN_BATCH_SIZE"
+echo " validation batch per device: $VAL_BATCH_SIZE"
+echo " gradient accumulation: $ACCUMULATE_GRAD_BATCHES"
+echo " effective global batch: $((NUM_DEVICES * TRAIN_BATCH_SIZE * ACCUMULATE_GRAD_BATCHES))"
+
+cd "$SCRIPT_DIR"
+srun python train.py "${TRAIN_ARGS[@]}"
diff --git a/train.py b/train.py
new file mode 100644
index 0000000000000000000000000000000000000000..779286805ed4249e7b76b93bb66aabf4e1b1ade6
--- /dev/null
+++ b/train.py
@@ -0,0 +1,37 @@
+"""Train MatterGen with the OneScience materials stack."""
+
+import os
+from datetime import datetime
+from pathlib import Path
+
+import hydra
+import omegaconf
+import torch
+from omegaconf import OmegaConf
+
+from model.diffusion.config import Config
+from model.diffusion.run import main as run_training
+
+EXAMPLE_DIR = Path(__file__).resolve().parent
+os.environ.setdefault(
+ "OUTPUT_DIR",
+ str(EXAMPLE_DIR / "outputs" / "train" / datetime.now().strftime("%Y-%m-%d_%H-%M-%S")),
+)
+
+
+@hydra.main(
+ config_path=str(EXAMPLE_DIR / "conf"),
+ config_name="default",
+ version_base="1.1",
+)
+def main(cfg: omegaconf.DictConfig) -> None:
+ """Resolve the Hydra config and run MatterGen training."""
+ torch.set_float32_matmul_precision("high")
+ config = OmegaConf.merge(OmegaConf.structured(Config), cfg)
+ OmegaConf.set_readonly(config, True)
+ print(OmegaConf.to_yaml(cfg, resolve=True))
+ run_training(config)
+
+
+if __name__ == "__main__":
+ main()
diff --git a/weight/.gitattributes b/weight/.gitattributes
new file mode 100644
index 0000000000000000000000000000000000000000..94c329c8fdd2b26cec0db159e9a14fad09459088
--- /dev/null
+++ b/weight/.gitattributes
@@ -0,0 +1,8 @@
+chemical_system/checkpoints/last.ckpt filter=lfs diff=lfs merge=lfs -text
+dft_band_gap/checkpoints/last.ckpt filter=lfs diff=lfs merge=lfs -text
+dft_mag_density/checkpoints/last.ckpt filter=lfs diff=lfs merge=lfs -text
+ml_bulk_modulus/checkpoints/last.ckpt filter=lfs diff=lfs merge=lfs -text
+space_group/checkpoints/last.ckpt filter=lfs diff=lfs merge=lfs -text
+chemical_system_energy_above_hull/checkpoints/last.ckpt filter=lfs diff=lfs merge=lfs -text
+dft_mag_density_hhi_score/checkpoints/last.ckpt filter=lfs diff=lfs merge=lfs -text
+mp_20_base/checkpoints/last.ckpt filter=lfs diff=lfs merge=lfs -text
diff --git a/weight/chemical_system/checkpoints/last.ckpt b/weight/chemical_system/checkpoints/last.ckpt
new file mode 100644
index 0000000000000000000000000000000000000000..0696587bc4cf1feebde014b1f8ebc3abec83c671
--- /dev/null
+++ b/weight/chemical_system/checkpoints/last.ckpt
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4ad21d977c2b776a31c92498c4d1c88d3e1e286deeb941e997af8e082310b80e
+size 512404276
diff --git a/weight/chemical_system/config.yaml b/weight/chemical_system/config.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..671d0abb3ec2af2721ab9f5dd8c7ac3a0e28151a
--- /dev/null
+++ b/weight/chemical_system/config.yaml
@@ -0,0 +1,224 @@
+adapter:
+ adapter:
+ _target_: mattergen.adapter.GemNetTAdapter
+ atom_type_diffusion: mask
+ denoise_atom_types: true
+ gemnet:
+ _target_: mattergen.common.gemnet.gemnet_ctrl.GemNetTCtrl
+ atom_embedding:
+ _target_: mattergen.common.gemnet.layers.embedding_block.AtomEmbedding
+ emb_size: 512
+ with_mask_type: true
+ condition_on_adapt:
+ - chemical_system
+ cutoff: 7.0
+ emb_size_atom: 512
+ emb_size_edge: 512
+ latent_dim: 512
+ max_cell_images_per_dim: 5
+ max_neighbors: 50
+ num_blocks: 4
+ num_targets: 1
+ otf_graph: true
+ regress_stress: true
+ scale_file: /scratch/amlt_code/mattergen/common/gemnet/gemnet-dT.json
+ hidden_dim: 512
+ property_embeddings: {}
+ property_embeddings_adapt:
+ chemical_system:
+ _target_: mattergen.property_embeddings.PropertyEmbedding
+ conditional_embedding_module:
+ _target_: mattergen.property_embeddings.ChemicalSystemMultiHotEmbedding
+ hidden_dim: 512
+ name: chemical_system
+ scaler:
+ _target_: torch.nn.Identity
+ unconditional_embedding_module:
+ _target_: mattergen.property_embeddings.EmbeddingVector
+ hidden_dim: 512
+ full_finetuning: true
+ load_epoch: last
+ model_path: checkpoints/mattergen_base
+data_module:
+ _recursive_: true
+ _target_: mattergen.common.data.datamodule.CrystDataModule
+ average_density: 0.05771451654022283
+ batch_size:
+ train: 64
+ val: 64
+ dataset_transforms:
+ - _partial_: true
+ _target_: mattergen.common.data.dataset_transform.filter_sparse_properties
+ max_epochs: 2200
+ num_workers:
+ train: 0
+ val: 0
+ properties:
+ - chemical_system
+ root_dir: datasets/cache/alex_mp_20/
+ train_dataset:
+ _target_: mattergen.common.data.dataset.CrystalDataset.from_cache_path
+ cache_path: datasets/cache/alex_mp_20/train
+ dataset_transforms:
+ - _partial_: true
+ _target_: mattergen.common.data.dataset_transform.filter_sparse_properties
+ properties:
+ - chemical_system
+ transforms:
+ - _partial_: true
+ _target_: mattergen.common.data.transform.symmetrize_lattice
+ - _partial_: true
+ _target_: mattergen.common.data.transform.set_chemical_system_string
+ transforms:
+ - _partial_: true
+ _target_: mattergen.common.data.transform.symmetrize_lattice
+ - _partial_: true
+ _target_: mattergen.common.data.transform.set_chemical_system_string
+ val_dataset:
+ _target_: mattergen.common.data.dataset.CrystalDataset.from_cache_path
+ cache_path: datasets/cache/alex_mp_20/val
+ dataset_transforms:
+ - _partial_: true
+ _target_: mattergen.common.data.dataset_transform.filter_sparse_properties
+ properties:
+ - chemical_system
+ transforms:
+ - _partial_: true
+ _target_: mattergen.common.data.transform.symmetrize_lattice
+ - _partial_: true
+ _target_: mattergen.common.data.transform.set_chemical_system_string
+lightning_module:
+ _target_: mattergen.diffusion.lightning_module.DiffusionLightningModule
+ diffusion_module:
+ _target_: mattergen.diffusion.diffusion_module.DiffusionModule
+ corruption:
+ _target_: mattergen.diffusion.corruption.multi_corruption.MultiCorruption
+ discrete_corruptions:
+ atomic_numbers:
+ _target_: mattergen.diffusion.corruption.d3pm_corruption.D3PMCorruption
+ d3pm:
+ _target_: mattergen.diffusion.d3pm.d3pm.MaskDiffusion
+ dim: 101
+ schedule:
+ _target_: mattergen.diffusion.d3pm.d3pm.create_discrete_diffusion_schedule
+ kind: standard
+ num_steps: 1000
+ offset: 1
+ sdes:
+ cell:
+ _target_: mattergen.common.diffusion.corruption.LatticeVPSDE.from_vpsde_config
+ vpsde_config:
+ beta_max: 20
+ beta_min: 0.1
+ limit_density: 0.05771451654022283
+ limit_var_scaling_constant: 0.25
+ pos:
+ _target_: mattergen.common.diffusion.corruption.NumAtomsVarianceAdjustedWrappedVESDE
+ limit_info_key: num_atoms
+ sigma_max: 5.0
+ wrapping_boundary: 1.0
+ loss_fn:
+ _target_: mattergen.common.loss.MaterialsLoss
+ d3pm_hybrid_lambda: 0.01
+ include_atomic_numbers: true
+ include_cell: true
+ include_pos: true
+ reduce: sum
+ weights:
+ atomic_numbers: 1.0
+ cell: 1.0
+ pos: 0.1
+ model:
+ _target_: mattergen.adapter.GemNetTAdapter
+ atom_type_diffusion: mask
+ denoise_atom_types: true
+ gemnet:
+ _target_: mattergen.common.gemnet.gemnet_ctrl.GemNetTCtrl
+ atom_embedding:
+ _target_: mattergen.common.gemnet.layers.embedding_block.AtomEmbedding
+ emb_size: 512
+ with_mask_type: true
+ condition_on_adapt:
+ - chemical_system
+ cutoff: 7.0
+ emb_size_atom: 512
+ emb_size_edge: 512
+ latent_dim: 512
+ max_cell_images_per_dim: 5
+ max_neighbors: 50
+ num_blocks: 4
+ num_targets: 1
+ otf_graph: true
+ regress_stress: true
+ scale_file: /scratch/amlt_code/mattergen/common/gemnet/gemnet-dT.json
+ hidden_dim: 512
+ property_embeddings: {}
+ property_embeddings_adapt:
+ chemical_system:
+ _target_: mattergen.property_embeddings.PropertyEmbedding
+ conditional_embedding_module:
+ _target_: mattergen.property_embeddings.ChemicalSystemMultiHotEmbedding
+ hidden_dim: 512
+ name: chemical_system
+ scaler:
+ _target_: torch.nn.Identity
+ unconditional_embedding_module:
+ _target_: mattergen.property_embeddings.EmbeddingVector
+ hidden_dim: 512
+ pre_corruption_fn:
+ _target_: mattergen.property_embeddings.SetEmbeddingType
+ dropout_fields_iid: false
+ p_unconditional: 0.2
+ optimizer_partial:
+ _partial_: true
+ _target_: torch.optim.Adam
+ lr: 5.0e-06
+ scheduler_partials:
+ - frequency: 1
+ interval: epoch
+ monitor: loss_train
+ scheduler:
+ _partial_: true
+ _target_: torch.optim.lr_scheduler.ReduceLROnPlateau
+ factor: 0.6
+ min_lr: 1.0e-06
+ patience: 100
+ verbose: true
+ strict: true
+trainer:
+ _target_: pytorch_lightning.Trainer
+ accelerator: gpu
+ accumulate_grad_batches: 1
+ callbacks:
+ - _target_: pytorch_lightning.callbacks.LearningRateMonitor
+ log_momentum: false
+ logging_interval: step
+ - _target_: pytorch_lightning.callbacks.ModelCheckpoint
+ every_n_epochs: 1
+ filename: '{epoch}-{loss_val:.2f}'
+ mode: min
+ monitor: loss_val
+ save_last: true
+ save_top_k: 1
+ verbose: false
+ - _target_: pytorch_lightning.callbacks.TQDMProgressBar
+ refresh_rate: 50
+ - _target_: mattergen.common.data.callback.SetPropertyScalers
+ check_val_every_n_epoch: 1
+ devices: 8
+ gradient_clip_algorithm: value
+ gradient_clip_val: 0.5
+ logger:
+ _target_: pytorch_lightning.loggers.WandbLogger
+ job_type: train_finetune
+ project: crystal-generation
+ settings:
+ _save_requirements: false
+ _target_: wandb.Settings
+ start_method: fork
+ max_epochs: 200
+ num_nodes: 1
+ precision: 32
+ strategy:
+ _target_: pytorch_lightning.strategies.ddp.DDPStrategy
+ find_unused_parameters: true
diff --git a/weight/chemical_system_energy_above_hull/checkpoints/last.ckpt b/weight/chemical_system_energy_above_hull/checkpoints/last.ckpt
new file mode 100644
index 0000000000000000000000000000000000000000..68c29a2a9ab31e0a6f2a71cbe2587dd0fe56a1f7
--- /dev/null
+++ b/weight/chemical_system_energy_above_hull/checkpoints/last.ckpt
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:50d4a51b109b178db98de82849d0ad14397b663526b2387caee3a643f27e326a
+size 562814862
diff --git a/weight/chemical_system_energy_above_hull/config.yaml b/weight/chemical_system_energy_above_hull/config.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..e0c953654fb5a59cc678f46403317b12dbffdbda
--- /dev/null
+++ b/weight/chemical_system_energy_above_hull/config.yaml
@@ -0,0 +1,251 @@
+adapter:
+ adapter:
+ _target_: mattergen.adapter.GemNetTAdapter
+ atom_type_diffusion: mask
+ denoise_atom_types: true
+ gemnet:
+ _target_: mattergen.common.gemnet.gemnet_ctrl.GemNetTCtrl
+ atom_embedding:
+ _target_: mattergen.common.gemnet.layers.embedding_block.AtomEmbedding
+ emb_size: 512
+ with_mask_type: true
+ condition_on_adapt:
+ - chemical_system
+ - energy_above_hull
+ cutoff: 7.0
+ emb_size_atom: 512
+ emb_size_edge: 512
+ latent_dim: 512
+ max_cell_images_per_dim: 5
+ max_neighbors: 50
+ num_blocks: 4
+ num_targets: 1
+ otf_graph: true
+ regress_stress: true
+ scale_file: /scratch/amlt_code/mattergen/common/gemnet/gemnet-dT.json
+ hidden_dim: 512
+ property_embeddings: {}
+ property_embeddings_adapt:
+ chemical_system:
+ _target_: mattergen.property_embeddings.PropertyEmbedding
+ conditional_embedding_module:
+ _target_: mattergen.property_embeddings.ChemicalSystemMultiHotEmbedding
+ hidden_dim: 512
+ name: chemical_system
+ scaler:
+ _target_: torch.nn.Identity
+ unconditional_embedding_module:
+ _target_: mattergen.property_embeddings.EmbeddingVector
+ hidden_dim: 512
+ energy_above_hull:
+ _target_: mattergen.property_embeddings.PropertyEmbedding
+ conditional_embedding_module:
+ _target_: mattergen.diffusion.model_utils.NoiseLevelEncoding
+ d_model: 512
+ name: energy_above_hull
+ scaler:
+ _target_: mattergen.common.utils.data_utils.StandardScalerTorch
+ unconditional_embedding_module:
+ _target_: mattergen.property_embeddings.EmbeddingVector
+ hidden_dim: 512
+ full_finetuning: true
+ load_epoch: last
+ model_path: checkpoints/mattergen_base
+data_module:
+ _recursive_: true
+ _target_: mattergen.common.data.datamodule.CrystDataModule
+ average_density: 0.05771451654022283
+ batch_size:
+ train: 64
+ val: 64
+ dataset_transforms:
+ - _partial_: true
+ _target_: mattergen.common.data.dataset_transform.filter_sparse_properties
+ max_epochs: 2200
+ num_workers:
+ train: 0
+ val: 0
+ properties:
+ - chemical_system
+ - energy_above_hull
+ root_dir: datasets/cache/alex_mp_20
+ train_dataset:
+ _target_: mattergen.common.data.dataset.CrystalDataset.from_cache_path
+ cache_path: datasets/cache/alex_mp_20/train
+ dataset_transforms:
+ - _partial_: true
+ _target_: mattergen.common.data.dataset_transform.filter_sparse_properties
+ properties:
+ - chemical_system
+ - energy_above_hull
+ transforms:
+ - _partial_: true
+ _target_: mattergen.common.data.transform.symmetrize_lattice
+ - _partial_: true
+ _target_: mattergen.common.data.transform.set_chemical_system_string
+ transforms:
+ - _partial_: true
+ _target_: mattergen.common.data.transform.symmetrize_lattice
+ - _partial_: true
+ _target_: mattergen.common.data.transform.set_chemical_system_string
+ val_dataset:
+ _target_: mattergen.common.data.dataset.CrystalDataset.from_cache_path
+ cache_path: datasets/cache/alex_mp_20/val
+ dataset_transforms:
+ - _partial_: true
+ _target_: mattergen.common.data.dataset_transform.filter_sparse_properties
+ properties:
+ - chemical_system
+ - energy_above_hull
+ transforms:
+ - _partial_: true
+ _target_: mattergen.common.data.transform.symmetrize_lattice
+ - _partial_: true
+ _target_: mattergen.common.data.transform.set_chemical_system_string
+lightning_module:
+ _target_: mattergen.diffusion.lightning_module.DiffusionLightningModule
+ diffusion_module:
+ _target_: mattergen.diffusion.diffusion_module.DiffusionModule
+ corruption:
+ _target_: mattergen.diffusion.corruption.multi_corruption.MultiCorruption
+ discrete_corruptions:
+ atomic_numbers:
+ _target_: mattergen.diffusion.corruption.d3pm_corruption.D3PMCorruption
+ d3pm:
+ _target_: mattergen.diffusion.d3pm.d3pm.MaskDiffusion
+ dim: 101
+ schedule:
+ _target_: mattergen.diffusion.d3pm.d3pm.create_discrete_diffusion_schedule
+ kind: standard
+ num_steps: 1000
+ offset: 1
+ sdes:
+ cell:
+ _target_: mattergen.common.diffusion.corruption.LatticeVPSDE.from_vpsde_config
+ vpsde_config:
+ beta_max: 20
+ beta_min: 0.1
+ limit_density: 0.05771451654022283
+ limit_var_scaling_constant: 0.25
+ pos:
+ _target_: mattergen.common.diffusion.corruption.NumAtomsVarianceAdjustedWrappedVESDE
+ limit_info_key: num_atoms
+ sigma_max: 5.0
+ wrapping_boundary: 1.0
+ loss_fn:
+ _target_: mattergen.common.loss.MaterialsLoss
+ d3pm_hybrid_lambda: 0.01
+ include_atomic_numbers: true
+ include_cell: true
+ include_pos: true
+ reduce: sum
+ weights:
+ atomic_numbers: 1.0
+ cell: 1.0
+ pos: 0.1
+ model:
+ _target_: mattergen.adapter.GemNetTAdapter
+ atom_type_diffusion: mask
+ denoise_atom_types: true
+ gemnet:
+ _target_: mattergen.common.gemnet.gemnet_ctrl.GemNetTCtrl
+ atom_embedding:
+ _target_: mattergen.common.gemnet.layers.embedding_block.AtomEmbedding
+ emb_size: 512
+ with_mask_type: true
+ condition_on_adapt:
+ - chemical_system
+ - energy_above_hull
+ cutoff: 7.0
+ emb_size_atom: 512
+ emb_size_edge: 512
+ latent_dim: 512
+ max_cell_images_per_dim: 5
+ max_neighbors: 50
+ num_blocks: 4
+ num_targets: 1
+ otf_graph: true
+ regress_stress: true
+ scale_file: /scratch/amlt_code/mattergen/common/gemnet/gemnet-dT.json
+ hidden_dim: 512
+ property_embeddings: {}
+ property_embeddings_adapt:
+ chemical_system:
+ _target_: mattergen.property_embeddings.PropertyEmbedding
+ conditional_embedding_module:
+ _target_: mattergen.property_embeddings.ChemicalSystemMultiHotEmbedding
+ hidden_dim: 512
+ name: chemical_system
+ scaler:
+ _target_: torch.nn.Identity
+ unconditional_embedding_module:
+ _target_: mattergen.property_embeddings.EmbeddingVector
+ hidden_dim: 512
+ energy_above_hull:
+ _target_: mattergen.property_embeddings.PropertyEmbedding
+ conditional_embedding_module:
+ _target_: mattergen.diffusion.model_utils.NoiseLevelEncoding
+ d_model: 512
+ name: energy_above_hull
+ scaler:
+ _target_: mattergen.common.utils.data_utils.StandardScalerTorch
+ unconditional_embedding_module:
+ _target_: mattergen.property_embeddings.EmbeddingVector
+ hidden_dim: 512
+ pre_corruption_fn:
+ _target_: mattergen.property_embeddings.SetEmbeddingType
+ dropout_fields_iid: false
+ p_unconditional: 0.2
+ optimizer_partial:
+ _partial_: true
+ _target_: torch.optim.Adam
+ lr: 5.0e-06
+ scheduler_partials:
+ - frequency: 1
+ interval: epoch
+ monitor: loss_train
+ scheduler:
+ _partial_: true
+ _target_: torch.optim.lr_scheduler.ReduceLROnPlateau
+ factor: 0.6
+ min_lr: 1.0e-06
+ patience: 100
+ verbose: true
+ strict: true
+trainer:
+ _target_: pytorch_lightning.Trainer
+ accelerator: gpu
+ accumulate_grad_batches: 1
+ callbacks:
+ - _target_: pytorch_lightning.callbacks.LearningRateMonitor
+ log_momentum: false
+ logging_interval: step
+ - _target_: pytorch_lightning.callbacks.ModelCheckpoint
+ every_n_epochs: 1
+ filename: '{epoch}-{loss_val:.2f}'
+ mode: min
+ monitor: loss_val
+ save_last: true
+ save_top_k: 1
+ verbose: false
+ - _target_: pytorch_lightning.callbacks.TQDMProgressBar
+ refresh_rate: 50
+ - _target_: mattergen.common.data.callback.SetPropertyScalers
+ check_val_every_n_epoch: 5
+ devices: 8
+ gradient_clip_algorithm: value
+ gradient_clip_val: 0.5
+ logger:
+ _target_: pytorch_lightning.loggers.WandbLogger
+ job_type: train_finetune
+ project: crystal-generation
+ settings:
+ _save_requirements: false
+ _target_: wandb.Settings
+ start_method: fork
+ max_epochs: 200
+ num_nodes: 1
+ precision: 32
+ strategy:
+ _target_: pytorch_lightning.strategies.ddp.DDPStrategy
+ find_unused_parameters: true
diff --git a/weight/dft_band_gap/checkpoints/last.ckpt b/weight/dft_band_gap/checkpoints/last.ckpt
new file mode 100644
index 0000000000000000000000000000000000000000..7322465cb236c31e434d3624e551d37202fa462b
--- /dev/null
+++ b/weight/dft_band_gap/checkpoints/last.ckpt
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:864ebbfd360e0a8287287d290a552da6b6bd92d67418a6a40be419ed3acd5e7e
+size 511777150
diff --git a/weight/dft_band_gap/config.yaml b/weight/dft_band_gap/config.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..1eebe5d53c3801709df446678ba2b8ad893647bd
--- /dev/null
+++ b/weight/dft_band_gap/config.yaml
@@ -0,0 +1,224 @@
+adapter:
+ adapter:
+ _target_: mattergen.adapter.GemNetTAdapter
+ atom_type_diffusion: mask
+ denoise_atom_types: true
+ gemnet:
+ _target_: mattergen.common.gemnet.gemnet_ctrl.GemNetTCtrl
+ atom_embedding:
+ _target_: mattergen.common.gemnet.layers.embedding_block.AtomEmbedding
+ emb_size: 512
+ with_mask_type: true
+ condition_on_adapt:
+ - dft_band_gap
+ cutoff: 7.0
+ emb_size_atom: 512
+ emb_size_edge: 512
+ latent_dim: 512
+ max_cell_images_per_dim: 5
+ max_neighbors: 50
+ num_blocks: 4
+ num_targets: 1
+ otf_graph: true
+ regress_stress: true
+ scale_file: /scratch/amlt_code/mattergen/common/gemnet/gemnet-dT.json
+ hidden_dim: 512
+ property_embeddings: {}
+ property_embeddings_adapt:
+ dft_band_gap:
+ _target_: mattergen.property_embeddings.PropertyEmbedding
+ conditional_embedding_module:
+ _target_: mattergen.diffusion.model_utils.NoiseLevelEncoding
+ d_model: 512
+ name: dft_band_gap
+ scaler:
+ _target_: mattergen.common.utils.data_utils.StandardScalerTorch
+ unconditional_embedding_module:
+ _target_: mattergen.property_embeddings.EmbeddingVector
+ hidden_dim: 512
+ full_finetuning: true
+ load_epoch: last
+ model_path: checkpoints/mattergen_base
+data_module:
+ _recursive_: true
+ _target_: mattergen.common.data.datamodule.CrystDataModule
+ average_density: 0.05771451654022283
+ batch_size:
+ train: 64
+ val: 64
+ dataset_transforms:
+ - _partial_: true
+ _target_: mattergen.common.data.dataset_transform.filter_sparse_properties
+ max_epochs: 2200
+ num_workers:
+ train: 0
+ val: 0
+ properties:
+ - dft_band_gap
+ root_dir: datasets/cache/alex_mp_20
+ train_dataset:
+ _target_: mattergen.common.data.dataset.CrystalDataset.from_cache_path
+ cache_path: datasets/cache/alex_mp_20/train
+ dataset_transforms:
+ - _partial_: true
+ _target_: mattergen.common.data.dataset_transform.filter_sparse_properties
+ properties:
+ - dft_band_gap
+ transforms:
+ - _partial_: true
+ _target_: mattergen.common.data.transform.symmetrize_lattice
+ - _partial_: true
+ _target_: mattergen.common.data.transform.set_chemical_system_string
+ transforms:
+ - _partial_: true
+ _target_: mattergen.common.data.transform.symmetrize_lattice
+ - _partial_: true
+ _target_: mattergen.common.data.transform.set_chemical_system_string
+ val_dataset:
+ _target_: mattergen.common.data.dataset.CrystalDataset.from_cache_path
+ cache_path: datasets/cache/alex_mp_20/val
+ dataset_transforms:
+ - _partial_: true
+ _target_: mattergen.common.data.dataset_transform.filter_sparse_properties
+ properties:
+ - dft_band_gap
+ transforms:
+ - _partial_: true
+ _target_: mattergen.common.data.transform.symmetrize_lattice
+ - _partial_: true
+ _target_: mattergen.common.data.transform.set_chemical_system_string
+lightning_module:
+ _target_: mattergen.diffusion.lightning_module.DiffusionLightningModule
+ diffusion_module:
+ _target_: mattergen.diffusion.diffusion_module.DiffusionModule
+ corruption:
+ _target_: mattergen.diffusion.corruption.multi_corruption.MultiCorruption
+ discrete_corruptions:
+ atomic_numbers:
+ _target_: mattergen.diffusion.corruption.d3pm_corruption.D3PMCorruption
+ d3pm:
+ _target_: mattergen.diffusion.d3pm.d3pm.MaskDiffusion
+ dim: 101
+ schedule:
+ _target_: mattergen.diffusion.d3pm.d3pm.create_discrete_diffusion_schedule
+ kind: standard
+ num_steps: 1000
+ offset: 1
+ sdes:
+ cell:
+ _target_: mattergen.common.diffusion.corruption.LatticeVPSDE.from_vpsde_config
+ vpsde_config:
+ beta_max: 20
+ beta_min: 0.1
+ limit_density: 0.05771451654022283
+ limit_var_scaling_constant: 0.25
+ pos:
+ _target_: mattergen.common.diffusion.corruption.NumAtomsVarianceAdjustedWrappedVESDE
+ limit_info_key: num_atoms
+ sigma_max: 5.0
+ wrapping_boundary: 1.0
+ loss_fn:
+ _target_: mattergen.common.loss.MaterialsLoss
+ d3pm_hybrid_lambda: 0.01
+ include_atomic_numbers: true
+ include_cell: true
+ include_pos: true
+ reduce: sum
+ weights:
+ atomic_numbers: 1.0
+ cell: 1.0
+ pos: 0.1
+ model:
+ _target_: mattergen.adapter.GemNetTAdapter
+ atom_type_diffusion: mask
+ denoise_atom_types: true
+ gemnet:
+ _target_: mattergen.common.gemnet.gemnet_ctrl.GemNetTCtrl
+ atom_embedding:
+ _target_: mattergen.common.gemnet.layers.embedding_block.AtomEmbedding
+ emb_size: 512
+ with_mask_type: true
+ condition_on_adapt:
+ - dft_band_gap
+ cutoff: 7.0
+ emb_size_atom: 512
+ emb_size_edge: 512
+ latent_dim: 512
+ max_cell_images_per_dim: 5
+ max_neighbors: 50
+ num_blocks: 4
+ num_targets: 1
+ otf_graph: true
+ regress_stress: true
+ scale_file: /scratch/amlt_code/mattergen/common/gemnet/gemnet-dT.json
+ hidden_dim: 512
+ property_embeddings: {}
+ property_embeddings_adapt:
+ dft_band_gap:
+ _target_: mattergen.property_embeddings.PropertyEmbedding
+ conditional_embedding_module:
+ _target_: mattergen.diffusion.model_utils.NoiseLevelEncoding
+ d_model: 512
+ name: dft_band_gap
+ scaler:
+ _target_: mattergen.common.utils.data_utils.StandardScalerTorch
+ unconditional_embedding_module:
+ _target_: mattergen.property_embeddings.EmbeddingVector
+ hidden_dim: 512
+ pre_corruption_fn:
+ _target_: mattergen.property_embeddings.SetEmbeddingType
+ dropout_fields_iid: false
+ p_unconditional: 0.2
+ optimizer_partial:
+ _partial_: true
+ _target_: torch.optim.Adam
+ lr: 5.0e-06
+ scheduler_partials:
+ - frequency: 1
+ interval: epoch
+ monitor: loss_train
+ scheduler:
+ _partial_: true
+ _target_: torch.optim.lr_scheduler.ReduceLROnPlateau
+ factor: 0.6
+ min_lr: 1.0e-06
+ patience: 100
+ verbose: true
+ strict: true
+trainer:
+ _target_: pytorch_lightning.Trainer
+ accelerator: gpu
+ accumulate_grad_batches: 1
+ callbacks:
+ - _target_: pytorch_lightning.callbacks.LearningRateMonitor
+ log_momentum: false
+ logging_interval: step
+ - _target_: pytorch_lightning.callbacks.ModelCheckpoint
+ every_n_epochs: 1
+ filename: '{epoch}-{loss_val:.2f}'
+ mode: min
+ monitor: loss_val
+ save_last: true
+ save_top_k: 1
+ verbose: false
+ - _target_: pytorch_lightning.callbacks.TQDMProgressBar
+ refresh_rate: 50
+ - _target_: mattergen.common.data.callback.SetPropertyScalers
+ check_val_every_n_epoch: 5
+ devices: 8
+ gradient_clip_algorithm: value
+ gradient_clip_val: 0.5
+ logger:
+ _target_: pytorch_lightning.loggers.WandbLogger
+ job_type: train_finetune
+ project: crystal-generation
+ settings:
+ _save_requirements: false
+ _target_: wandb.Settings
+ start_method: fork
+ max_epochs: 200
+ num_nodes: 1
+ precision: 32
+ strategy:
+ _target_: pytorch_lightning.strategies.ddp.DDPStrategy
+ find_unused_parameters: true
diff --git a/weight/dft_mag_density/checkpoints/last.ckpt b/weight/dft_mag_density/checkpoints/last.ckpt
new file mode 100644
index 0000000000000000000000000000000000000000..88ee4dd136a8021c8a3b8fd6a43545c5d27b3eb0
--- /dev/null
+++ b/weight/dft_mag_density/checkpoints/last.ckpt
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:01dd3e86805165412e0810e2a77a4756f8e1020f3ff2707c74af0a3f88a1bb8e
+size 511777278
diff --git a/weight/dft_mag_density/config.yaml b/weight/dft_mag_density/config.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..bcb67a9b17ab7ad0d33020a3509f51e06121d9f8
--- /dev/null
+++ b/weight/dft_mag_density/config.yaml
@@ -0,0 +1,224 @@
+adapter:
+ adapter:
+ _target_: mattergen.adapter.GemNetTAdapter
+ atom_type_diffusion: mask
+ denoise_atom_types: true
+ gemnet:
+ _target_: mattergen.common.gemnet.gemnet_ctrl.GemNetTCtrl
+ atom_embedding:
+ _target_: mattergen.common.gemnet.layers.embedding_block.AtomEmbedding
+ emb_size: 512
+ with_mask_type: true
+ condition_on_adapt:
+ - dft_mag_density
+ cutoff: 7.0
+ emb_size_atom: 512
+ emb_size_edge: 512
+ latent_dim: 512
+ max_cell_images_per_dim: 5
+ max_neighbors: 50
+ num_blocks: 4
+ num_targets: 1
+ otf_graph: true
+ regress_stress: true
+ scale_file: /scratch/amlt_code/mattergen/common/gemnet/gemnet-dT.json
+ hidden_dim: 512
+ property_embeddings: {}
+ property_embeddings_adapt:
+ dft_mag_density:
+ _target_: mattergen.property_embeddings.PropertyEmbedding
+ conditional_embedding_module:
+ _target_: mattergen.diffusion.model_utils.NoiseLevelEncoding
+ d_model: 512
+ name: dft_mag_density
+ scaler:
+ _target_: mattergen.common.utils.data_utils.StandardScalerTorch
+ unconditional_embedding_module:
+ _target_: mattergen.property_embeddings.EmbeddingVector
+ hidden_dim: 512
+ full_finetuning: true
+ load_epoch: last
+ model_path: checkpoints/mattergen_base
+data_module:
+ _recursive_: true
+ _target_: mattergen.common.data.datamodule.CrystDataModule
+ average_density: 0.05771451654022283
+ batch_size:
+ train: 64
+ val: 64
+ dataset_transforms:
+ - _partial_: true
+ _target_: mattergen.common.data.dataset_transform.filter_sparse_properties
+ max_epochs: 2200
+ num_workers:
+ train: 0
+ val: 0
+ properties:
+ - dft_mag_density
+ root_dir: datasets/cache/alex_mp_20
+ train_dataset:
+ _target_: mattergen.common.data.dataset.CrystalDataset.from_cache_path
+ cache_path: datasets/cache/alex_mp_20/train
+ dataset_transforms:
+ - _partial_: true
+ _target_: mattergen.common.data.dataset_transform.filter_sparse_properties
+ properties:
+ - dft_mag_density
+ transforms:
+ - _partial_: true
+ _target_: mattergen.common.data.transform.symmetrize_lattice
+ - _partial_: true
+ _target_: mattergen.common.data.transform.set_chemical_system_string
+ transforms:
+ - _partial_: true
+ _target_: mattergen.common.data.transform.symmetrize_lattice
+ - _partial_: true
+ _target_: mattergen.common.data.transform.set_chemical_system_string
+ val_dataset:
+ _target_: mattergen.common.data.dataset.CrystalDataset.from_cache_path
+ cache_path: datasets/cache/alex_mp_20/val
+ dataset_transforms:
+ - _partial_: true
+ _target_: mattergen.common.data.dataset_transform.filter_sparse_properties
+ properties:
+ - dft_mag_density
+ transforms:
+ - _partial_: true
+ _target_: mattergen.common.data.transform.symmetrize_lattice
+ - _partial_: true
+ _target_: mattergen.common.data.transform.set_chemical_system_string
+lightning_module:
+ _target_: mattergen.diffusion.lightning_module.DiffusionLightningModule
+ diffusion_module:
+ _target_: mattergen.diffusion.diffusion_module.DiffusionModule
+ corruption:
+ _target_: mattergen.diffusion.corruption.multi_corruption.MultiCorruption
+ discrete_corruptions:
+ atomic_numbers:
+ _target_: mattergen.diffusion.corruption.d3pm_corruption.D3PMCorruption
+ d3pm:
+ _target_: mattergen.diffusion.d3pm.d3pm.MaskDiffusion
+ dim: 101
+ schedule:
+ _target_: mattergen.diffusion.d3pm.d3pm.create_discrete_diffusion_schedule
+ kind: standard
+ num_steps: 1000
+ offset: 1
+ sdes:
+ cell:
+ _target_: mattergen.common.diffusion.corruption.LatticeVPSDE.from_vpsde_config
+ vpsde_config:
+ beta_max: 20
+ beta_min: 0.1
+ limit_density: 0.05771451654022283
+ limit_var_scaling_constant: 0.25
+ pos:
+ _target_: mattergen.common.diffusion.corruption.NumAtomsVarianceAdjustedWrappedVESDE
+ limit_info_key: num_atoms
+ sigma_max: 5.0
+ wrapping_boundary: 1.0
+ loss_fn:
+ _target_: mattergen.common.loss.MaterialsLoss
+ d3pm_hybrid_lambda: 0.01
+ include_atomic_numbers: true
+ include_cell: true
+ include_pos: true
+ reduce: sum
+ weights:
+ atomic_numbers: 1.0
+ cell: 1.0
+ pos: 0.1
+ model:
+ _target_: mattergen.adapter.GemNetTAdapter
+ atom_type_diffusion: mask
+ denoise_atom_types: true
+ gemnet:
+ _target_: mattergen.common.gemnet.gemnet_ctrl.GemNetTCtrl
+ atom_embedding:
+ _target_: mattergen.common.gemnet.layers.embedding_block.AtomEmbedding
+ emb_size: 512
+ with_mask_type: true
+ condition_on_adapt:
+ - dft_mag_density
+ cutoff: 7.0
+ emb_size_atom: 512
+ emb_size_edge: 512
+ latent_dim: 512
+ max_cell_images_per_dim: 5
+ max_neighbors: 50
+ num_blocks: 4
+ num_targets: 1
+ otf_graph: true
+ regress_stress: true
+ scale_file: /scratch/amlt_code/mattergen/common/gemnet/gemnet-dT.json
+ hidden_dim: 512
+ property_embeddings: {}
+ property_embeddings_adapt:
+ dft_mag_density:
+ _target_: mattergen.property_embeddings.PropertyEmbedding
+ conditional_embedding_module:
+ _target_: mattergen.diffusion.model_utils.NoiseLevelEncoding
+ d_model: 512
+ name: dft_mag_density
+ scaler:
+ _target_: mattergen.common.utils.data_utils.StandardScalerTorch
+ unconditional_embedding_module:
+ _target_: mattergen.property_embeddings.EmbeddingVector
+ hidden_dim: 512
+ pre_corruption_fn:
+ _target_: mattergen.property_embeddings.SetEmbeddingType
+ dropout_fields_iid: false
+ p_unconditional: 0.2
+ optimizer_partial:
+ _partial_: true
+ _target_: torch.optim.Adam
+ lr: 5.0e-06
+ scheduler_partials:
+ - frequency: 1
+ interval: epoch
+ monitor: loss_train
+ scheduler:
+ _partial_: true
+ _target_: torch.optim.lr_scheduler.ReduceLROnPlateau
+ factor: 0.6
+ min_lr: 1.0e-06
+ patience: 100
+ verbose: true
+ strict: true
+trainer:
+ _target_: pytorch_lightning.Trainer
+ accelerator: gpu
+ accumulate_grad_batches: 1
+ callbacks:
+ - _target_: pytorch_lightning.callbacks.LearningRateMonitor
+ log_momentum: false
+ logging_interval: step
+ - _target_: pytorch_lightning.callbacks.ModelCheckpoint
+ every_n_epochs: 1
+ filename: '{epoch}-{loss_val:.2f}'
+ mode: min
+ monitor: loss_val
+ save_last: true
+ save_top_k: 1
+ verbose: false
+ - _target_: pytorch_lightning.callbacks.TQDMProgressBar
+ refresh_rate: 50
+ - _target_: mattergen.common.data.callback.SetPropertyScalers
+ check_val_every_n_epoch: 1
+ devices: 8
+ gradient_clip_algorithm: value
+ gradient_clip_val: 0.5
+ logger:
+ _target_: pytorch_lightning.loggers.WandbLogger
+ job_type: train_finetune
+ project: crystal-generation
+ settings:
+ _save_requirements: false
+ _target_: wandb.Settings
+ start_method: fork
+ max_epochs: 200
+ num_nodes: 1
+ precision: 32
+ strategy:
+ _target_: pytorch_lightning.strategies.ddp.DDPStrategy
+ find_unused_parameters: true
diff --git a/weight/dft_mag_density_hhi_score/checkpoints/last.ckpt b/weight/dft_mag_density_hhi_score/checkpoints/last.ckpt
new file mode 100644
index 0000000000000000000000000000000000000000..a3b443b6f8f67a14599523108cfc14ec9b7a8c98
--- /dev/null
+++ b/weight/dft_mag_density_hhi_score/checkpoints/last.ckpt
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:98a4f2f9263b013d7d94a04c08e6928f90933d5f02c75f9b58c12609dddef016
+size 562187480
diff --git a/weight/dft_mag_density_hhi_score/config.yaml b/weight/dft_mag_density_hhi_score/config.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..f86c90fd2949b8f7b1e9e8550b066ea11b8b772c
--- /dev/null
+++ b/weight/dft_mag_density_hhi_score/config.yaml
@@ -0,0 +1,251 @@
+adapter:
+ adapter:
+ _target_: mattergen.adapter.GemNetTAdapter
+ atom_type_diffusion: mask
+ denoise_atom_types: true
+ gemnet:
+ _target_: mattergen.common.gemnet.gemnet_ctrl.GemNetTCtrl
+ atom_embedding:
+ _target_: mattergen.common.gemnet.layers.embedding_block.AtomEmbedding
+ emb_size: 512
+ with_mask_type: true
+ condition_on_adapt:
+ - dft_mag_density
+ - hhi_score
+ cutoff: 7.0
+ emb_size_atom: 512
+ emb_size_edge: 512
+ latent_dim: 512
+ max_cell_images_per_dim: 5
+ max_neighbors: 50
+ num_blocks: 4
+ num_targets: 1
+ otf_graph: true
+ regress_stress: true
+ scale_file: /scratch/amlt_code/mattergen/common/gemnet/gemnet-dT.json
+ hidden_dim: 512
+ property_embeddings: {}
+ property_embeddings_adapt:
+ dft_mag_density:
+ _target_: mattergen.property_embeddings.PropertyEmbedding
+ conditional_embedding_module:
+ _target_: mattergen.diffusion.model_utils.NoiseLevelEncoding
+ d_model: 512
+ name: dft_mag_density
+ scaler:
+ _target_: mattergen.common.utils.data_utils.StandardScalerTorch
+ unconditional_embedding_module:
+ _target_: mattergen.property_embeddings.EmbeddingVector
+ hidden_dim: 512
+ hhi_score:
+ _target_: mattergen.property_embeddings.PropertyEmbedding
+ conditional_embedding_module:
+ _target_: mattergen.diffusion.model_utils.NoiseLevelEncoding
+ d_model: 512
+ name: hhi_score
+ scaler:
+ _target_: mattergen.common.utils.data_utils.StandardScalerTorch
+ unconditional_embedding_module:
+ _target_: mattergen.property_embeddings.EmbeddingVector
+ hidden_dim: 512
+ full_finetuning: true
+ load_epoch: last
+ model_path: checkpoints/mattergen_base
+data_module:
+ _recursive_: true
+ _target_: mattergen.common.data.datamodule.CrystDataModule
+ average_density: 0.05771451654022283
+ batch_size:
+ train: 64
+ val: 64
+ dataset_transforms:
+ - _partial_: true
+ _target_: mattergen.common.data.dataset_transform.filter_sparse_properties
+ max_epochs: 2200
+ num_workers:
+ train: 0
+ val: 0
+ properties:
+ - dft_mag_density
+ - hhi_score
+ root_dir: datasets/cache/alex_mp_20
+ train_dataset:
+ _target_: mattergen.common.data.dataset.CrystalDataset.from_cache_path
+ cache_path: datasets/cache/alex_mp_20/train
+ dataset_transforms:
+ - _partial_: true
+ _target_: mattergen.common.data.dataset_transform.filter_sparse_properties
+ properties:
+ - dft_mag_density
+ - hhi_score
+ transforms:
+ - _partial_: true
+ _target_: mattergen.common.data.transform.symmetrize_lattice
+ - _partial_: true
+ _target_: mattergen.common.data.transform.set_chemical_system_string
+ transforms:
+ - _partial_: true
+ _target_: mattergen.common.data.transform.symmetrize_lattice
+ - _partial_: true
+ _target_: mattergen.common.data.transform.set_chemical_system_string
+ val_dataset:
+ _target_: mattergen.common.data.dataset.CrystalDataset.from_cache_path
+ cache_path: datasets/cache/alex_mp_20/val
+ dataset_transforms:
+ - _partial_: true
+ _target_: mattergen.common.data.dataset_transform.filter_sparse_properties
+ properties:
+ - dft_mag_density
+ - hhi_score
+ transforms:
+ - _partial_: true
+ _target_: mattergen.common.data.transform.symmetrize_lattice
+ - _partial_: true
+ _target_: mattergen.common.data.transform.set_chemical_system_string
+lightning_module:
+ _target_: mattergen.diffusion.lightning_module.DiffusionLightningModule
+ diffusion_module:
+ _target_: mattergen.diffusion.diffusion_module.DiffusionModule
+ corruption:
+ _target_: mattergen.diffusion.corruption.multi_corruption.MultiCorruption
+ discrete_corruptions:
+ atomic_numbers:
+ _target_: mattergen.diffusion.corruption.d3pm_corruption.D3PMCorruption
+ d3pm:
+ _target_: mattergen.diffusion.d3pm.d3pm.MaskDiffusion
+ dim: 101
+ schedule:
+ _target_: mattergen.diffusion.d3pm.d3pm.create_discrete_diffusion_schedule
+ kind: standard
+ num_steps: 1000
+ offset: 1
+ sdes:
+ cell:
+ _target_: mattergen.common.diffusion.corruption.LatticeVPSDE.from_vpsde_config
+ vpsde_config:
+ beta_max: 20
+ beta_min: 0.1
+ limit_density: 0.05771451654022283
+ limit_var_scaling_constant: 0.25
+ pos:
+ _target_: mattergen.common.diffusion.corruption.NumAtomsVarianceAdjustedWrappedVESDE
+ limit_info_key: num_atoms
+ sigma_max: 5.0
+ wrapping_boundary: 1.0
+ loss_fn:
+ _target_: mattergen.common.loss.MaterialsLoss
+ d3pm_hybrid_lambda: 0.01
+ include_atomic_numbers: true
+ include_cell: true
+ include_pos: true
+ reduce: sum
+ weights:
+ atomic_numbers: 1.0
+ cell: 1.0
+ pos: 0.1
+ model:
+ _target_: mattergen.adapter.GemNetTAdapter
+ atom_type_diffusion: mask
+ denoise_atom_types: true
+ gemnet:
+ _target_: mattergen.common.gemnet.gemnet_ctrl.GemNetTCtrl
+ atom_embedding:
+ _target_: mattergen.common.gemnet.layers.embedding_block.AtomEmbedding
+ emb_size: 512
+ with_mask_type: true
+ condition_on_adapt:
+ - dft_mag_density
+ - hhi_score
+ cutoff: 7.0
+ emb_size_atom: 512
+ emb_size_edge: 512
+ latent_dim: 512
+ max_cell_images_per_dim: 5
+ max_neighbors: 50
+ num_blocks: 4
+ num_targets: 1
+ otf_graph: true
+ regress_stress: true
+ scale_file: /scratch/amlt_code/mattergen/common/gemnet/gemnet-dT.json
+ hidden_dim: 512
+ property_embeddings: {}
+ property_embeddings_adapt:
+ dft_mag_density:
+ _target_: mattergen.property_embeddings.PropertyEmbedding
+ conditional_embedding_module:
+ _target_: mattergen.diffusion.model_utils.NoiseLevelEncoding
+ d_model: 512
+ name: dft_mag_density
+ scaler:
+ _target_: mattergen.common.utils.data_utils.StandardScalerTorch
+ unconditional_embedding_module:
+ _target_: mattergen.property_embeddings.EmbeddingVector
+ hidden_dim: 512
+ hhi_score:
+ _target_: mattergen.property_embeddings.PropertyEmbedding
+ conditional_embedding_module:
+ _target_: mattergen.diffusion.model_utils.NoiseLevelEncoding
+ d_model: 512
+ name: hhi_score
+ scaler:
+ _target_: mattergen.common.utils.data_utils.StandardScalerTorch
+ unconditional_embedding_module:
+ _target_: mattergen.property_embeddings.EmbeddingVector
+ hidden_dim: 512
+ pre_corruption_fn:
+ _target_: mattergen.property_embeddings.SetEmbeddingType
+ dropout_fields_iid: false
+ p_unconditional: 0.2
+ optimizer_partial:
+ _partial_: true
+ _target_: torch.optim.Adam
+ lr: 5.0e-06
+ scheduler_partials:
+ - frequency: 1
+ interval: epoch
+ monitor: loss_train
+ scheduler:
+ _partial_: true
+ _target_: torch.optim.lr_scheduler.ReduceLROnPlateau
+ factor: 0.6
+ min_lr: 1.0e-06
+ patience: 100
+ verbose: true
+ strict: true
+trainer:
+ _target_: pytorch_lightning.Trainer
+ accelerator: gpu
+ accumulate_grad_batches: 1
+ callbacks:
+ - _target_: pytorch_lightning.callbacks.LearningRateMonitor
+ log_momentum: false
+ logging_interval: step
+ - _target_: pytorch_lightning.callbacks.ModelCheckpoint
+ every_n_epochs: 1
+ filename: '{epoch}-{loss_val:.2f}'
+ mode: min
+ monitor: loss_val
+ save_last: true
+ save_top_k: 1
+ verbose: false
+ - _target_: pytorch_lightning.callbacks.TQDMProgressBar
+ refresh_rate: 50
+ - _target_: mattergen.common.data.callback.SetPropertyScalers
+ check_val_every_n_epoch: 5
+ devices: 8
+ gradient_clip_algorithm: value
+ gradient_clip_val: 0.5
+ logger:
+ _target_: pytorch_lightning.loggers.WandbLogger
+ job_type: train_finetune
+ project: crystal-generation
+ settings:
+ _save_requirements: false
+ _target_: wandb.Settings
+ start_method: fork
+ max_epochs: 200
+ num_nodes: 1
+ precision: 32
+ strategy:
+ _target_: pytorch_lightning.strategies.ddp.DDPStrategy
+ find_unused_parameters: true
diff --git a/weight/mattergen_base/checkpoints/.gitattributes b/weight/mattergen_base/checkpoints/.gitattributes
new file mode 100644
index 0000000000000000000000000000000000000000..109114b1655849e379b65453433bea5051f7d710
--- /dev/null
+++ b/weight/mattergen_base/checkpoints/.gitattributes
@@ -0,0 +1 @@
+last.ckpt filter=lfs diff=lfs merge=lfs -text
diff --git a/weight/mattergen_base/checkpoints/last.ckpt b/weight/mattergen_base/checkpoints/last.ckpt
new file mode 100644
index 0000000000000000000000000000000000000000..386097d9280cc3659b1c31d1812b3fd3a8bf8d77
--- /dev/null
+++ b/weight/mattergen_base/checkpoints/last.ckpt
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:81668ee12afc1ee1b037f362420730de3460bfd2d36e547585fdcb911a3dfdef
+size 461370350
diff --git a/weight/mattergen_base/config.yaml b/weight/mattergen_base/config.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..04012f3a1f9ea12db3d46eef8e41ec745361f09b
--- /dev/null
+++ b/weight/mattergen_base/config.yaml
@@ -0,0 +1,184 @@
+auto_resume: true
+checkpoint_path: null
+data_module:
+ _recursive_: true
+ _target_: mattergen.common.data.datamodule.CrystDataModule
+ average_density: 0.05771451654022283
+ batch_size:
+ train: 32
+ val: 32
+ max_epochs: 2200
+ num_workers:
+ train: 0
+ val: 0
+ properties:
+ - dft_bulk_modulus
+ - dft_band_gap
+ - dft_mag_density
+ - ml_bulk_modulus
+ - hhi_score
+ - space_group
+ - energy_above_hull
+ root_dir: datasets/cache/alex_mp_20/
+ train_dataset:
+ _target_: mattergen.common.data.dataset.CrystalDataset.from_cache_path
+ cache_path: datasets/cache/alex_mp_20/train
+ properties:
+ - dft_bulk_modulus
+ - dft_band_gap
+ - dft_mag_density
+ - ml_bulk_modulus
+ - hhi_score
+ - space_group
+ - energy_above_hull
+ transforms:
+ - _partial_: true
+ _target_: mattergen.common.data.transform.symmetrize_lattice
+ - _partial_: true
+ _target_: mattergen.common.data.transform.set_chemical_system_string
+ transforms:
+ - _partial_: true
+ _target_: mattergen.common.data.transform.symmetrize_lattice
+ - _partial_: true
+ _target_: mattergen.common.data.transform.set_chemical_system_string
+ val_dataset:
+ _target_: mattergen.common.data.dataset.CrystalDataset.from_cache_path
+ cache_path: datasets/cache/alex_mp_20/val
+ properties:
+ - dft_bulk_modulus
+ - dft_band_gap
+ - dft_mag_density
+ - ml_bulk_modulus
+ - hhi_score
+ - space_group
+ - energy_above_hull
+ transforms:
+ - _partial_: true
+ _target_: mattergen.common.data.transform.symmetrize_lattice
+ - _partial_: true
+ _target_: mattergen.common.data.transform.set_chemical_system_string
+lightning_module:
+ _target_: mattergen.diffusion.lightning_module.DiffusionLightningModule
+ diffusion_module:
+ _target_: mattergen.diffusion.diffusion_module.DiffusionModule
+ corruption:
+ _target_: mattergen.diffusion.corruption.multi_corruption.MultiCorruption
+ discrete_corruptions:
+ atomic_numbers:
+ _target_: mattergen.diffusion.corruption.d3pm_corruption.D3PMCorruption
+ d3pm:
+ _target_: mattergen.diffusion.d3pm.d3pm.MaskDiffusion
+ dim: 101
+ schedule:
+ _target_: mattergen.diffusion.d3pm.d3pm.create_discrete_diffusion_schedule
+ kind: standard
+ num_steps: 1000
+ offset: 1
+ sdes:
+ cell:
+ _target_: mattergen.common.diffusion.corruption.LatticeVPSDE.from_vpsde_config
+ vpsde_config:
+ beta_max: 20
+ beta_min: 0.1
+ limit_density: 0.05771451654022283
+ limit_var_scaling_constant: 0.25
+ pos:
+ _target_: mattergen.common.diffusion.corruption.NumAtomsVarianceAdjustedWrappedVESDE
+ limit_info_key: num_atoms
+ sigma_max: 5.0
+ wrapping_boundary: 1.0
+ loss_fn:
+ _target_: mattergen.common.loss.MaterialsLoss
+ d3pm_hybrid_lambda: 0.01
+ include_atomic_numbers: true
+ include_cell: true
+ include_pos: true
+ reduce: sum
+ weights:
+ atomic_numbers: 1.0
+ cell: 1.0
+ pos: 0.1
+ model:
+ _target_: mattergen.denoiser.GemNetTDenoiser
+ atom_type_diffusion: mask
+ denoise_atom_types: true
+ gemnet:
+ _target_: mattergen.common.gemnet.gemnet.GemNetT
+ atom_embedding:
+ _target_: mattergen.common.gemnet.layers.embedding_block.AtomEmbedding
+ emb_size: 512
+ with_mask_type: true
+ cutoff: 7.0
+ emb_size_atom: 512
+ emb_size_edge: 512
+ latent_dim: 512
+ max_cell_images_per_dim: 5
+ max_neighbors: 50
+ num_blocks: 4
+ num_targets: 1
+ otf_graph: true
+ regress_stress: true
+ scale_file: /scratch/amlt_code/mattergen/common/gemnet/gemnet-dT.json
+ hidden_dim: 512
+ property_embeddings: {}
+ property_embeddings_adapt: {}
+ pre_corruption_fn:
+ _target_: mattergen.property_embeddings.SetEmbeddingType
+ dropout_fields_iid: false
+ p_unconditional: 0.2
+ optimizer_partial:
+ _partial_: true
+ _target_: torch.optim.Adam
+ lr: 0.0001
+ scheduler_partials:
+ - frequency: 1
+ interval: epoch
+ monitor: loss_train
+ scheduler:
+ _partial_: true
+ _target_: torch.optim.lr_scheduler.ReduceLROnPlateau
+ factor: 0.6
+ min_lr: 1.0e-06
+ patience: 100
+ verbose: true
+ strict: true
+load_original: false
+params: {}
+train: true
+trainer:
+ _target_: pytorch_lightning.Trainer
+ accelerator: gpu
+ accumulate_grad_batches: 1
+ callbacks:
+ - _target_: pytorch_lightning.callbacks.LearningRateMonitor
+ log_momentum: false
+ logging_interval: step
+ - _target_: pytorch_lightning.callbacks.ModelCheckpoint
+ every_n_epochs: 1
+ filename: '{epoch}-{loss_val:.2f}'
+ mode: min
+ monitor: loss_val
+ save_last: true
+ save_top_k: 1
+ verbose: false
+ - _target_: pytorch_lightning.callbacks.TQDMProgressBar
+ refresh_rate: 50
+ - _target_: mattergen.common.data.callback.SetPropertyScalers
+ check_val_every_n_epoch: 5
+ devices: 8
+ gradient_clip_algorithm: value
+ gradient_clip_val: 0.5
+ logger:
+ _target_: pytorch_lightning.loggers.WandbLogger
+ job_type: train
+ project: crystal-generation
+ settings:
+ _save_requirements: false
+ _target_: wandb.Settings
+ start_method: fork
+ max_epochs: 2200
+ num_nodes: 2
+ precision: 32
+ strategy:
+ _target_: pytorch_lightning.strategies.ddp.DDPStrategy
+ find_unused_parameters: true
diff --git a/weight/ml_bulk_modulus/checkpoints/last.ckpt b/weight/ml_bulk_modulus/checkpoints/last.ckpt
new file mode 100644
index 0000000000000000000000000000000000000000..4a4ada5ca8cf99e83cbcc8f9a8f7e34143e24143
--- /dev/null
+++ b/weight/ml_bulk_modulus/checkpoints/last.ckpt
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:67f7f5e28695776ed285a293b30160f8b546063cbac4a763a36f6ca3f6242db6
+size 511777406
diff --git a/weight/ml_bulk_modulus/config.yaml b/weight/ml_bulk_modulus/config.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..413ffcf1f474bba620fd62c9d7ee8dae8e5ed63d
--- /dev/null
+++ b/weight/ml_bulk_modulus/config.yaml
@@ -0,0 +1,227 @@
+adapter:
+ adapter:
+ _target_: mattergen.adapter.GemNetTAdapter
+ atom_type_diffusion: mask
+ denoise_atom_types: true
+ gemnet:
+ _target_: mattergen.common.gemnet.gemnet_ctrl.GemNetTCtrl
+ atom_embedding:
+ _target_: mattergen.common.gemnet.layers.embedding_block.AtomEmbedding
+ emb_size: 512
+ with_mask_type: true
+ condition_on_adapt:
+ - ml_bulk_modulus
+ cutoff: 7.0
+ emb_size_atom: 512
+ emb_size_edge: 512
+ latent_dim: 512
+ max_cell_images_per_dim: 5
+ max_neighbors: 50
+ num_blocks: 4
+ num_targets: 1
+ otf_graph: true
+ regress_stress: true
+ scale_file: /scratch/amlt_code/mattergen/common/gemnet/gemnet-dT.json
+ hidden_dim: 512
+ property_embeddings: {}
+ property_embeddings_adapt:
+ ml_bulk_modulus:
+ _target_: mattergen.property_embeddings.PropertyEmbedding
+ conditional_embedding_module:
+ _target_: mattergen.diffusion.model_utils.NoiseLevelEncoding
+ d_model: 512
+ name: ml_bulk_modulus
+ scaler:
+ _target_: mattergen.common.utils.data_utils.StandardScalerTorch
+ log10_transform: true
+ unconditional_embedding_module:
+ _target_: mattergen.property_embeddings.EmbeddingVector
+ hidden_dim: 512
+ full_finetuning: true
+ load_epoch: last
+ model_path: null
+ pretrained_name: mattergen_base
+data_module:
+ _recursive_: true
+ _target_: mattergen.common.data.datamodule.CrystDataModule
+ average_density: 0.05771451654022283
+ batch_size:
+ train: 64
+ val: 64
+ dataset_transforms:
+ - _partial_: true
+ _target_: mattergen.common.data.dataset_transform.filter_sparse_properties
+ max_epochs: 2200
+ num_workers:
+ train: 0
+ val: 0
+ properties:
+ - ml_bulk_modulus
+ root_dir: datasets/cache/alex_mp_20
+ train_dataset:
+ _target_: mattergen.common.data.dataset.CrystalDataset.from_cache_path
+ cache_path: datasets/cache/alex_mp_20/train
+ dataset_transforms:
+ - _partial_: true
+ _target_: mattergen.common.data.dataset_transform.filter_sparse_properties
+ properties:
+ - ml_bulk_modulus
+ transforms:
+ - _partial_: true
+ _target_: mattergen.common.data.transform.symmetrize_lattice
+ - _partial_: true
+ _target_: mattergen.common.data.transform.set_chemical_system_string
+ transforms:
+ - _partial_: true
+ _target_: mattergen.common.data.transform.symmetrize_lattice
+ - _partial_: true
+ _target_: mattergen.common.data.transform.set_chemical_system_string
+ val_dataset:
+ _target_: mattergen.common.data.dataset.CrystalDataset.from_cache_path
+ cache_path: datasets/cache/alex_mp_20/val
+ dataset_transforms:
+ - _partial_: true
+ _target_: mattergen.common.data.dataset_transform.filter_sparse_properties
+ properties:
+ - ml_bulk_modulus
+ transforms:
+ - _partial_: true
+ _target_: mattergen.common.data.transform.symmetrize_lattice
+ - _partial_: true
+ _target_: mattergen.common.data.transform.set_chemical_system_string
+lightning_module:
+ _target_: mattergen.diffusion.lightning_module.DiffusionLightningModule
+ diffusion_module:
+ _target_: mattergen.diffusion.diffusion_module.DiffusionModule
+ corruption:
+ _target_: mattergen.diffusion.corruption.multi_corruption.MultiCorruption
+ discrete_corruptions:
+ atomic_numbers:
+ _target_: mattergen.diffusion.corruption.d3pm_corruption.D3PMCorruption
+ d3pm:
+ _target_: mattergen.diffusion.d3pm.d3pm.MaskDiffusion
+ dim: 101
+ schedule:
+ _target_: mattergen.diffusion.d3pm.d3pm.create_discrete_diffusion_schedule
+ kind: standard
+ num_steps: 1000
+ offset: 1
+ sdes:
+ cell:
+ _target_: mattergen.common.diffusion.corruption.LatticeVPSDE.from_vpsde_config
+ vpsde_config:
+ beta_max: 20
+ beta_min: 0.1
+ limit_density: 0.05771451654022283
+ limit_var_scaling_constant: 0.25
+ pos:
+ _target_: mattergen.common.diffusion.corruption.NumAtomsVarianceAdjustedWrappedVESDE
+ limit_info_key: num_atoms
+ sigma_max: 5.0
+ wrapping_boundary: 1.0
+ loss_fn:
+ _target_: mattergen.common.loss.MaterialsLoss
+ d3pm_hybrid_lambda: 0.01
+ include_atomic_numbers: true
+ include_cell: true
+ include_pos: true
+ reduce: sum
+ weights:
+ atomic_numbers: 1.0
+ cell: 1.0
+ pos: 0.1
+ model:
+ _target_: mattergen.adapter.GemNetTAdapter
+ atom_type_diffusion: mask
+ denoise_atom_types: true
+ gemnet:
+ _target_: mattergen.common.gemnet.gemnet_ctrl.GemNetTCtrl
+ atom_embedding:
+ _target_: mattergen.common.gemnet.layers.embedding_block.AtomEmbedding
+ emb_size: 512
+ with_mask_type: true
+ condition_on_adapt:
+ - ml_bulk_modulus
+ cutoff: 7.0
+ emb_size_atom: 512
+ emb_size_edge: 512
+ latent_dim: 512
+ max_cell_images_per_dim: 5
+ max_neighbors: 50
+ num_blocks: 4
+ num_targets: 1
+ otf_graph: true
+ regress_stress: true
+ scale_file: /scratch/amlt_code/mattergen/common/gemnet/gemnet-dT.json
+ hidden_dim: 512
+ property_embeddings: {}
+ property_embeddings_adapt:
+ ml_bulk_modulus:
+ _target_: mattergen.property_embeddings.PropertyEmbedding
+ conditional_embedding_module:
+ _target_: mattergen.diffusion.model_utils.NoiseLevelEncoding
+ d_model: 512
+ name: ml_bulk_modulus
+ scaler:
+ _target_: mattergen.common.utils.data_utils.StandardScalerTorch
+ log10_transform: true
+ unconditional_embedding_module:
+ _target_: mattergen.property_embeddings.EmbeddingVector
+ hidden_dim: 512
+ pre_corruption_fn:
+ _target_: mattergen.property_embeddings.SetEmbeddingType
+ dropout_fields_iid: false
+ p_unconditional: 0.2
+ optimizer_partial:
+ _partial_: true
+ _target_: torch.optim.Adam
+ lr: 5.0e-06
+ scheduler_partials:
+ - frequency: 1
+ interval: epoch
+ monitor: loss_train
+ scheduler:
+ _partial_: true
+ _target_: torch.optim.lr_scheduler.ReduceLROnPlateau
+ factor: 0.6
+ min_lr: 1.0e-06
+ patience: 100
+ verbose: true
+ strict: true
+trainer:
+ _target_: pytorch_lightning.Trainer
+ accelerator: gpu
+ accumulate_grad_batches: 1
+ callbacks:
+ - _target_: pytorch_lightning.callbacks.LearningRateMonitor
+ log_momentum: false
+ logging_interval: step
+ - _target_: pytorch_lightning.callbacks.ModelCheckpoint
+ every_n_epochs: 1
+ filename: '{epoch}-{loss_val:.2f}'
+ mode: min
+ monitor: loss_val
+ save_last: true
+ save_top_k: 1
+ verbose: false
+ - _target_: pytorch_lightning.callbacks.TQDMProgressBar
+ refresh_rate: 50
+ - _target_: mattergen.common.data.callback.SetPropertyScalers
+ check_val_every_n_epoch: 5
+ devices: 8
+ gradient_clip_algorithm: value
+ gradient_clip_val: 0.5
+ logger:
+ _target_: pytorch_lightning.loggers.WandbLogger
+ job_type: train_finetune
+ project: crystal-generation
+ settings:
+ _save_requirements: false
+ _target_: wandb.Settings
+ start_method: fork
+ max_epochs: 200
+ num_nodes: 1
+ precision: 32
+ strategy:
+ _target_: pytorch_lightning.strategies.ddp.DDPStrategy
+ find_unused_parameters: true
diff --git a/weight/mp_20_base/checkpoints/last.ckpt b/weight/mp_20_base/checkpoints/last.ckpt
new file mode 100644
index 0000000000000000000000000000000000000000..9633905fbbdf5899fda0e25b26cb969d5abfa35f
--- /dev/null
+++ b/weight/mp_20_base/checkpoints/last.ckpt
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ffb80e4425a6f99f479a67b8cd111885d45117234e8947ff77eb3a55df420b9a
+size 461369442
diff --git a/weight/mp_20_base/config.yaml b/weight/mp_20_base/config.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..814a3656fd8224d412b811f62dd5a0161f9bdf16
--- /dev/null
+++ b/weight/mp_20_base/config.yaml
@@ -0,0 +1,185 @@
+auto_resume: true
+checkpoint_path: null
+data_module:
+ _recursive_: true
+ _target_: mattergen.common.data.datamodule.CrystDataModule
+ average_density: 0.05771451654022283
+ batch_size:
+ test: 8
+ train: 64
+ val: 8
+ dataset_transforms:
+ - _partial_: true
+ _target_: mattergen.common.data.dataset_transform.filter_sparse_properties
+ max_epochs: 900
+ num_workers:
+ test: 0
+ train: 0
+ val: 0
+ properties: []
+ root_dir: /mnt/data_cache/mattergen-release-cache/mp_20
+ test_dataset:
+ _target_: mattergen.common.data.dataset.CrystalDataset.from_cache_path
+ cache_path: /mnt/data_cache/mattergen-release-cache/mp_20/test
+ dataset_transforms:
+ - _partial_: true
+ _target_: mattergen.common.data.dataset_transform.filter_sparse_properties
+ properties: []
+ transforms:
+ - _partial_: true
+ _target_: mattergen.common.data.transform.symmetrize_lattice
+ - _partial_: true
+ _target_: mattergen.common.data.transform.set_chemical_system_string
+ train_dataset:
+ _target_: mattergen.common.data.dataset.CrystalDataset.from_cache_path
+ cache_path: /mnt/data_cache/mattergen-release-cache/mp_20/train
+ dataset_transforms:
+ - _partial_: true
+ _target_: mattergen.common.data.dataset_transform.filter_sparse_properties
+ properties: []
+ transforms:
+ - _partial_: true
+ _target_: mattergen.common.data.transform.symmetrize_lattice
+ - _partial_: true
+ _target_: mattergen.common.data.transform.set_chemical_system_string
+ transforms:
+ - _partial_: true
+ _target_: mattergen.common.data.transform.symmetrize_lattice
+ - _partial_: true
+ _target_: mattergen.common.data.transform.set_chemical_system_string
+ val_dataset:
+ _target_: mattergen.common.data.dataset.CrystalDataset.from_cache_path
+ cache_path: /mnt/data_cache/mattergen-release-cache/mp_20/val
+ dataset_transforms:
+ - _partial_: true
+ _target_: mattergen.common.data.dataset_transform.filter_sparse_properties
+ properties: []
+ transforms:
+ - _partial_: true
+ _target_: mattergen.common.data.transform.symmetrize_lattice
+ - _partial_: true
+ _target_: mattergen.common.data.transform.set_chemical_system_string
+lightning_module:
+ _target_: mattergen.diffusion.lightning_module.DiffusionLightningModule
+ diffusion_module:
+ _target_: mattergen.diffusion.diffusion_module.DiffusionModule
+ corruption:
+ _target_: mattergen.diffusion.corruption.multi_corruption.MultiCorruption
+ discrete_corruptions:
+ atomic_numbers:
+ _target_: mattergen.diffusion.corruption.d3pm_corruption.D3PMCorruption
+ d3pm:
+ _target_: mattergen.diffusion.d3pm.d3pm.MaskDiffusion
+ dim: 101
+ schedule:
+ _target_: mattergen.diffusion.d3pm.d3pm.create_discrete_diffusion_schedule
+ kind: standard
+ num_steps: 1000
+ offset: 1
+ sdes:
+ cell:
+ _target_: mattergen.common.diffusion.corruption.LatticeVPSDE.from_vpsde_config
+ vpsde_config:
+ beta_max: 20
+ beta_min: 0.1
+ limit_density: 0.05771451654022283
+ limit_var_scaling_constant: 0.25
+ pos:
+ _target_: mattergen.common.diffusion.corruption.NumAtomsVarianceAdjustedWrappedVESDE
+ limit_info_key: num_atoms
+ sigma_max: 5.0
+ wrapping_boundary: 1.0
+ loss_fn:
+ _target_: mattergen.common.loss.MaterialsLoss
+ d3pm_hybrid_lambda: 0.01
+ include_atomic_numbers: true
+ include_cell: true
+ include_pos: true
+ reduce: sum
+ weights:
+ atomic_numbers: 1.0
+ cell: 1.0
+ pos: 0.1
+ model:
+ _target_: mattergen.denoiser.GemNetTDenoiser
+ atom_type_diffusion: mask
+ denoise_atom_types: true
+ gemnet:
+ _target_: mattergen.common.gemnet.gemnet.GemNetT
+ atom_embedding:
+ _target_: mattergen.common.gemnet.layers.embedding_block.AtomEmbedding
+ emb_size: 512
+ with_mask_type: true
+ cutoff: 7.0
+ emb_size_atom: 512
+ emb_size_edge: 512
+ latent_dim: 512
+ max_cell_images_per_dim: 5
+ max_neighbors: 50
+ num_blocks: 4
+ num_targets: 1
+ otf_graph: true
+ regress_stress: true
+ scale_file: /scratch/amlt_code/mattergen/common/gemnet/gemnet-dT.json
+ hidden_dim: 512
+ property_embeddings: {}
+ property_embeddings_adapt: {}
+ pre_corruption_fn:
+ _target_: mattergen.property_embeddings.SetEmbeddingType
+ dropout_fields_iid: false
+ p_unconditional: 0.2
+ optimizer_partial:
+ _partial_: true
+ _target_: torch.optim.Adam
+ lr: 0.0001
+ scheduler_partials:
+ - frequency: 1
+ interval: epoch
+ monitor: loss_train
+ scheduler:
+ _partial_: true
+ _target_: torch.optim.lr_scheduler.ReduceLROnPlateau
+ factor: 0.6
+ min_lr: 1.0e-06
+ patience: 100
+ verbose: true
+ strict: true
+load_original: false
+params: {}
+trainer:
+ _target_: pytorch_lightning.Trainer
+ accelerator: gpu
+ accumulate_grad_batches: 1
+ callbacks:
+ - _target_: pytorch_lightning.callbacks.LearningRateMonitor
+ log_momentum: false
+ logging_interval: step
+ - _target_: pytorch_lightning.callbacks.ModelCheckpoint
+ every_n_epochs: 1
+ filename: '{epoch}-{loss_val:.2f}'
+ mode: min
+ monitor: loss_val
+ save_last: true
+ save_top_k: 1
+ verbose: false
+ - _target_: pytorch_lightning.callbacks.TQDMProgressBar
+ refresh_rate: 50
+ - _target_: mattergen.common.data.callback.SetPropertyScalers
+ check_val_every_n_epoch: 5
+ devices: 8
+ gradient_clip_algorithm: value
+ gradient_clip_val: 0.5
+ logger:
+ _target_: pytorch_lightning.loggers.WandbLogger
+ job_type: train
+ project: crystal-generation
+ settings:
+ _save_requirements: false
+ _target_: wandb.Settings
+ start_method: fork
+ max_epochs: 900
+ num_nodes: 1
+ precision: 32
+ strategy:
+ _target_: pytorch_lightning.strategies.ddp.DDPStrategy
+ find_unused_parameters: true
diff --git a/weight/space_group/checkpoints/last.ckpt b/weight/space_group/checkpoints/last.ckpt
new file mode 100644
index 0000000000000000000000000000000000000000..2b6adf095e7528dd58a59001d6a0c13c6fd7075c
--- /dev/null
+++ b/weight/space_group/checkpoints/last.ckpt
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1b6ab06739f4319d75f4a19960840ec3c8d3d32ee842c996876ef9d8b188e4bb
+size 513189244
diff --git a/weight/space_group/config.yaml b/weight/space_group/config.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..2861dcee39acb23fbbd4668cf17c1679f34209ac
--- /dev/null
+++ b/weight/space_group/config.yaml
@@ -0,0 +1,224 @@
+adapter:
+ adapter:
+ _target_: mattergen.adapter.GemNetTAdapter
+ atom_type_diffusion: mask
+ denoise_atom_types: true
+ gemnet:
+ _target_: mattergen.common.gemnet.gemnet_ctrl.GemNetTCtrl
+ atom_embedding:
+ _target_: mattergen.common.gemnet.layers.embedding_block.AtomEmbedding
+ emb_size: 512
+ with_mask_type: true
+ condition_on_adapt:
+ - space_group
+ cutoff: 7.0
+ emb_size_atom: 512
+ emb_size_edge: 512
+ latent_dim: 512
+ max_cell_images_per_dim: 5
+ max_neighbors: 50
+ num_blocks: 4
+ num_targets: 1
+ otf_graph: true
+ regress_stress: true
+ scale_file: /scratch/amlt_code/mattergen/common/gemnet/gemnet-dT.json
+ hidden_dim: 512
+ property_embeddings: {}
+ property_embeddings_adapt:
+ space_group:
+ _target_: mattergen.property_embeddings.PropertyEmbedding
+ conditional_embedding_module:
+ _target_: mattergen.property_embeddings.SpaceGroupEmbeddingVector
+ hidden_dim: 512
+ name: space_group
+ scaler:
+ _target_: torch.nn.Identity
+ unconditional_embedding_module:
+ _target_: mattergen.property_embeddings.EmbeddingVector
+ hidden_dim: 512
+ full_finetuning: true
+ load_epoch: last
+ model_path: checkpoints/mattergen_base
+data_module:
+ _recursive_: true
+ _target_: mattergen.common.data.datamodule.CrystDataModule
+ average_density: 0.05771451654022283
+ batch_size:
+ train: 64
+ val: 64
+ dataset_transforms:
+ - _partial_: true
+ _target_: mattergen.common.data.dataset_transform.filter_sparse_properties
+ max_epochs: 2200
+ num_workers:
+ train: 0
+ val: 0
+ properties:
+ - space_group
+ root_dir: datasets/cache/alex_mp_20
+ train_dataset:
+ _target_: mattergen.common.data.dataset.CrystalDataset.from_cache_path
+ cache_path: datasets/cache/alex_mp_20/train
+ dataset_transforms:
+ - _partial_: true
+ _target_: mattergen.common.data.dataset_transform.filter_sparse_properties
+ properties:
+ - space_group
+ transforms:
+ - _partial_: true
+ _target_: mattergen.common.data.transform.symmetrize_lattice
+ - _partial_: true
+ _target_: mattergen.common.data.transform.set_chemical_system_string
+ transforms:
+ - _partial_: true
+ _target_: mattergen.common.data.transform.symmetrize_lattice
+ - _partial_: true
+ _target_: mattergen.common.data.transform.set_chemical_system_string
+ val_dataset:
+ _target_: mattergen.common.data.dataset.CrystalDataset.from_cache_path
+ cache_path: datasets/cache/alex_mp_20/val
+ dataset_transforms:
+ - _partial_: true
+ _target_: mattergen.common.data.dataset_transform.filter_sparse_properties
+ properties:
+ - space_group
+ transforms:
+ - _partial_: true
+ _target_: mattergen.common.data.transform.symmetrize_lattice
+ - _partial_: true
+ _target_: mattergen.common.data.transform.set_chemical_system_string
+lightning_module:
+ _target_: mattergen.diffusion.lightning_module.DiffusionLightningModule
+ diffusion_module:
+ _target_: mattergen.diffusion.diffusion_module.DiffusionModule
+ corruption:
+ _target_: mattergen.diffusion.corruption.multi_corruption.MultiCorruption
+ discrete_corruptions:
+ atomic_numbers:
+ _target_: mattergen.diffusion.corruption.d3pm_corruption.D3PMCorruption
+ d3pm:
+ _target_: mattergen.diffusion.d3pm.d3pm.MaskDiffusion
+ dim: 101
+ schedule:
+ _target_: mattergen.diffusion.d3pm.d3pm.create_discrete_diffusion_schedule
+ kind: standard
+ num_steps: 1000
+ offset: 1
+ sdes:
+ cell:
+ _target_: mattergen.common.diffusion.corruption.LatticeVPSDE.from_vpsde_config
+ vpsde_config:
+ beta_max: 20
+ beta_min: 0.1
+ limit_density: 0.05771451654022283
+ limit_var_scaling_constant: 0.25
+ pos:
+ _target_: mattergen.common.diffusion.corruption.NumAtomsVarianceAdjustedWrappedVESDE
+ limit_info_key: num_atoms
+ sigma_max: 5.0
+ wrapping_boundary: 1.0
+ loss_fn:
+ _target_: mattergen.common.loss.MaterialsLoss
+ d3pm_hybrid_lambda: 0.01
+ include_atomic_numbers: true
+ include_cell: true
+ include_pos: true
+ reduce: sum
+ weights:
+ atomic_numbers: 1.0
+ cell: 1.0
+ pos: 0.1
+ model:
+ _target_: mattergen.adapter.GemNetTAdapter
+ atom_type_diffusion: mask
+ denoise_atom_types: true
+ gemnet:
+ _target_: mattergen.common.gemnet.gemnet_ctrl.GemNetTCtrl
+ atom_embedding:
+ _target_: mattergen.common.gemnet.layers.embedding_block.AtomEmbedding
+ emb_size: 512
+ with_mask_type: true
+ condition_on_adapt:
+ - space_group
+ cutoff: 7.0
+ emb_size_atom: 512
+ emb_size_edge: 512
+ latent_dim: 512
+ max_cell_images_per_dim: 5
+ max_neighbors: 50
+ num_blocks: 4
+ num_targets: 1
+ otf_graph: true
+ regress_stress: true
+ scale_file: /scratch/amlt_code/mattergen/common/gemnet/gemnet-dT.json
+ hidden_dim: 512
+ property_embeddings: {}
+ property_embeddings_adapt:
+ space_group:
+ _target_: mattergen.property_embeddings.PropertyEmbedding
+ conditional_embedding_module:
+ _target_: mattergen.property_embeddings.SpaceGroupEmbeddingVector
+ hidden_dim: 512
+ name: space_group
+ scaler:
+ _target_: torch.nn.Identity
+ unconditional_embedding_module:
+ _target_: mattergen.property_embeddings.EmbeddingVector
+ hidden_dim: 512
+ pre_corruption_fn:
+ _target_: mattergen.property_embeddings.SetEmbeddingType
+ dropout_fields_iid: false
+ p_unconditional: 0.2
+ optimizer_partial:
+ _partial_: true
+ _target_: torch.optim.Adam
+ lr: 5.0e-06
+ scheduler_partials:
+ - frequency: 1
+ interval: epoch
+ monitor: loss_train
+ scheduler:
+ _partial_: true
+ _target_: torch.optim.lr_scheduler.ReduceLROnPlateau
+ factor: 0.6
+ min_lr: 1.0e-06
+ patience: 100
+ verbose: true
+ strict: true
+trainer:
+ _target_: pytorch_lightning.Trainer
+ accelerator: gpu
+ accumulate_grad_batches: 1
+ callbacks:
+ - _target_: pytorch_lightning.callbacks.LearningRateMonitor
+ log_momentum: false
+ logging_interval: step
+ - _target_: pytorch_lightning.callbacks.ModelCheckpoint
+ every_n_epochs: 1
+ filename: '{epoch}-{loss_val:.2f}'
+ mode: min
+ monitor: loss_val
+ save_last: true
+ save_top_k: 1
+ verbose: false
+ - _target_: pytorch_lightning.callbacks.TQDMProgressBar
+ refresh_rate: 50
+ - _target_: mattergen.common.data.callback.SetPropertyScalers
+ check_val_every_n_epoch: 1
+ devices: 8
+ gradient_clip_algorithm: value
+ gradient_clip_val: 0.5
+ logger:
+ _target_: pytorch_lightning.loggers.WandbLogger
+ job_type: train_finetune
+ project: crystal-generation
+ settings:
+ _save_requirements: false
+ _target_: wandb.Settings
+ start_method: fork
+ max_epochs: 200
+ num_nodes: 1
+ precision: 32
+ strategy:
+ _target_: pytorch_lightning.strategies.ddp.DDPStrategy
+ find_unused_parameters: true