A MACE fornece potenciais interatômicos rápidos e precisos de aprendizado de máquina com a passagem de mensagens equivárias de ordem superior.
Este repositório contém a implementação de referência da MACE desenvolvida por Ilyes Batatia, Gregor Simm, David Kovacs e o grupo de Gabor Csanyi e amigos (consulte colaboradores).
Também disponível:
Uma documentação parcial está disponível em: https://mace-docs.readthedocs.io
Certifique -se de instalar o Pytorch. Consulte a instalação oficial do Pytorch para obter as instruções de instalação. Selecione as opções apropriadas para o seu sistema.
Esta é a maneira recomendada de instalar o MACE.
pip install --upgrade pip
pip install mace-torch
Nota: O pacote homônimo no Pypi não tem nada a ver com este.
git clone https://github.com/ACEsuit/mace.git
pip install ./mace
Para treinar um modelo MACE, você pode usar o script mace_run_train
, que deve estar no local habitual que Pip coloca os binários (ou você pode executar explicitamente python3 <path_to_cloned_dir>/mace/cli/run_train.py
)
mace_run_train
--name= " MACE_model "
--train_file= " train.xyz "
--valid_fraction=0.05
--test_file= " test.xyz "
--config_type_weights= ' {"Default":1.0} '
--E0s= ' {1:-13.663181292231226, 6:-1029.2809654211628, 7:-1484.1187695035828, 8:-2042.0330099956639} '
--model= " MACE "
--hidden_irreps= ' 128x0e + 128x1o '
--r_max=5.0
--batch_size=10
--max_num_epochs=1500
--swa
--start_swa=1200
--ema
--ema_decay=0.99
--amsgrad
--restart_latest
--device=cuda
Para fornecer um conjunto de validação específico, use o argumento --valid_file
. Para definir um tamanho de lote maior para avaliar o conjunto de validação, especifique --valid_batch_size
.
Para controlar o tamanho do modelo, você precisa alterar --hidden_irreps
. Para a maioria dos aplicativos, o tamanho do modelo padrão recomendado é --hidden_irreps='256x0e'
(significando 256 mensagens invariantes) ou --hidden_irreps='128x0e + 128x1o'
. Se o modelo não for preciso o suficiente, você poderá incluir recursos de ordem superior, por exemplo, 128x0e + 128x1o + 128x2e
ou aumentar o número de canais para 256
. Também é possível especificar o modelo usando as teclas --num_channels=128
e --max_L=1
.
Geralmente é preferido adicionar os átomos isolados ao conjunto de treinamento, em vez de ler em suas energias através da linha de comando, como no exemplo acima. Para rotulá -los no conjunto de treinamento, set config_type=IsolatedAtom
em seus campos de informação. Se você preferir não usar ou não conhece as energias dos átomos isolados, você pode usar a opção --E0s="average"
, que estima as energias atômicas usando a regressão de mínimos quadrados.
Se a palavra -chave --swa
estiver ativada, o peso energético da perda será aumentado para os últimos ~ 20% das épocas de treinamento (de --start_swa
Epochs). Essa configuração geralmente ajuda a diminuir os erros de energia.
A precisão pode ser alterada usando a palavra-chave --default_dtype
, o padrão é float64
, mas float32
fornece uma aceleração significativa (geralmente um fator de x2 no treinamento).
As palavras -chave --batch_size
e --max_num_epochs
devem ser adaptadas com base no tamanho do conjunto de treinamento. O tamanho do lote deve ser aumentado quando o número de dados de treinamento aumentar e o número de épocas deve ser diminuído. Uma heurística para as configurações iniciais é considerar o número de atualização de gradiente constante a 200.000, que pode ser calculada como
O código pode lidar com o conjunto de treinamento com rótulos heterogêneos, por exemplo, contendo estruturas a granel com estresse e moléculas isoladas. Neste exemplo, para fazer com que o código ignore a tensão nas moléculas, anexe à sua configuração de moléculas a config_stress_weight = 0.0
.
Para usar a aceleração da GPU da Apple Silicon, instale a versão mais recente do Pytorch e especifique --device=mps
.
Para treinamento multi-GPU, use o sinalizador --distributed
. Isso usará o módulo DistributedDataParallelel de Pytorch para treinar o modelo em várias GPUs. Combine com o carregamento de dados on-line para grandes conjuntos de dados (veja abaixo). Um exemplo de script slurm pode ser encontrado em mace/scripts/distributed_example.sbatch
.
A opção para analisar todos ou alguns argumentos usando um YAML está disponível. Por exemplo, para treinar um modelo usando os argumentos acima, você pode criar um arquivo yaml your_configs.yaml
com o seguinte conteúdo:
name : nacl
seed : 2024
train_file : train.xyz
swa : yes
start_swa : 1200
max_num_epochs : 1500
device : cpu
test_file : test.xyz
E0s :
41 : -1029.2809654211628
38 : -1484.1187695035828
8 : -2042.0330099956639
config_type_weights :
Default : 1.0
E anexar à linha de comando --config="your_configs.yaml"
. Qualquer argumento especificado na linha de comando substituirá o do arquivo YAML.
Para avaliar seu modelo MACE em um arquivo xyz, execute o mace_eval_configs
:
mace_eval_configs
--configs= " your_configs.xyz "
--model= " your_model.model "
--output= " ./your_output.xyz "
Você pode executar nosso tutorial do Colab para começar rapidamente com a Mace.
Também temos um tutorial de colab mais detalhado sobre:
Se você possui um conjunto de dados grande que pode não se encaixar na memória da GPU, é recomendável pré-processar os dados em uma CPU e usar dados on-line de dados para treinar o modelo. Para pré -processamento, seu conjunto de dados especificado como um arquivo xyz executa o script preprocess_data.py
. Um exemplo é dado aqui:
mkdir processed_data
python ./mace/scripts/preprocess_data.py
--train_file= " /path/to/train_large.xyz "
--valid_fraction=0.05
--test_file= " /path/to/test_large.xyz "
--atomic_numbers= " [1, 6, 7, 8, 9, 15, 16, 17, 35, 53] "
--r_max=4.5
--h5_prefix= " processed_data/ "
--compute_statistics
--E0s= " average "
--seed=123
Para ver todas as opções e uma pequena descrição deles executando python ./mace/scripts/preprocess_data.py --help
. O script criará vários arquivos HDF5 na pasta processed_data
, que pode ser usada para treinamento. Haverá uma pasta para treinamento, uma para validação e uma separada para cada config_type
no conjunto de testes. Para treinar o modelo, use o script run_train.py
da seguinte forma:
python ./mace/scripts/run_train.py
--name= " MACE_on_big_data "
--num_workers=16
--train_file= " ./processed_data/train.h5 "
--valid_file= " ./processed_data/valid.h5 "
--test_dir= " ./processed_data "
--statistics_file= " ./processed_data/statistics.json "
--model= " ScaleShiftMACE "
--num_interactions=2
--num_channels=128
--max_L=1
--correlation=3
--batch_size=32
--valid_batch_size=32
--max_num_epochs=100
--swa
--start_swa=60
--ema
--ema_decay=0.99
--amsgrad
--error_table= ' PerAtomMAE '
--device=cuda
--seed=123
Se você deseja usar maça com pesos e vieses para registrar seus experimentos, simplesmente instale com
pip install ./mace[wandb]
E especifique os argumentos de palavra -chave necessários ( --wandb
, --wandb_project
, --wandb_entity
, --wandb_name
, --wandb_log_hypers
)
Colaboramos com o Projeto de Materiais (MP) para treinar um potencial universal de maça, cobrindo 89 elementos em cristais a granel de 1,6 m no conjunto de dados MPTRJ selecionado nas trajetórias de relaxamento MP. Os modelos são retirados no github em https://github.com/acesuit/mace-mp. Se você os usar, cite nosso artigo, que também contém uma grande variedade de aplicativos de exemplo e benchmarks.
Cuidado
Os modelos MACE-MP são treinados nas energias DFT RAW MPTRJ de saídas VASP e não são diretamente comparáveis às energias DFT do MP ou às energias do ChGNET, que foram aplicadas Correções de compatibilidade de MP2020 para alguns óxidos de metais de transição, fluorídeos (GGA/GGA+U u correções) e 14 espécies de ânions (correções de ânions). Para mais detalhes, consulte a documentação do MP e o MP2020Compatibility.yaml.
from mace . calculators import mace_mp
from ase import build
atoms = build . molecule ( 'H2O' )
calc = mace_mp ( model = "medium" , dispersion = False , default_dtype = "float32" , device = 'cuda' )
atoms . calc = calc
print ( atoms . get_potential_energy ())
Há uma série (pequenos, médios, grandes) campos de força orgânica transferível. Estes podem ser usados para a simulação de moléculas orgânicas, cristais e líquidos moleculares, ou como ponto de partida para ajustar fino em um novo conjunto de dados. Os modelos são liberados sob a licença ASL. Os modelos são retirados no github em https://github.com/acesuit/mace-ft. Se você os usar, cite nosso artigo, que também contém benchmarks detalhados e aplicativos de exemplo.
from mace . calculators import mace_off
from ase import build
atoms = build . molecule ( 'H2O' )
calc = mace_off ( model = "medium" , device = 'cuda' )
atoms . calc = calc
print ( atoms . get_potential_energy ())
Para o modelo de fundação MACE-MP-0, você pode usar o script mace_run_train
com o argumento extra --foundation_model=model_type
. Por exemplo, para finalizar o modelo pequeno em um novo conjunto de dados, você pode usar:
mace_run_train
--name= " MACE "
--foundation_model= " small "
--train_file= " train.xyz "
--valid_fraction=0.05
--test_file= " test.xyz "
--energy_weight=1.0
--forces_weight=1.0
--E0s= " average "
--lr=0.01
--scaling= " rms_forces_scaling "
--batch_size=2
--max_num_epochs=6
--ema
--ema_decay=0.99
--amsgrad
--default_dtype= " float32 "
--device=cuda
--seed=3
Outras opções são "médias" e "grandes", ou o caminho para um modelo de fundação. Se você deseja definir outro modelo, o modelo será carregado a partir do caminho fornecido --foundation_model=$path_model
, mas você precisará fornecer o conjunto completo de hiperparameters (Irreps ocultos, R_Max, etc.) correspondendo ao modelo.
Este projeto usa o pré-compromisso para executar a formatação e o linhagem de código no comprometimento. Também usamos black
, isort
, pylint
e mypy
. Recomendamos configurar seu ambiente de desenvolvimento instalando os pacotes dev
em seu ambiente Python:
pip install -e " .[dev] "
pre-commit install
A segunda linha inicializará as verificações de código de execução pre-commit
para automaticamente no comprometimento. Temos o CI configurado para verificar isso, mas é altamente recomendável que você execute esses comandos antes de se comprometer (e empurrar) para evitar cometer um código ruim acidentalmente.
Estamos felizes em aceitar solicitações de tração sob uma licença do MIT. Copie/cole o texto da licença como um comentário em sua solicitação de tração.
Se você usar este código, cite nossos papéis:
@inproceedings { Batatia2022mace ,
title = { {MACE}: Higher Order Equivariant Message Passing Neural Networks for Fast and Accurate Force Fields } ,
author = { Ilyes Batatia and David Peter Kovacs and Gregor N. C. Simm and Christoph Ortner and Gabor Csanyi } ,
booktitle = { Advances in Neural Information Processing Systems } ,
editor = { Alice H. Oh and Alekh Agarwal and Danielle Belgrave and Kyunghyun Cho } ,
year = { 2022 } ,
url = { https://openreview.net/forum?id=YPpSngE-ZU }
}
@misc { Batatia2022Design ,
title = { The Design Space of E(3)-Equivariant Atom-Centered Interatomic Potentials } ,
author = { Batatia, Ilyes and Batzner, Simon and Kov{'a}cs, D{'a}vid P{'e}ter and Musaelian, Albert and Simm, Gregor N. C. and Drautz, Ralf and Ortner, Christoph and Kozinsky, Boris and Cs{'a}nyi, G{'a}bor } ,
year = { 2022 } ,
number = { arXiv:2205.06643 } ,
eprint = { 2205.06643 } ,
eprinttype = { arxiv } ,
doi = { 10.48550/arXiv.2205.06643 } ,
archiveprefix = { arXiv }
}
Se você tiver alguma dúvida, entre em contato conosco em [email protected].
Para insetos ou solicitações de recursos, use problemas do GitHub.
Mace é publicado e distribuído sob a licença do MIT.