Upload folder using huggingface_hub

.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

.gitignore

@@ -0,0 +1,6 @@
+__pycache__/
+*.py[cod]
+
+outputs/
+demo/slurm_*.out
+demo/slurm_*.err

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: []
+---
+
+<p align="center">
+  <strong>
+    <span style="font-size: 30px;">MatterGen</span>
+  </strong>
+</p>
+
+# 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*<br>
+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/<model_name>` 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.

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

configuration.json

@@ -0,0 +1 @@
+{"framework":"Pytorch","task":"other"}

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()

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

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

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

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

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/<name>.yaml [--submit]"; exit 0 ;;
+        *) echo "Unknown argument: $1" >&2; exit 2 ;;
+    esac
+done
+[[ -n "$CONFIG" ]] || { echo "Please specify --config configs/<name>.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"

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()

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()

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}")

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 <property_embeddings_adapt> 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 <z_per_crystal> with <self.condition_embs_adapt>.
+        """
+        (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)

model/common/__init__.py

@@ -0,0 +1,3 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+__version__ = "0.1.0"

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

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

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
+        }
+    ]
+}

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
+}

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 <True>, 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())

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 <condition_on_adapt>:
+
+    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())

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

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

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"

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,
+        )

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)