diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000..b7526e3
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1,163 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+# Usually these files are written by a python script from a template
+# before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+# For a library or package, you might want to ignore these files since the code is
+# intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+# However, in case of collaboration, if having platform-specific dependencies or dependencies
+# having no cross-platform support, pipenv may install dependencies that don't work, or not
+# install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+# This is especially recommended for binary packages to ensure reproducibility, and is more
+# commonly ignored for libraries.
+# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+# pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+# in version control.
+# https://pdm.fming.dev/#use-with-ide
+.pdm.toml
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+# JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+# and can be added to the global gitignore or merged into this file. For a more nuclear
+# option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/
+
+# Results directory
+results/
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
new file mode 100644
index 0000000..c7a2d66
--- /dev/null
+++ b/.pre-commit-config.yaml
@@ -0,0 +1,17 @@
+repos:
+ - repo: https://github.com/pre-commit/pre-commit-hooks
+ rev: v4.6.0
+ hooks:
+ - id: check-added-large-files
+ - id: check-yaml
+ - id: debug-statements
+ - id: end-of-file-fixer
+ - id: trailing-whitespace
+ - repo: https://github.com/astral-sh/ruff-pre-commit
+ rev: v0.4.2
+ hooks:
+ - id: ruff
+ types_or: [ python, pyi, jupyter ]
+ args: [ --fix ]
+ - id: ruff-format
+ types_or: [ python, pyi, jupyter ]
\ No newline at end of file
diff --git a/LICENSE b/LICENSE
new file mode 100644
index 0000000..5c5d1d7
--- /dev/null
+++ b/LICENSE
@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 Juno Nam and Rafael Gomez-Bombarelli
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
\ No newline at end of file
diff --git a/README.md b/README.md
index 490cb2f..15b5a0b 100644
--- a/README.md
+++ b/README.md
@@ -1,10 +1,90 @@
# Alchemical MLIP
+[![arXiv](https://img.shields.io/badge/arXiv-2404.10746-84cc16)](https://arxiv.org/abs/2404.10746)
+[![Zenodo](https://img.shields.io/badge/Zenodo-10.5281/zenodo.11081395-14b8a6.svg)](https://zenodo.org/doi/10.5281/zenodo.11081395)
+[![MIT](https://img.shields.io/badge/License-MIT-3b82f6.svg)](https://opensource.org/license/mit)
+
This repository contains the code to modify machine learning interatomic potentials (MLIPs) to enable continuous and differentiable alchemical transformations.
+Currently, we provide the alchemical modification for the [MACE](https://github.com/ACEsuit/mace) model.
The details of the method are described in the paper: [Interpolation and differentiation of alchemical degrees of freedom in machine learning interatomic potentials](https://arxiv.org/abs/2404.10746).
-> **Note:** The code is currently under final preparation and will be uploaded here shortly.
-> We anticipate having the code ready by the end of this week.
-> Please stay tuned for updates!
+## Installation
+We tested the code with Python 3.10 and the packages in `requirements.txt`.
+For example, you can create a conda environment and install the required packages as follows (assuming CUDA 11.8):
+```bash
+conda create -n alchemical-mlip python=3.10
+conda activate alchemical-mlip
+pip install torch==2.0.1 --index-url https://download.pytorch.org/whl/cu118
+pip install -r requirements.txt
+pip install -e .
+```
+
+## Static calculations
+We provide the jupyter notebooks for the lattice parameter calculations (Fig. 2 in the paper) and the compositional optimization (Fig. 3) in the `notebook` directory.
+```
+notebook/
+├── 1_solid_solution.ipynb
+└── 2_compositional_optimization.ipynb
+```
+
+## Free energy calculations
+We provide the scripts for the free energy calculations for the vacancy (Fig. 4) and perovskites (Fig. 5) in the `scripts` directory.
+```
+scripts/
+├── vacancy_frenkel_ladd.py
+├── perovskite_frenkel_ladd.py
+└── perovskite_alchemy.py
+```
+
+The arguments for the scripts are as follows:
+```bash
+# Vacancy Frenkel-Ladd calculation
+python vacancy_frenkel_ladd.py \
+ --structure-file data/structures/Fe.cif \
+ --supercell 5 5 5 \
+ --temperature 100 \
+ --output-dir data/results/vacancy/Fe_5x5x5_100K/0
+
+# Perovskite Frenkel-Ladd calculation (alpha phase)
+python perovskite_frenkel_ladd.py \
+ --structure-file data/structures/CsPbI3_alpha.cif \
+ --supercell 6 6 6 \
+ --temperature 400 \
+ --output-dir data/results/perovskite/frenkel_ladd/CsPbI3_alpha_6x6x6_400K/0
+
+# Perovskite Frenkel-Ladd calculation (delta phase)
+python perovskite_frenkel_ladd.py \
+ --structure-file data/structures/CsPbI3_delta.cif \
+ --supercell 6 3 3 \
+ --temperature 400 \
+ --output-dir data/results/perovskite/frenkel_ladd/CsPbI3_delta_6x3x3_400K/0
+
+# Perovskite alchemy calculation (alpha phase)
+python -u perovskite_alchemy.py \
+ --structure-file data/structures/CsPbI3_alpha.cif \
+ --supercell 6 6 6 \
+ --switch-pair Pb Sn \
+ --temperature 400 \
+ --output-dir data/results/perovskite/alchemy/CsPbI3_CsSnI3_alpha_400K/0
+
+# Perovskite alchemy calculation (delta phase)
+python -u perovskite_alchemy.py \
+ --structure-file data/structures/CsPbI3_delta.cif \
+ --supercell 6 3 3 \
+ --switch-pair Pb Sn \
+ --temperature 400 \
+ --output-dir data/results/perovskite/alchemy/CsPbI3_CsSnI3_delta_400K/0
+```
+
+The result files are large and not included in the repository.
+If you want to reproduce the results without running the calculations, the result files are uploaded in the [Zenodo repository](https://zenodo.org/doi/10.5281/zenodo.11081395).
+Please download the files and place them in the `data/results` directory.
+
+The post-processing scripts for the free energy calculations are provided in the `notebook` directory.
+```
+notebook/
+├── 3_vacancy_analysis.ipynb
+└── 4_perovskite_analysis.ipynb
+```
## Citation
```
diff --git a/THIRD-PARTY-LICENSES b/THIRD-PARTY-LICENSES
new file mode 100644
index 0000000..859c10b
--- /dev/null
+++ b/THIRD-PARTY-LICENSES
@@ -0,0 +1,60 @@
+Code in alchemical_mace/{calculator,model}.py is adapted from
+https://github.com/ACEsuit/mace
+
+MIT License
+
+Copyright (c) 2022 ACEsuit/mace
+
+Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
+
+-------------------------------------------------------------------------------
+
+Code in alchemical_mace/utils.py is adapted from
+https://github.com/CederGroupHub/chgnet
+
+Crystal Hamiltonian Graph neural Network (CHGNet) Copyright (c) 2023, The Regents
+of the University of California, through Lawrence Berkeley National
+Laboratory (subject to receipt of any required approvals from the U.S.
+Dept. of Energy) and the University of California, Berkeley. All rights reserved.
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+(1) Redistributions of source code must retain the above copyright notice,
+this list of conditions and the following disclaimer.
+
+(2) Redistributions in binary form must reproduce the above copyright
+notice, this list of conditions and the following disclaimer in the
+documentation and/or other materials provided with the distribution.
+
+(3) Neither the name of the University of California, Lawrence Berkeley
+National Laboratory, U.S. Dept. of Energy, University of California,
+Berkeley nor the names of its contributors may be used to endorse or
+promote products derived from this software without specific prior written
+permission.
+
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
+LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+POSSIBILITY OF SUCH DAMAGE.
+
+You are under no obligation whatsoever to provide any bug fixes, patches,
+or upgrades to the features, functionality or performance of the source
+code ("Enhancements") to anyone; however, if you choose to make your
+Enhancements available either publicly, or directly to Lawrence Berkeley
+National Laboratory, without imposing a separate written license agreement
+for such Enhancements, then you hereby grant the following license: a
+non-exclusive, royalty-free perpetual license to install, use, modify,
+prepare derivative works, incorporate into other computer software,
+distribute, and sublicense such enhancements or derivative works thereof,
+in binary and source code form.
diff --git a/alchemical_mace/__init__.py b/alchemical_mace/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/alchemical_mace/calculator.py b/alchemical_mace/calculator.py
new file mode 100644
index 0000000..ffc06c1
--- /dev/null
+++ b/alchemical_mace/calculator.py
@@ -0,0 +1,432 @@
+from typing import Sequence, Tuple
+
+import ase
+import numpy as np
+import torch
+import torch.nn.functional as F
+from ase.calculators.calculator import Calculator, all_changes
+from ase.constraints import ExpCellFilter
+from ase.optimize import FIRE
+from ase.stress import full_3x3_to_voigt_6_stress
+from mace import data
+from mace.calculators import mace_mp
+from mace.tools import torch_geometric
+
+from alchemical_mace.model import (
+ AlchemicalPair,
+ AlchemyManager,
+ alchemical_mace_mp,
+ get_z_table_and_r_max,
+)
+
+################################################################################
+# Alchemical MACE calculator
+################################################################################
+
+
+class AlchemicalMACECalculator(Calculator):
+ """
+ Alchemical MACE calculator for ASE.
+ """
+
+ def __init__(
+ self,
+ atoms: ase.Atoms,
+ alchemical_pairs: Sequence[Sequence[Tuple[int, int]]],
+ alchemical_weights: Sequence[float],
+ device: str = "cpu",
+ model: str = "medium",
+ ):
+ """
+ Initialize the Alchemical MACE calculator.
+
+ Args:
+ atoms (ase.Atoms): Atoms object.
+ alchemical_pairs (Sequence[Sequence[Tuple[int, int]]]): List of
+ alchemical pairs. Each pair is a tuple of the atom index and
+ atomic number of an alchemical atom.
+ alchemical_weights (Sequence[float]): List of alchemical weights.
+ device (str): Device to run the calculations on.
+ model (str): Model to use for the MACE calculator.
+ """
+ Calculator.__init__(self)
+ self.results = {}
+ self.implemented_properties = ["energy", "free_energy", "forces", "stress"]
+
+ # Build the alchemical MACE model
+ self.device = device
+ self.model = alchemical_mace_mp(
+ model=model, device=device, default_dtype="float32"
+ )
+ for param in self.model.parameters():
+ param.requires_grad = False
+
+ # Set AlchemyManager
+ z_table, r_max = get_z_table_and_r_max(self.model)
+ alchemical_weights = torch.tensor(alchemical_weights, dtype=torch.float32)
+ self.alchemy_manager = AlchemyManager(
+ atoms=atoms,
+ alchemical_pairs=alchemical_pairs,
+ alchemical_weights=alchemical_weights,
+ z_table=z_table,
+ r_max=r_max,
+ ).to(self.device)
+
+ # Disable alchemical weights gradients by default
+ self.alchemy_manager.alchemical_weights.requires_grad = False
+ self.calculate_alchemical_grad = False
+
+ self.num_atoms = len(atoms)
+
+ def set_alchemical_weights(self, alchemical_weights: Sequence[float]):
+ alchemical_weights = torch.tensor(
+ alchemical_weights,
+ dtype=torch.float32,
+ device=self.device,
+ )
+ self.alchemy_manager.alchemical_weights.data = alchemical_weights
+
+ def get_alchemical_atomic_masses(self) -> np.ndarray:
+ # Get atomic masses for alchemical atoms
+ node_masses = ase.data.atomic_masses[self.alchemy_manager.atomic_numbers]
+ weights = self.alchemy_manager.alchemical_weights.data
+ weights = F.pad(weights, (1, 0), "constant", 1.0).cpu().numpy()
+ node_weights = weights[self.alchemy_manager.weight_indices]
+
+ # Scatter sum to get the atomic masses
+ atom_masses = np.zeros(self.num_atoms, dtype=np.float32)
+ np.add.at(
+ atom_masses, self.alchemy_manager.atom_indices, node_masses * node_weights
+ )
+ return atom_masses
+
+ # pylint: disable=dangerous-default-value
+ def calculate(self, atoms=None, properties=None, system_changes=all_changes):
+ # call to base-class to set atoms attribute
+ Calculator.calculate(self, atoms)
+
+ # prepare data
+ tensor_kwargs = {"dtype": torch.float32, "device": self.device}
+ positions = torch.tensor(atoms.get_positions(), **tensor_kwargs)
+ cell = torch.tensor(atoms.get_cell().array, **tensor_kwargs)
+ if self.calculate_alchemical_grad:
+ self.alchemy_manager.alchemical_weights.requires_grad = True
+ batch = self.alchemy_manager(positions, cell).to(self.device)
+
+ # get outputs
+ if self.calculate_alchemical_grad:
+ out = self.model(batch, compute_stress=True, compute_alchemical_grad=True)
+ (grad,) = torch.autograd.grad(
+ outputs=[batch["node_weights"], batch["edge_weights"]],
+ inputs=[self.alchemy_manager.alchemical_weights],
+ grad_outputs=[out["node_grad"], out["edge_grad"]],
+ retain_graph=False,
+ create_graph=False,
+ )
+ grad = grad.cpu().numpy()
+ self.alchemy_manager.alchemical_weights.requires_grad = False
+ else:
+ out = self.model(batch, retain_graph=False, compute_stress=True)
+ grad = np.zeros(
+ self.alchemy_manager.alchemical_weights.shape[0], dtype=np.float32
+ )
+
+ # store results
+ self.results = {}
+ self.results["energy"] = out["energy"][0].item()
+ self.results["free_energy"] = self.results["energy"]
+ self.results["forces"] = out["forces"].detach().cpu().numpy()
+ self.results["stress"] = full_3x3_to_voigt_6_stress(
+ out["stress"][0].detach().cpu().numpy()
+ )
+ self.results["alchemical_grad"] = grad
+
+
+class NVTMACECalculator(Calculator):
+ def __init__(self, model: str = "medium", device: str = "cuda"):
+ Calculator.__init__(self)
+ self.results = {}
+ self.implemented_properties = ["energy", "free_energy", "forces", "stress"]
+ self.device = device
+ self.model = mace_mp(
+ model=model, device=device, default_dtype="float32"
+ ).models[0]
+ self.z_table, self.r_max = get_z_table_and_r_max(self.model)
+ for param in self.model.parameters():
+ param.requires_grad = False
+
+ # pylint: disable=dangerous-default-value
+ def calculate(self, atoms=None, properties=None, system_changes=all_changes):
+ # call to base-class to set atoms attribute
+ Calculator.calculate(self, atoms)
+
+ # prepare data
+ config = data.config_from_atoms(atoms)
+ atomic_data = data.AtomicData.from_config(
+ config, z_table=self.z_table, cutoff=self.r_max
+ )
+ data_loader = torch_geometric.dataloader.DataLoader(
+ dataset=[atomic_data],
+ batch_size=1,
+ shuffle=False,
+ drop_last=False,
+ )
+ batch = next(iter(data_loader)).to(self.device)
+
+ out = self.model(batch, compute_stress=False)
+ self.results = {}
+ self.results["energy"] = out["energy"][0].item()
+ self.results["free_energy"] = self.results["energy"]
+ self.results["forces"] = out["forces"].detach().cpu().numpy()
+
+
+class FrenkelLaddCalculator(Calculator):
+ """
+ Frenkel-Ladd calculator for ASE.
+ """
+
+ def __init__(
+ self,
+ spring_constants: np.ndarray,
+ initial_positions: np.ndarray,
+ device: str,
+ model: str = "medium",
+ ):
+ """
+ Initialize the Frenkel-Ladd calculator.
+
+ Args:
+ spring_constants (np.ndarray): Spring constants for each atom.
+ initial_positions (np.ndarray): Initial positions of the atoms.
+ device (str): Device to run the calculations on.
+ model (str): Model to use for the MACE calculator.
+ """
+ Calculator.__init__(self)
+ self.results = {}
+ self.implemented_properties = ["energy", "free_energy", "forces"]
+ self.device = device
+ self.model = mace_mp(
+ model=model, device=device, default_dtype="float32"
+ ).models[0]
+ self.z_table, self.r_max = get_z_table_and_r_max(self.model)
+ for param in self.model.parameters():
+ param.requires_grad = False
+
+ # Spring constants
+ self.spring_constants = spring_constants
+ self.initial_positions = initial_positions
+
+ # Reversible scaling factor
+ self.weights = [1.0, 0.0]
+ self.compute_mace = True
+
+ def set_weights(self, lambda_value: float):
+ self.weights = [1.0 - lambda_value, lambda_value]
+
+ # pylint: disable=dangerous-default-value
+ def calculate(self, atoms=None, properties=None, system_changes=all_changes):
+ # call to base-class to set atoms attribute
+ Calculator.calculate(self, atoms)
+
+ # Get MACE results if needed
+ if self.compute_mace:
+ config = data.config_from_atoms(atoms)
+ atomic_data = data.AtomicData.from_config(
+ config, z_table=self.z_table, cutoff=self.r_max
+ )
+ data_loader = torch_geometric.dataloader.DataLoader(
+ dataset=[atomic_data],
+ batch_size=1,
+ shuffle=False,
+ drop_last=False,
+ )
+ batch = next(iter(data_loader)).to(self.device)
+ out = self.model(batch, compute_stress=False) # Frenkel-Ladd is NVT
+ mace_energy = out["energy"][0].item()
+ mace_forces = out["forces"].detach().cpu().numpy()
+ else:
+ mace_energy = 0.0
+ mace_forces = np.zeros((len(atoms), 3), dtype=np.float32)
+
+ # Get spring energy and forces
+ displacement = atoms.get_positions() - self.initial_positions
+ spring_energy = 0.5 * np.sum(
+ self.spring_constants * np.sum(displacement**2, axis=1)
+ )
+ spring_forces = -self.spring_constants[:, None] * displacement
+
+ # Combine energies and forces
+ total_energy = self.weights[0] * spring_energy + self.weights[1] * mace_energy
+ total_forces = self.weights[0] * spring_forces + self.weights[1] * mace_forces
+ if self.compute_mace:
+ energy_diff = mace_energy - spring_energy
+ else:
+ energy_diff = 0.0
+
+ self.results = {}
+ self.results["energy"] = total_energy
+ self.results["free_energy"] = total_energy
+ self.results["forces"] = total_forces
+ self.results["energy_diff"] = energy_diff
+
+
+class DefectFrenkelLaddCalculator(Calculator):
+ """
+ Frenkel-Ladd calculator for ASE, for a crystal with a defect.
+ """
+
+ def __init__(
+ self,
+ atoms: ase.Atoms,
+ spring_constant: float,
+ defect_index: int,
+ device: str = "cpu",
+ model: str = "medium",
+ ):
+ """
+ Initialize the Frenkel-Ladd calculator.
+
+ Args:
+ atoms (ase.Atoms): Atoms object.
+ spring_constant (float): Spring constant for the defect atom.
+ defect_index (int): Index of the defect atom.
+ device (str): Device to run the calculations on.
+ model (str): Model to use for the MACE calculator.
+ """
+ Calculator.__init__(self)
+ self.results = {}
+ self.implemented_properties = ["energy", "free_energy", "forces", "stress"]
+
+ # Build the alchemical MACE model
+ self.device = device
+ self.model = alchemical_mace_mp(
+ model=model, device=device, default_dtype="float32"
+ )
+ for param in self.model.parameters():
+ param.requires_grad = False
+
+ # Set AlchemyManager
+ z_table, r_max = get_z_table_and_r_max(self.model)
+ alchemical_weights = torch.tensor([1.0], dtype=torch.float32)
+ atomic_number = atoms.get_atomic_numbers()[defect_index]
+ alchemical_pairs = [[AlchemicalPair(defect_index, atomic_number)]]
+ self.alchemy_manager = AlchemyManager(
+ atoms=atoms,
+ alchemical_pairs=alchemical_pairs,
+ alchemical_weights=alchemical_weights,
+ z_table=z_table,
+ r_max=r_max,
+ ).to(self.device)
+
+ # Disable alchemical weights gradients by default
+ self.alchemy_manager.alchemical_weights.requires_grad = False
+ self.calculate_alchemical_grad = False
+
+ self.num_atoms = len(atoms)
+
+ # Switching
+ self.defect_index = defect_index
+ self.spring_constant = spring_constant
+
+ def set_alchemical_weight(self, alchemical_weight: float):
+ # Set alchemical weights
+ alchemical_weights = torch.tensor(
+ [1.0 - alchemical_weight], # initial = original atoms = 1 - 0
+ dtype=torch.float32,
+ device=self.device,
+ )
+ self.alchemy_manager.alchemical_weights.data = alchemical_weights
+
+ # pylint: disable=dangerous-default-value
+ def calculate(self, atoms=None, properties=None, system_changes=all_changes):
+ # call to base-class to set atoms attribute
+ Calculator.calculate(self, atoms)
+
+ # prepare data
+ tensor_kwargs = {"dtype": torch.float32, "device": self.device}
+ positions = torch.tensor(atoms.get_positions(), **tensor_kwargs)
+ cell = torch.tensor(atoms.get_cell().array, **tensor_kwargs)
+ if self.calculate_alchemical_grad:
+ self.alchemy_manager.alchemical_weights.requires_grad = True
+ batch = self.alchemy_manager(positions, cell).to(self.device)
+
+ # get outputs
+ if self.calculate_alchemical_grad:
+ out = self.model(batch, retain_graph=True, compute_stress=True)
+ out["energy"].backward()
+ grad = self.alchemy_manager.alchemical_weights.grad.item()
+ self.alchemy_manager.alchemical_weights.grad.zero_()
+ self.alchemy_manager.alchemical_weights.requires_grad = False
+ else:
+ out = self.model(batch, retain_graph=False, compute_stress=True)
+ grad = 0.0
+ mace_energy = out["energy"][0].item()
+ mace_forces = out["forces"].detach().cpu().numpy()
+ mace_stress = out["stress"][0].detach().cpu().numpy()
+
+ # Get spring energy and forces
+ cell_center = np.array([0.5, 0.5, 0.5]) @ atoms.get_cell().array
+ displacement = atoms.get_positions()[self.defect_index] - cell_center
+ spring_energy = 0.5 * self.spring_constant * np.sum(displacement**2)
+ spring_forces = -self.spring_constant * displacement
+
+ # Combine energies and forces
+ # Note: weight here is 1 - lambda, and we're not weighting the mace
+ # energy because it's already weighted by the alchemical weight
+ weight = self.alchemy_manager.alchemical_weights.item()
+ total_energy = mace_energy + (1 - weight) * spring_energy
+ total_forces = mace_forces
+ total_forces[self.defect_index] += (1 - weight) * spring_forces
+ if self.calculate_alchemical_grad:
+ # H(lambda) = E(1 - lambda) + lambda * spring_energy
+ # dH/d(lambda) = -dE/d(1 - lambda) + spring_energy
+ grad = -grad + spring_energy
+
+ # store results
+ self.results = {}
+ self.results["energy"] = total_energy
+ self.results["free_energy"] = total_energy
+ self.results["forces"] = total_forces
+ self.results["stress"] = full_3x3_to_voigt_6_stress(mace_stress)
+ self.results["alchemical_grad"] = grad
+
+
+def get_alchemical_optimized_cellpar(
+ atoms: ase.Atoms,
+ alchemical_pairs: Sequence[Sequence[Tuple[int, int]]],
+ alchemical_weights: Sequence[float],
+ model: str = "medium",
+ device: str = "cpu",
+ **kwargs,
+):
+ """
+ Optimize the cell parameters of a crystal with alchemical atoms using the
+ Alchemical MACE calculator.
+
+ Args:
+ atoms (ase.Atoms): Atoms object.
+ alchemical_pairs (Sequence[Sequence[Tuple[int, int]]]): List of
+ alchemical pairs. Each pair is a tuple of the atom index and
+ atomic number of an alchemical atom.
+ alchemical_weights (Sequence[float]): List of alchemical weights.
+ model (str): Model to use for the MACE calculator.
+ device (str): Device to run the calculations on.
+
+ Returns:
+ np.ndarray: Optimized cell parameters.
+ """
+ # Make a copy of the atoms object
+ atoms = atoms.copy()
+
+ # Load Alchemical MACE calculator and relax the structure
+ calc = AlchemicalMACECalculator(
+ atoms, alchemical_pairs, alchemical_weights, device=device, model=model
+ )
+ atoms.set_calculator(calc)
+ atoms.set_masses(calc.get_alchemical_atomic_masses())
+ atoms = ExpCellFilter(atoms)
+ optimizer = FIRE(atoms)
+ optimizer.run(fmax=kwargs.get("fmax", 0.01), steps=kwargs.get("steps", 500))
+
+ # Return the optimized cell parameters
+ return atoms.atoms.get_cell().cellpar()
diff --git a/alchemical_mace/model.py b/alchemical_mace/model.py
new file mode 100644
index 0000000..e9e8e62
--- /dev/null
+++ b/alchemical_mace/model.py
@@ -0,0 +1,537 @@
+import ast
+from collections import namedtuple
+from typing import Dict, List, Optional, Sequence, Tuple
+
+import ase
+import numpy as np
+import torch
+import torch.nn.functional as F
+from e3nn import o3
+from e3nn.util.jit import compile_mode
+from mace import modules, tools
+from mace.calculators import mace_mp
+from mace.data.neighborhood import get_neighborhood
+from mace.modules import RealAgnosticResidualInteractionBlock, ScaleShiftMACE
+from mace.modules.utils import get_edge_vectors_and_lengths, get_symmetric_displacement
+from mace.tools import (
+ AtomicNumberTable,
+ atomic_numbers_to_indices,
+ to_one_hot,
+ torch_geometric,
+ utils,
+)
+from mace.tools.scatter import scatter_sum
+
+################################################################################
+# Alchemy manager class for handling alchemical weights
+################################################################################
+
+AlchemicalPair = namedtuple("AlchemicalPair", ["atom_index", "atomic_number"])
+
+
+class AlchemyManager(torch.nn.Module):
+ """
+ Class for managing alchemical weights and building alchemical graphs for MACE.
+ """
+
+ def __init__(
+ self,
+ atoms: ase.Atoms,
+ alchemical_pairs: Sequence[Sequence[Tuple[int, int]]],
+ alchemical_weights: torch.Tensor,
+ z_table: AtomicNumberTable,
+ r_max: float,
+ ):
+ """
+ Initialize the alchemy manager.
+
+ Args:
+ atoms: ASE atoms object
+ alchemical_pairs: List of lists of tuples, where each tuple contains
+ the atom index and atomic number of an alchemical atom
+ alchemical_weights: Tensor of alchemical weights
+ z_table: Atomic number table
+ r_max: Maximum cutoff radius for the alchemical graph
+ """
+ super().__init__()
+ self.alchemical_weights = torch.nn.Parameter(alchemical_weights)
+ self.r_max = r_max
+
+ # Process alchemical pairs into atom indices and atomic numbers
+ # Alchemical weights are 1-indexed, 0 is reserved for non-alchemical atoms
+ alchemical_atom_indices = []
+ alchemical_atomic_numbers = []
+ alchemical_weight_indices = []
+
+ for weight_idx, pairs in enumerate(alchemical_pairs):
+ for pair in pairs:
+ alchemical_atom_indices.append(pair.atom_index)
+ alchemical_atomic_numbers.append(pair.atomic_number)
+ alchemical_weight_indices.append(weight_idx + 1)
+
+ non_alchemical_atom_indices = [
+ i for i in range(len(atoms)) if i not in alchemical_atom_indices
+ ]
+ non_alchemical_atomic_numbers = atoms.get_atomic_numbers()[
+ non_alchemical_atom_indices
+ ].tolist()
+ non_alchemical_weight_indices = [0] * len(non_alchemical_atom_indices)
+
+ self.atom_indices = alchemical_atom_indices + non_alchemical_atom_indices
+ self.atomic_numbers = alchemical_atomic_numbers + non_alchemical_atomic_numbers
+ self.weight_indices = alchemical_weight_indices + non_alchemical_weight_indices
+
+ self.atom_indices = np.array(self.atom_indices)
+ self.atomic_numbers = np.array(self.atomic_numbers)
+ self.weight_indices = np.array(self.weight_indices)
+
+ sort_idx = np.argsort(self.atom_indices)
+ self.atom_indices = self.atom_indices[sort_idx]
+ self.atomic_numbers = self.atomic_numbers[sort_idx]
+ self.weight_indices = self.weight_indices[sort_idx]
+
+ # Array to map original atom indices to alchemical indices
+ # -1 means the atom does not have a corresponding alchemical atom
+ # [n_atoms, n_weights + 1]
+ self.original_to_alchemical_index = np.full(
+ (len(atoms), len(alchemical_pairs) + 1), -1
+ )
+ for i, (atom_idx, weight_idx) in enumerate(
+ zip(self.atom_indices, self.weight_indices)
+ ):
+ self.original_to_alchemical_index[atom_idx, weight_idx] = i
+
+ self.is_original_atom_alchemical = np.any(
+ self.original_to_alchemical_index[:, 1:] != -1, axis=1
+ )
+
+ # Extract common node features
+ z_indices = atomic_numbers_to_indices(self.atomic_numbers, z_table=z_table)
+ node_attrs = to_one_hot(
+ torch.tensor(z_indices, dtype=torch.long).unsqueeze(-1),
+ num_classes=len(z_table),
+ )
+ self.register_buffer("node_attrs", node_attrs)
+ self.pbc = atoms.get_pbc()
+
+ def forward(
+ self,
+ positions: torch.Tensor,
+ cell: torch.Tensor,
+ ) -> Dict[str, torch.Tensor]:
+ """
+ Build an alchemical graph for the given positions and cell.
+
+ Args:
+ positions: Tensor of atomic positions
+ cell: Tensor of cell vectors
+
+ Returns:
+ Dictionary containing the alchemical graph data
+ """
+
+ # Build original atom graph
+ orig_edge_index, shifts, unit_shifts = get_neighborhood(
+ positions=positions.detach().cpu().numpy(),
+ cutoff=self.r_max,
+ pbc=self.pbc,
+ cell=cell.detach().cpu().numpy(),
+ )
+
+ # Extend edges to alchemical pairs
+ edge_index = []
+ orig_edge_loc = []
+ edge_weight_indices = []
+
+ is_alchemical = self.is_original_atom_alchemical[orig_edge_index]
+ src_non_dst_non = ~is_alchemical[0] & ~is_alchemical[1]
+ src_non_dst_alch = ~is_alchemical[0] & is_alchemical[1]
+ src_alch_dst_non = is_alchemical[0] & ~is_alchemical[1]
+ src_alch_dst_alch = is_alchemical[0] & is_alchemical[1]
+
+ # Both non-alchemical: keep as is
+ _orig_edge_index = orig_edge_index[:, src_non_dst_non]
+ edge_index.append(self.original_to_alchemical_index[_orig_edge_index, 0])
+ orig_edge_loc.append(np.where(src_non_dst_non)[0])
+ edge_weight_indices.append(np.zeros_like(_orig_edge_index[0]))
+
+ # Source non-alchemical, destination alchemical: pair all, weights are 1
+ _src, _dst = orig_edge_index[:, src_non_dst_alch]
+ _orig_edge_loc = np.where(src_non_dst_alch)[0]
+ _src = self.original_to_alchemical_index[_src, 0]
+ _dst = self.original_to_alchemical_index[_dst, :]
+ _dst_mask = _dst != -1
+ _dst = _dst[_dst_mask]
+ _repeat = _dst_mask.sum(axis=1)
+ _src = np.repeat(_src, _repeat)
+ edge_index.append(np.stack((_src, _dst), axis=0))
+ orig_edge_loc.append(np.repeat(_orig_edge_loc, _repeat))
+ edge_weight_indices.append(np.zeros_like(_src))
+
+ # Source alchemical, destination non-alchemical: pair all, follow src weights
+ _src, _dst = orig_edge_index[:, src_alch_dst_non]
+ _orig_edge_loc = np.where(src_alch_dst_non)[0]
+ _src = self.original_to_alchemical_index[_src, :]
+ _dst = self.original_to_alchemical_index[_dst, 0]
+ _src_mask = _src != -1
+ _src = _src[_src_mask]
+ _repeat = _src_mask.sum(axis=1)
+ _dst = np.repeat(_dst, _repeat)
+ edge_index.append(np.stack((_src, _dst), axis=0))
+ orig_edge_loc.append(np.repeat(_orig_edge_loc, _repeat))
+ edge_weight_indices.append(np.where(_src_mask)[1])
+
+ # Both alchemical: pair according to alchemical indices, weights are 1
+ _orig_edge_index = orig_edge_index[:, src_alch_dst_alch]
+ _orig_edge_loc = np.where(src_alch_dst_alch)[0]
+ _alch_edge_index = self.original_to_alchemical_index[_orig_edge_index, :]
+ _idx = np.where((_alch_edge_index != -1).all(axis=0))
+ edge_index.append(_alch_edge_index[:, _idx[0], _idx[1]])
+ orig_edge_loc.append(_orig_edge_loc[_idx[0]])
+ edge_weight_indices.append(np.zeros_like(_idx[0]))
+
+ # Collect all edges
+ edge_index = np.concatenate(edge_index, axis=1)
+ orig_edge_loc = np.concatenate(orig_edge_loc)
+ edge_weight_indices = np.concatenate(edge_weight_indices)
+
+ # Convert to torch tensors
+ edge_index = torch.tensor(edge_index, dtype=torch.long)
+ shifts = torch.tensor(shifts[orig_edge_loc], dtype=torch.float32)
+ unit_shifts = torch.tensor(unit_shifts[orig_edge_loc], dtype=torch.float32)
+
+ # Alchemical weights for nodes and edges
+ weights = F.pad(self.alchemical_weights, (1, 0), "constant", 1.0)
+ node_weights = weights[self.weight_indices]
+ edge_weights = weights[edge_weight_indices]
+
+ # Build data batch
+ atomic_data = torch_geometric.data.Data(
+ num_nodes=len(self.atom_indices),
+ edge_index=edge_index,
+ node_attrs=self.node_attrs,
+ positions=positions[self.atom_indices],
+ shifts=shifts,
+ unit_shifts=unit_shifts,
+ cell=cell,
+ node_weights=node_weights,
+ edge_weights=edge_weights,
+ node_atom_indices=torch.tensor(self.atom_indices, dtype=torch.long),
+ )
+ data_loader = torch_geometric.dataloader.DataLoader(
+ dataset=[atomic_data],
+ batch_size=1,
+ shuffle=False,
+ drop_last=False,
+ )
+ batch = next(iter(data_loader))
+
+ return batch
+
+
+################################################################################
+# Alchemical MACE model
+################################################################################
+
+# get_outputs function from mace.modules.utils is modified to calculate also
+# the alchemical gradients
+
+
+def get_outputs(
+ energy: torch.Tensor,
+ positions: torch.Tensor,
+ displacement: torch.Tensor,
+ cell: torch.Tensor,
+ node_weights: torch.Tensor,
+ edge_weights: torch.Tensor,
+ retain_graph: bool = False,
+ create_graph: bool = False,
+ compute_force: bool = True,
+ compute_stress: bool = False,
+ compute_alchemical_grad: bool = False,
+) -> Tuple[
+ Optional[torch.Tensor],
+ Optional[torch.Tensor],
+ Optional[torch.Tensor],
+ Optional[torch.Tensor],
+ Optional[torch.Tensor],
+]:
+ grad_outputs: List[Optional[torch.Tensor]] = [torch.ones_like(energy)]
+ if not compute_force:
+ return None, None, None, None, None
+ inputs = [positions]
+ if compute_stress:
+ inputs.append(displacement)
+ if compute_alchemical_grad:
+ inputs.extend([node_weights, edge_weights])
+ gradients = torch.autograd.grad(
+ outputs=[energy],
+ inputs=inputs,
+ grad_outputs=grad_outputs,
+ retain_graph=retain_graph,
+ create_graph=create_graph,
+ allow_unused=True,
+ )
+
+ forces = gradients[0]
+ stress = torch.zeros_like(displacement)
+ virials = gradients[1] if compute_stress else None
+ if compute_alchemical_grad:
+ node_grad, edge_grad = gradients[-2], gradients[-1]
+ else:
+ node_grad, edge_grad = None, None
+ if compute_stress and virials is not None:
+ cell = cell.view(-1, 3, 3)
+ volume = torch.einsum(
+ "zi,zi->z",
+ cell[:, 0, :],
+ torch.cross(cell[:, 1, :], cell[:, 2, :], dim=1),
+ ).unsqueeze(-1)
+ stress = virials / volume.view(-1, 1, 1)
+
+ if forces is not None:
+ forces = -1 * forces
+ if virials is not None:
+ virials = -1 * virials
+ return forces, virials, stress, node_grad, edge_grad
+
+
+@compile_mode("script")
+class AlchemicalResidualInteractionBlock(RealAgnosticResidualInteractionBlock):
+ def forward(
+ self,
+ node_attrs: torch.Tensor,
+ node_feats: torch.Tensor,
+ edge_attrs: torch.Tensor,
+ edge_feats: torch.Tensor,
+ edge_index: torch.Tensor,
+ edge_weights: torch.Tensor, # alchemy
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
+ sender = edge_index[0]
+ receiver = edge_index[1]
+ num_nodes = node_feats.shape[0]
+ sc = self.skip_tp(node_feats, node_attrs)
+ node_feats = self.linear_up(node_feats)
+ tp_weights = self.conv_tp_weights(edge_feats)
+ tp_weights = tp_weights * edge_weights[:, None] # alchemy
+ mji = self.conv_tp(
+ node_feats[sender], edge_attrs, tp_weights
+ ) # [n_edges, irreps]
+ message = scatter_sum(
+ src=mji, index=receiver, dim=0, dim_size=num_nodes
+ ) # [n_nodes, irreps]
+ message = self.linear(message) / self.avg_num_neighbors
+ return (
+ self.reshape(message),
+ sc,
+ ) # [n_nodes, channels, (lmax + 1)**2]
+
+
+@compile_mode("script")
+class AlchemicalMACE(ScaleShiftMACE):
+ def forward(
+ self,
+ data: Dict[str, torch.Tensor],
+ retain_graph: bool = False, # alchemy
+ create_graph: bool = False, # alchemy
+ compute_force: bool = True,
+ compute_stress: bool = False,
+ compute_displacement: bool = False,
+ compute_alchemical_grad: bool = False, # alchemy
+ map_to_original_atoms: bool = True, # alchemy
+ ) -> Dict[str, Optional[torch.Tensor]]:
+ # Setup
+ data["positions"].requires_grad_(True)
+ data["node_attrs"].requires_grad_(True)
+ num_graphs = data["ptr"].numel() - 1
+ displacement = torch.zeros(
+ (num_graphs, 3, 3),
+ dtype=data["positions"].dtype,
+ device=data["positions"].device,
+ )
+ if compute_stress or compute_displacement:
+ (
+ data["positions"],
+ data["shifts"],
+ displacement,
+ ) = get_symmetric_displacement(
+ positions=data["positions"],
+ unit_shifts=data["unit_shifts"],
+ cell=data["cell"],
+ edge_index=data["edge_index"],
+ num_graphs=num_graphs,
+ batch=data["batch"],
+ )
+
+ # Atomic energies
+ node_e0 = self.atomic_energies_fn(data["node_attrs"])
+ node_e0 = node_e0 * data["node_weights"] # alchemy
+ e0 = scatter_sum(
+ src=node_e0, index=data["batch"], dim=-1, dim_size=num_graphs
+ ) # [n_graphs,]
+
+ # Embeddings
+ node_feats = self.node_embedding(data["node_attrs"])
+ vectors, lengths = get_edge_vectors_and_lengths(
+ positions=data["positions"],
+ edge_index=data["edge_index"],
+ shifts=data["shifts"],
+ )
+ edge_attrs = self.spherical_harmonics(vectors)
+ edge_feats = self.radial_embedding(lengths)
+
+ # Interactions
+ node_es_list = []
+ node_feats_list = []
+ for interaction, product, readout in zip(
+ self.interactions, self.products, self.readouts
+ ):
+ node_feats, sc = interaction(
+ node_attrs=data["node_attrs"],
+ node_feats=node_feats,
+ edge_attrs=edge_attrs,
+ edge_feats=edge_feats,
+ edge_index=data["edge_index"],
+ edge_weights=data["edge_weights"], # alchemy
+ )
+ node_feats = product(
+ node_feats=node_feats, sc=sc, node_attrs=data["node_attrs"]
+ )
+ node_feats_list.append(node_feats)
+ node_es_list.append(readout(node_feats).squeeze(-1)) # {[n_nodes, ], }
+
+ # Concatenate node features
+ # node_feats_out = torch.cat(node_feats_list, dim=-1)
+
+ # Sum over interactions
+ node_inter_es = torch.sum(
+ torch.stack(node_es_list, dim=0), dim=0
+ ) # [n_nodes, ]
+ node_inter_es = self.scale_shift(node_inter_es)
+ node_inter_es = node_inter_es * data["node_weights"] # alchemy
+
+ # Sum over nodes in graph
+ inter_e = scatter_sum(
+ src=node_inter_es, index=data["batch"], dim=-1, dim_size=num_graphs
+ ) # [n_graphs,]
+
+ # Add E_0 and (scaled) interaction energy
+ total_energy = e0 + inter_e
+ node_energy = node_e0 + node_inter_es
+
+ forces, virials, stress, node_grad, edge_grad = get_outputs(
+ energy=total_energy, # alchemy
+ positions=data["positions"],
+ displacement=displacement,
+ cell=data["cell"],
+ node_weights=data["node_weights"], # alchemy
+ edge_weights=data["edge_weights"], # alchemy
+ retain_graph=retain_graph, # alchemy
+ create_graph=create_graph, # alchemy
+ compute_force=compute_force,
+ # compute_virials=compute_virials, # alchemy
+ compute_stress=compute_stress,
+ compute_alchemical_grad=compute_alchemical_grad, # alchemy
+ )
+
+ # Map to original atoms (node energies and forces): alchemy
+ if map_to_original_atoms:
+ # Note: we're not giving the dim_size, as we assume that all
+ # original atoms are present in the batch
+ node_index = data["node_atom_indices"]
+ node_energy = scatter_sum(src=node_energy, dim=0, index=node_index)
+ if compute_force:
+ forces = scatter_sum(src=forces, dim=0, index=node_index)
+
+ output = {
+ "energy": total_energy,
+ "node_energy": node_energy,
+ "interaction_energy": inter_e,
+ "forces": forces,
+ "virials": virials,
+ "stress": stress,
+ "displacement": displacement,
+ "node_grad": node_grad,
+ "edge_grad": edge_grad,
+ }
+
+ return output
+
+
+################################################################################
+# Alchemical MACE universal model
+################################################################################
+
+
+def alchemical_mace_mp(
+ model: str,
+ device: str,
+ default_dtype: str = "float32",
+):
+ """
+ Load a pre-trained alchemical MACE model.
+
+ Args:
+ model: Model size (small, medium)
+ device: Device to load the model onto
+ default_dtype: Default data type for the model
+
+ Returns:
+ Alchemical MACE model
+ """
+
+ # Load foundation MACE model and extract initial parameters
+ assert model in ("small", "medium") # TODO: support large model
+ mace = mace_mp(model=model, device=device, default_dtype=default_dtype).models[0]
+ atomic_energies = mace.atomic_energies_fn.atomic_energies.detach().clone()
+ z_table = utils.AtomicNumberTable([int(z) for z in mace.atomic_numbers])
+ atomic_inter_scale = mace.scale_shift.scale.detach().clone()
+ atomic_inter_shift = mace.scale_shift.shift.detach().clone()
+
+ # Prepare arguments for building the model
+ placeholder_args = ["--name", "None", "--train_file", "None"]
+ args = tools.build_default_arg_parser().parse_args(placeholder_args)
+ args.max_L = {"small": 0, "medium": 1, "large": 2}[model]
+ args.num_channels = 128
+ args.hidden_irreps = o3.Irreps(
+ (args.num_channels * o3.Irreps.spherical_harmonics(args.max_L))
+ .sort()
+ .irreps.simplify()
+ )
+
+ # Build the alchemical MACE model
+ model = AlchemicalMACE(
+ r_max=6.0,
+ num_bessel=10,
+ num_polynomial_cutoff=5,
+ max_ell=3,
+ interaction_cls=AlchemicalResidualInteractionBlock,
+ interaction_cls_first=AlchemicalResidualInteractionBlock,
+ num_interactions=2,
+ num_elements=len(z_table),
+ hidden_irreps=o3.Irreps(args.hidden_irreps),
+ MLP_irreps=o3.Irreps(args.MLP_irreps),
+ atomic_energies=atomic_energies,
+ avg_num_neighbors=args.avg_num_neighbors,
+ atomic_numbers=z_table.zs,
+ correlation=args.correlation,
+ gate=modules.gate_dict[args.gate],
+ radial_MLP=ast.literal_eval(args.radial_MLP),
+ radial_type=args.radial_type,
+ atomic_inter_scale=atomic_inter_scale,
+ atomic_inter_shift=atomic_inter_shift,
+ )
+
+ # Load foundation model parameters
+ model.load_state_dict(mace.state_dict(), strict=True)
+ for i in range(int(model.num_interactions)):
+ model.interactions[i].avg_num_neighbors = mace.interactions[i].avg_num_neighbors
+ model = model.to(device)
+ return model
+
+
+def get_z_table_and_r_max(model: ScaleShiftMACE) -> Tuple[AtomicNumberTable, float]:
+ """Extract the atomic number table and maximum cutoff radius from a MACE model."""
+ z_table = AtomicNumberTable([int(z) for z in model.atomic_numbers])
+ r_max = model.r_max.item()
+ return z_table, r_max
diff --git a/alchemical_mace/optimize.py b/alchemical_mace/optimize.py
new file mode 100644
index 0000000..87171be
--- /dev/null
+++ b/alchemical_mace/optimize.py
@@ -0,0 +1,29 @@
+import torch
+from typing import Union, Iterable, Dict, Any
+
+
+class ExponentiatedGradientDescent(torch.optim.Optimizer):
+ """
+ Implements Exponentiated Gradient Descent.
+
+ Args:
+ params (iterable of torch.Tensor or dict): iterable of parameters to optimize or
+ dicts defining parameter groups.
+ lr (float, optional): learning rate. Defaults to 1e-3.
+ eps (float, optional): small constant for numerical stability. Defaults to 1e-8.
+ """
+ def __init__(
+ self,
+ params: Union[Iterable[torch.Tensor], Iterable[Dict[str, Any]]],
+ lr: float = 1e-3,
+ eps: float = 1e-8,
+ ):
+ super().__init__(params, defaults=dict(lr=lr, eps=eps))
+
+ def step(self):
+ for group in self.param_groups:
+ for p in group["params"]:
+ if p.grad is None:
+ continue
+ p.data.mul_(torch.exp(-group["lr"] * p.grad))
+ p.data.div_(p.data.sum() + group["eps"])
diff --git a/alchemical_mace/utils.py b/alchemical_mace/utils.py
new file mode 100644
index 0000000..d37cf9c
--- /dev/null
+++ b/alchemical_mace/utils.py
@@ -0,0 +1,43 @@
+import os
+from contextlib import ExitStack, contextmanager, redirect_stderr, redirect_stdout
+
+from ase import Atoms
+
+
+@contextmanager
+def suppress_print(out: bool = True, err: bool = False):
+ """Suppress stdout and/or stderr."""
+
+ with ExitStack() as stack:
+ devnull = stack.enter_context(open(os.devnull, "w"))
+ if out:
+ stack.enter_context(redirect_stdout(devnull))
+ if err:
+ stack.enter_context(redirect_stderr(devnull))
+ yield
+
+
+# From CHGNet
+def upper_triangular_cell(atoms: Atoms):
+ """Transform to upper-triangular cell."""
+ import numpy as np
+ from ase.md.npt import NPT
+
+ if NPT._isuppertriangular(atoms.get_cell()):
+ return
+
+ a, b, c, alpha, beta, gamma = atoms.cell.cellpar()
+ angles = np.radians((alpha, beta, gamma))
+ sin_a, sin_b, _sin_g = np.sin(angles)
+ cos_a, cos_b, cos_g = np.cos(angles)
+ cos_p = (cos_g - cos_a * cos_b) / (sin_a * sin_b)
+ cos_p = np.clip(cos_p, -1, 1)
+ sin_p = (1 - cos_p**2) ** 0.5
+
+ new_basis = [
+ (a * sin_b * sin_p, a * sin_b * cos_p, a * cos_b),
+ (0, b * sin_a, b * cos_a),
+ (0, 0, c),
+ ]
+
+ atoms.set_cell(new_basis, scale_atoms=True)
diff --git a/data/structures/AlN_hex.cif b/data/structures/AlN_hex.cif
new file mode 100644
index 0000000..ef9bf10
--- /dev/null
+++ b/data/structures/AlN_hex.cif
@@ -0,0 +1,30 @@
+# generated using pymatgen
+data_AlN
+_symmetry_space_group_name_H-M 'P 1'
+_cell_length_a 3.12664153
+_cell_length_b 3.12664143
+_cell_length_c 5.00715332
+_cell_angle_alpha 90.00000043
+_cell_angle_beta 89.99999986
+_cell_angle_gamma 119.99999965
+_symmetry_Int_Tables_number 1
+_chemical_formula_structural AlN
+_chemical_formula_sum 'Al2 N2'
+_cell_volume 42.39139351
+_cell_formula_units_Z 2
+loop_
+ _symmetry_equiv_pos_site_id
+ _symmetry_equiv_pos_as_xyz
+ 1 'x, y, z'
+loop_
+ _atom_site_type_symbol
+ _atom_site_label
+ _atom_site_symmetry_multiplicity
+ _atom_site_fract_x
+ _atom_site_fract_y
+ _atom_site_fract_z
+ _atom_site_occupancy
+ Al Al0 1 0.66666667 0.33333333 0.49932157 1
+ Al Al1 1 0.33333333 0.66666667 0.99932156 1
+ N N2 1 0.66666667 0.33333333 0.88067843 1
+ N N3 1 0.33333333 0.66666667 0.38067844 1
diff --git a/data/structures/BiSBr.cif b/data/structures/BiSBr.cif
new file mode 100644
index 0000000..e9bbad3
--- /dev/null
+++ b/data/structures/BiSBr.cif
@@ -0,0 +1,38 @@
+# generated using pymatgen
+data_BiSBr
+_symmetry_space_group_name_H-M 'P 1'
+_cell_length_a 4.10679210
+_cell_length_b 8.22484633
+_cell_length_c 11.05701800
+_cell_angle_alpha 89.99999573
+_cell_angle_beta 90.00000347
+_cell_angle_gamma 90.00001971
+_symmetry_Int_Tables_number 1
+_chemical_formula_structural BiSBr
+_chemical_formula_sum 'Bi4 S4 Br4'
+_cell_volume 373.48101173
+_cell_formula_units_Z 4
+loop_
+ _symmetry_equiv_pos_site_id
+ _symmetry_equiv_pos_as_xyz
+ 1 'x, y, z'
+loop_
+ _atom_site_type_symbol
+ _atom_site_label
+ _atom_site_symmetry_multiplicity
+ _atom_site_fract_x
+ _atom_site_fract_y
+ _atom_site_fract_z
+ _atom_site_occupancy
+ Bi Bi0 1 0.24999992 0.88556832 0.87177467 1
+ Bi Bi1 1 0.75000010 0.11443170 0.12822534 1
+ Bi Bi2 1 0.24999988 0.38556830 0.62822532 1
+ Bi Bi3 1 0.75000009 0.61443169 0.37177466 1
+ S S4 1 0.74999990 0.82391374 0.03167750 1
+ S S5 1 0.25000008 0.67608625 0.53167751 1
+ S S6 1 0.74999991 0.32391375 0.46832249 1
+ S S7 1 0.25000010 0.17608625 0.96832251 1
+ Br Br8 1 0.24999990 0.96479764 0.30006197 1
+ Br Br9 1 0.75000009 0.53520235 0.80006197 1
+ Br Br10 1 0.24999991 0.46479765 0.19993803 1
+ Br Br11 1 0.75000013 0.03520235 0.69993804 1
diff --git a/data/structures/CeO2.cif b/data/structures/CeO2.cif
new file mode 100644
index 0000000..ad4b557
--- /dev/null
+++ b/data/structures/CeO2.cif
@@ -0,0 +1,38 @@
+# generated using pymatgen
+data_CeO2
+_symmetry_space_group_name_H-M 'P 1'
+_cell_length_a 5.46789061
+_cell_length_b 5.46789009
+_cell_length_c 5.46788984
+_cell_angle_alpha 90.00000180
+_cell_angle_beta 90.00000282
+_cell_angle_gamma 90.00000072
+_symmetry_Int_Tables_number 1
+_chemical_formula_structural CeO2
+_chemical_formula_sum 'Ce4 O8'
+_cell_volume 163.47801292
+_cell_formula_units_Z 4
+loop_
+ _symmetry_equiv_pos_site_id
+ _symmetry_equiv_pos_as_xyz
+ 1 'x, y, z'
+loop_
+ _atom_site_type_symbol
+ _atom_site_label
+ _atom_site_symmetry_multiplicity
+ _atom_site_fract_x
+ _atom_site_fract_y
+ _atom_site_fract_z
+ _atom_site_occupancy
+ Ce Ce0 1 0.00000000 -0.00000000 -0.00000000 1
+ Ce Ce1 1 0.00000000 0.50000000 0.50000000 1
+ Ce Ce2 1 0.50000000 0.00000000 0.50000000 1
+ Ce Ce3 1 0.50000000 0.50000000 0.00000000 1
+ O O4 1 0.25000000 0.75000000 0.74999999 1
+ O O5 1 0.74999999 0.25000000 0.25000000 1
+ O O6 1 0.25000000 0.25000000 0.25000000 1
+ O O7 1 0.75000000 0.75000000 0.75000000 1
+ O O8 1 0.74999999 0.75000000 0.25000000 1
+ O O9 1 0.25000000 0.25000000 0.74999999 1
+ O O10 1 0.75000000 0.25000000 0.75000000 1
+ O O11 1 0.25000000 0.75000000 0.25000000 1
diff --git a/data/structures/CsPbI3_alpha.cif b/data/structures/CsPbI3_alpha.cif
new file mode 100644
index 0000000..8f0984d
--- /dev/null
+++ b/data/structures/CsPbI3_alpha.cif
@@ -0,0 +1,31 @@
+# generated using pymatgen
+data_CsPbI3
+_symmetry_space_group_name_H-M 'P 1'
+_cell_length_a 6.37904246
+_cell_length_b 6.37904251
+_cell_length_c 6.37904247
+_cell_angle_alpha 90.00000013
+_cell_angle_beta 89.99999979
+_cell_angle_gamma 89.99999953
+_symmetry_Int_Tables_number 1
+_chemical_formula_structural CsPbI3
+_chemical_formula_sum 'Cs1 Pb1 I3'
+_cell_volume 259.57716340
+_cell_formula_units_Z 1
+loop_
+ _symmetry_equiv_pos_site_id
+ _symmetry_equiv_pos_as_xyz
+ 1 'x, y, z'
+loop_
+ _atom_site_type_symbol
+ _atom_site_label
+ _atom_site_symmetry_multiplicity
+ _atom_site_fract_x
+ _atom_site_fract_y
+ _atom_site_fract_z
+ _atom_site_occupancy
+ Cs Cs0 1 0.00000000 0.00000000 -0.00000000 1
+ I I1 1 0.50000000 0.50000000 -0.00000000 1
+ I I2 1 0.50000000 -0.00000000 0.50000000 1
+ I I3 1 0.00000000 0.50000000 0.50000000 1
+ Pb Pb4 1 0.50000000 0.50000000 0.50000000 1
diff --git a/data/structures/CsPbI3_delta.cif b/data/structures/CsPbI3_delta.cif
new file mode 100644
index 0000000..1313893
--- /dev/null
+++ b/data/structures/CsPbI3_delta.cif
@@ -0,0 +1,46 @@
+# generated using pymatgen
+data_CsPbI3
+_symmetry_space_group_name_H-M 'P 1'
+_cell_length_a 4.91187765
+_cell_length_b 10.78602159
+_cell_length_c 18.12941751
+_cell_angle_alpha 90.00000734
+_cell_angle_beta 89.99999876
+_cell_angle_gamma 89.99999900
+_symmetry_Int_Tables_number 1
+_chemical_formula_structural CsPbI3
+_chemical_formula_sum 'Cs4 Pb4 I12'
+_cell_volume 960.48962177
+_cell_formula_units_Z 4
+loop_
+ _symmetry_equiv_pos_site_id
+ _symmetry_equiv_pos_as_xyz
+ 1 'x, y, z'
+loop_
+ _atom_site_type_symbol
+ _atom_site_label
+ _atom_site_symmetry_multiplicity
+ _atom_site_fract_x
+ _atom_site_fract_y
+ _atom_site_fract_z
+ _atom_site_occupancy
+ Cs Cs0 1 0.75000002 0.57436836 0.17329066 1
+ Cs Cs1 1 0.25000000 0.42563165 0.82670935 1
+ Cs Cs2 1 0.75000002 0.07436835 0.32670935 1
+ Cs Cs3 1 0.24999999 0.92563166 0.67329065 1
+ I I4 1 0.74999999 0.20706041 0.71324789 1
+ I I5 1 0.25000000 0.79293960 0.28675211 1
+ I I6 1 0.75000001 0.97331566 0.10968932 1
+ I I7 1 0.24999998 0.02668433 0.89031069 1
+ I I8 1 0.74999999 0.47331567 0.39031069 1
+ I I9 1 0.25000001 0.52668433 0.60968931 1
+ I I10 1 0.25000000 0.66460269 0.00429655 1
+ I I11 1 0.25000000 0.16460269 0.49570344 1
+ I I12 1 0.25000000 0.29293960 0.21324788 1
+ I I13 1 0.75000001 0.83539729 0.50429655 1
+ I I14 1 0.75000000 0.33539731 0.99570345 1
+ I I15 1 0.75000001 0.70706040 0.78675212 1
+ Pb Pb16 1 0.74999998 0.83743203 0.94239546 1
+ Pb Pb17 1 0.24999999 0.16256797 0.05760453 1
+ Pb Pb18 1 0.75000002 0.33743205 0.55760454 1
+ Pb Pb19 1 0.24999998 0.66256796 0.44239545 1
diff --git a/data/structures/CsSnI3_alpha.cif b/data/structures/CsSnI3_alpha.cif
new file mode 100644
index 0000000..87f7b76
--- /dev/null
+++ b/data/structures/CsSnI3_alpha.cif
@@ -0,0 +1,31 @@
+# generated using pymatgen
+data_CsSnI3
+_symmetry_space_group_name_H-M 'P 1'
+_cell_length_a 6.27081715
+_cell_length_b 6.27081724
+_cell_length_c 6.27081711
+_cell_angle_alpha 89.99999964
+_cell_angle_beta 89.99999969
+_cell_angle_gamma 89.99999972
+_symmetry_Int_Tables_number 1
+_chemical_formula_structural CsSnI3
+_chemical_formula_sum 'Cs1 Sn1 I3'
+_cell_volume 246.58827085
+_cell_formula_units_Z 1
+loop_
+ _symmetry_equiv_pos_site_id
+ _symmetry_equiv_pos_as_xyz
+ 1 'x, y, z'
+loop_
+ _atom_site_type_symbol
+ _atom_site_label
+ _atom_site_symmetry_multiplicity
+ _atom_site_fract_x
+ _atom_site_fract_y
+ _atom_site_fract_z
+ _atom_site_occupancy
+ Cs Cs0 1 0.00000000 0.00000000 -0.00000000 1
+ I I1 1 0.50000000 0.50000000 0.00000000 1
+ I I2 1 0.50000000 -0.00000000 0.50000000 1
+ I I3 1 -0.00000000 0.50000000 0.50000000 1
+ Sn Sn4 1 0.50000000 0.50000000 0.50000000 1
diff --git a/data/structures/CsSnI3_delta.cif b/data/structures/CsSnI3_delta.cif
new file mode 100644
index 0000000..e7054a0
--- /dev/null
+++ b/data/structures/CsSnI3_delta.cif
@@ -0,0 +1,46 @@
+# generated using pymatgen
+data_CsSnI3
+_symmetry_space_group_name_H-M 'P 1'
+_cell_length_a 4.84918880
+_cell_length_b 10.69167480
+_cell_length_c 18.19616575
+_cell_angle_alpha 89.99999030
+_cell_angle_beta 90.00000262
+_cell_angle_gamma 89.99999762
+_symmetry_Int_Tables_number 1
+_chemical_formula_structural CsSnI3
+_chemical_formula_sum 'Cs4 Sn4 I12'
+_cell_volume 943.39749344
+_cell_formula_units_Z 4
+loop_
+ _symmetry_equiv_pos_site_id
+ _symmetry_equiv_pos_as_xyz
+ 1 'x, y, z'
+loop_
+ _atom_site_type_symbol
+ _atom_site_label
+ _atom_site_symmetry_multiplicity
+ _atom_site_fract_x
+ _atom_site_fract_y
+ _atom_site_fract_z
+ _atom_site_occupancy
+ Cs Cs0 1 0.75000000 0.57552416 0.17186806 1
+ Cs Cs1 1 0.25000001 0.42447584 0.82813194 1
+ Cs Cs2 1 0.75000000 0.07552414 0.32813195 1
+ Cs Cs3 1 0.25000000 0.92447584 0.67186806 1
+ I I4 1 0.75000000 0.21397041 0.70606567 1
+ I I5 1 0.25000000 0.78602961 0.29393431 1
+ I I6 1 0.75000000 0.97123818 0.11227438 1
+ I I7 1 0.25000000 0.02876181 0.88772566 1
+ I I8 1 0.74999999 0.47123818 0.38772564 1
+ I I9 1 0.25000001 0.52876182 0.61227437 1
+ I I10 1 0.25000002 0.66813711 0.00020614 1
+ I I11 1 0.24999999 0.16813710 0.49979388 1
+ I I12 1 0.25000001 0.28602959 0.20606567 1
+ I I13 1 0.74999999 0.83186292 0.50020614 1
+ I I14 1 0.75000000 0.33186292 0.99979386 1
+ I I15 1 0.74999999 0.71397042 0.79393433 1
+ Sn Sn16 1 0.75000000 0.84420702 0.94393743 1
+ Sn Sn17 1 0.25000000 0.15579296 0.05606252 1
+ Sn Sn18 1 0.74999998 0.34420702 0.55606253 1
+ Sn Sn19 1 0.24999999 0.65579297 0.44393746 1
diff --git a/data/structures/Fe.cif b/data/structures/Fe.cif
new file mode 100644
index 0000000..18422ee
--- /dev/null
+++ b/data/structures/Fe.cif
@@ -0,0 +1,28 @@
+# generated using pymatgen
+data_Fe
+_symmetry_space_group_name_H-M 'P 1'
+_cell_length_a 2.86106543
+_cell_length_b 2.86106544
+_cell_length_c 2.86106538
+_cell_angle_alpha 90.00000018
+_cell_angle_beta 89.99999992
+_cell_angle_gamma 90.00000009
+_symmetry_Int_Tables_number 1
+_chemical_formula_structural Fe
+_chemical_formula_sum Fe2
+_cell_volume 23.41980977
+_cell_formula_units_Z 2
+loop_
+ _symmetry_equiv_pos_site_id
+ _symmetry_equiv_pos_as_xyz
+ 1 'x, y, z'
+loop_
+ _atom_site_type_symbol
+ _atom_site_label
+ _atom_site_symmetry_multiplicity
+ _atom_site_fract_x
+ _atom_site_fract_y
+ _atom_site_fract_z
+ _atom_site_occupancy
+ Fe Fe0 1 0.00000000 -0.00000000 0.00000000 1
+ Fe Fe1 1 0.50000000 0.50000000 0.50000000 1
diff --git a/data/structures/GaN_hex.cif b/data/structures/GaN_hex.cif
new file mode 100644
index 0000000..86d63dc
--- /dev/null
+++ b/data/structures/GaN_hex.cif
@@ -0,0 +1,30 @@
+# generated using pymatgen
+data_GaN
+_symmetry_space_group_name_H-M 'P 1'
+_cell_length_a 3.21192371
+_cell_length_b 3.21192377
+_cell_length_c 5.21628467
+_cell_angle_alpha 90.00000055
+_cell_angle_beta 90.00000024
+_cell_angle_gamma 119.99999991
+_symmetry_Int_Tables_number 1
+_chemical_formula_structural GaN
+_chemical_formula_sum 'Ga2 N2'
+_cell_volume 46.60391121
+_cell_formula_units_Z 2
+loop_
+ _symmetry_equiv_pos_site_id
+ _symmetry_equiv_pos_as_xyz
+ 1 'x, y, z'
+loop_
+ _atom_site_type_symbol
+ _atom_site_label
+ _atom_site_symmetry_multiplicity
+ _atom_site_fract_x
+ _atom_site_fract_y
+ _atom_site_fract_z
+ _atom_site_occupancy
+ Ga Ga0 1 0.66666667 0.33333333 0.49900050 1
+ Ga Ga1 1 0.33333333 0.66666667 0.99900050 1
+ N N2 1 0.66666667 0.33333333 0.87599950 1
+ N N3 1 0.33333333 0.66666667 0.37599950 1
diff --git a/data/structures/NaCl.cif b/data/structures/NaCl.cif
new file mode 100644
index 0000000..ece9ba4
--- /dev/null
+++ b/data/structures/NaCl.cif
@@ -0,0 +1,34 @@
+# generated using pymatgen
+data_NaCl
+_symmetry_space_group_name_H-M 'P 1'
+_cell_length_a 5.68304678
+_cell_length_b 5.68304679
+_cell_length_c 5.68304672
+_cell_angle_alpha 90.00000009
+_cell_angle_beta 90.00000015
+_cell_angle_gamma 90.00000005
+_symmetry_Int_Tables_number 1
+_chemical_formula_structural NaCl
+_chemical_formula_sum 'Na4 Cl4'
+_cell_volume 183.54547783
+_cell_formula_units_Z 4
+loop_
+ _symmetry_equiv_pos_site_id
+ _symmetry_equiv_pos_as_xyz
+ 1 'x, y, z'
+loop_
+ _atom_site_type_symbol
+ _atom_site_label
+ _atom_site_symmetry_multiplicity
+ _atom_site_fract_x
+ _atom_site_fract_y
+ _atom_site_fract_z
+ _atom_site_occupancy
+ Na Na0 1 -0.00000000 -0.00000000 -0.00000000 1
+ Na Na1 1 -0.00000000 0.50000000 0.50000000 1
+ Na Na2 1 0.50000000 -0.00000000 0.50000000 1
+ Na Na3 1 0.50000000 0.50000000 -0.00000000 1
+ Cl Cl4 1 0.00000000 0.00000000 0.50000000 1
+ Cl Cl5 1 -0.00000000 0.50000000 0.00000000 1
+ Cl Cl6 1 0.50000000 0.00000000 -0.00000000 1
+ Cl Cl7 1 0.50000000 0.50000000 0.50000000 1
diff --git a/notebooks/1_solid_solution.ipynb b/notebooks/1_solid_solution.ipynb
new file mode 100644
index 0000000..3c2101b
--- /dev/null
+++ b/notebooks/1_solid_solution.ipynb
@@ -0,0 +1,253 @@
+{
+ "cells": [
+ {
+ "cell_type": "code",
+ "execution_count": 1,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "import ase\n",
+ "import matplotlib.pyplot as plt\n",
+ "import numpy as np\n",
+ "from tqdm import tqdm\n",
+ "\n",
+ "from alchemical_mace.calculator import get_alchemical_optimized_cellpar\n",
+ "from alchemical_mace.model import AlchemicalPair\n",
+ "from alchemical_mace.utils import suppress_print\n",
+ "\n",
+ "plt.style.use(\"default\")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "### CeO2"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 2,
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stderr",
+ "output_type": "stream",
+ "text": [
+ "100%|██████████| 21/21 [01:15<00:00, 3.60s/it]\n"
+ ]
+ }
+ ],
+ "source": [
+ "# Default settings\n",
+ "model = \"medium\"\n",
+ "device = \"cpu\"\n",
+ "\n",
+ "# Load structure\n",
+ "atoms = ase.io.read(\"../data/structures/CeO2.cif\")\n",
+ "alch_elements = [\"Ce\", \"Sn\"]\n",
+ "alch_idx = [i for i, atom in enumerate(atoms) if atom.symbol in alch_elements]\n",
+ "alch_atomic_numbers = [ase.Atoms(el).numbers[0] for el in alch_elements]\n",
+ "alchemical_pairs = [\n",
+ " [AlchemicalPair(atom_index=idx, atomic_number=z) for idx in alch_idx]\n",
+ " for z in alch_atomic_numbers\n",
+ "]\n",
+ "\n",
+ "comp_grid = [[1 - x, x] for x in np.linspace(0, 0.5, 21)]\n",
+ "lat_params_CeSn = []\n",
+ "for comp in tqdm(comp_grid):\n",
+ " with suppress_print(out=True, err=True):\n",
+ " cellpar = get_alchemical_optimized_cellpar(\n",
+ " atoms, alchemical_pairs, comp, model=model, device=device\n",
+ " )\n",
+ " lat_params_CeSn.append(cellpar[:3])"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 3,
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stderr",
+ "output_type": "stream",
+ "text": [
+ "100%|██████████| 21/21 [01:35<00:00, 4.52s/it]\n"
+ ]
+ }
+ ],
+ "source": [
+ "# Load structure\n",
+ "alch_elements = [\"Ce\", \"Zr\"]\n",
+ "alch_idx = [i for i, atom in enumerate(atoms) if atom.symbol in alch_elements]\n",
+ "alch_atomic_numbers = [ase.Atoms(el).numbers[0] for el in alch_elements]\n",
+ "alchemical_pairs = [\n",
+ " [AlchemicalPair(atom_index=idx, atomic_number=z) for idx in alch_idx]\n",
+ " for z in alch_atomic_numbers\n",
+ "]\n",
+ "\n",
+ "comp_grid = [[1 - x, x] for x in np.linspace(0, 0.5, 21)]\n",
+ "lat_params_CeZr = []\n",
+ "for comp in tqdm(comp_grid):\n",
+ " with suppress_print(out=True, err=True):\n",
+ " cellpar = get_alchemical_optimized_cellpar(\n",
+ " atoms, alchemical_pairs, comp, model=model, device=device\n",
+ " )\n",
+ " lat_params_CeZr.append(cellpar[:3])"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 4,
+ "metadata": {},
+ "outputs": [
+ {
+ "data": {
+ "image/png": "iVBORw0KGgoAAAANSUhEUgAAATkAAAD/CAYAAACKJ6HsAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjguMCwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy81sbWrAAAACXBIWXMAAA9hAAAPYQGoP6dpAAA5pklEQVR4nO3deVxU9f7H8dcsrLKJ7IgoiuKGu/7cUlPTMrdMc8kltxYrq6u3xbqYWXrNlqtZlloulWZWZmq44JJ7aaEIrigiAqLsiKzz/f1xgkQBQQfOzPB9Ph7zuHdmzpz5cIK355zvphFCCCRJkiyUVu0CJEmSqpIMOUmSLJoMOUmSLJoMOUmSLJoMOUmSLJoMOUmSLJoMOUmSLJpe7QJMlcFgID4+HkdHRzQajdrlSJJ0CyEEmZmZ+Pj4oNWWf64mQ64M8fHx+Pn5qV2GJEnluHz5MnXr1i13GxlyZXB0dASUg+jk5KRyNZIk3SojIwM/P7/iv9PyyJArQ9ElqpOTkww5STJRFbmVJBseJEmyaDLkJEmyaDLkJEmyaDLkJEmyaDLk7lP6zXzSonYhDAa1S5EkqRSydfU+ndy0iK6n5vCDoScrXZ7Hs05t/OvYU89Vefi52lO3th22Vjq1S5WkGkmG3H0quJlBodAwTLuHJmkxPHttOjuFZ4ltNBrwcrLFz9WeALdaNPd1pqWvM0FejjL8JKmKaeT056XLyMjA2dmZ9PT0u/aTyzu3G92Pk9DdTCZX78RPDULYbWhNbMpNYpNvcCOvsNTP6bUaAj0daeHjRMu6zrTwdaaZt5MMPkm6i8r8fcqQK0NlDiIA6Vdg/Ti4chTQQI9/Q49XERotKTfyiE3JJjYlm3NXs4i4ks7JK+kk38i7Yzc6rYZADwda+DrTpp4L7f1dCfRwQKuV42clqYgMOSOodMgBFOTCtjfgj+XK80Z94LFlYO96x6ZCCBLSc4oDr+h/r2fdGXxOtnra13elnX9tOtR3Jbiuszzbk2o0GXJGcE8hV+T4OvjlJSi4CS71YMQa8Gl9148JIUjMyCEiTgm9Y5dS+Ss2jZv5JS93rXQaWvo6076+K+39a9O+viuutawrV6MkmTEZckZwXyEHkHgSvnsSUi+CzgYGfABtx1Z6N/mFBk4lZPBHTCrHLqXwR0wq1zJzS2yj0UBwXRcebOLBg0EeNPdxkpe3kkWTIWcE9x1yADfT4Kdn4OyvyvO24+Dh98HK9p7rEkJwOeUmf8SkcPRSKkdjUjiXlFViG3dHG3o1cefBIA+6BbrjYCMb0SXLIkPOCIwScgAGA+z/EHbNBQR4t4aR34Bz+XNgVcbVjBz2nEki7FQS+89fJ/uW1lwrnYZODerQK0g5y2vgVsto3ytJapEhZwRGC7ki0btgwyS4mQK13JX7dP6d73+/t8ktKOT3iynsOp3E7tNJxCRnl3g/wL0W/Zt70a+5F8F1neWsx5JZkiFnBEYPOYC0WFg7Gq5GgNYKHnkf2j9lnH2X4cK1LCXwziRx5EIKBYZ//nN7O9vyUDNP+jX3omMDV/Q6OcpPMg8y5IygSkIOIO8G/DwNIn9SnrefBP3ng77qW0czc/LZfeYa204msvtMUonL2tr2VvRuqgRe90A32UVFMmky5IygykIOQAjlPl3YO4AA/64wfBU4uBv3e8qRk1/IgfPXCT2ZyM5TV0nNzi9+z95aR88m7gxo6cODQR7YWcvAk0yLDDkjqNKQK3ImFH6YDHmZ4OynNEh4t6qa7ypHQaGBP2JS2RaZyPbIROLTc4rfs7fW0aepJwNb+fBAYzds9DLwJPXJkDOCagk5gGtnYO0oSIkGvR0MWQIthlXd992FEIKIK+lsjUhk84l44lJvFr/naKunX3MvHg32pmsjN6zkPTxJJTLkjKDaQg7gZqrS8hodpjzv9go8+CZo1T1rEkIQfjmNzScS2HwinqsZ/3RCrm1vxcMtvXk02JtODeqgk52PpWokQ84IqjXkAAyFsHM2HFykPA/sB8OWga1z1X93BRgMgj9iUth8IoGtEQklJhfwdrZlRHs/Rnb0w9vZTsUqpZpChpwRVHvIFTmxHja9AAU5UKcRjPwW3JtU3/dXQEGhgcMXUvjleDyhkYmk31QaLbQaeDDIg9Gd6tGjsYc8u5OqjAw5I1At5ACu/KmMe824AtYOMHQpNB1YvTVUUG5BIaEnE/n2SCxHLqYUv+7jbMsTHerxRAc/vJzvfRibJJWmMn+fqt85nj17NhqNpsQjKCioQp9dt24dGo2GIUOGlHh9woQJd+yzf//+VVB9FfFtC1P3gn83yMtSAi/sHeWS1sTY6HUMbu3Ld093ZucrPZjcrQEu9lbEp+fw0c6zdP3vLqasPsruM0kUGuS/p1L1M4mR282bN2fnzp3Fz/X6u5cVExPDjBkz6N69e6nv9+/fn6+++qr4uY2Nzf0XWp0c3GHcRtj+Fhz5DPYthITjyn06u9pqV1eqRh4OvPloM2b0a1J8dvd7TAo7oq6yI+oqvi52jO/iz6iO9XC0tVK7XKmGMImQ0+v1eHl5VXj7wsJCxowZw9tvv82+fftIS0u7YxsbG5tK7dMk6azg4fng0wZ+eRHO74Avein96Tybq11dmWytdAxp48uQNr6cu5rJt7/H8sOxOK6k3eS9radZHHaeUZ3q8VTX+rKhQqpyql+uApw7dw4fHx8CAgIYM2YMsbGx5W4/Z84cPDw8mDRpUpnb7NmzBw8PD5o0acKzzz5LcnJyufvMzc0lIyOjxMNktHoCJm0H53rK/HTL+8DJH9WuqkICPR0JGdic32f14b/DWtLIw4HM3AK++O0C3f+7m1e+Cycq3oSOtWRxVG94+PXXX8nKyqJJkyYkJCTw9ttvc+XKFU6ePImjo+Md2+/fv5+RI0cSHh6Om5sbEyZMIC0tjY0bNxZvs27dOuzt7WnQoAHR0dG88cYbODg4cOjQIXS60vuezZ49m7fffvuO11VpeCjLjWT4YSJc2KM87zodHvwP6EzihLxCDAbBnrNJfL73QomGiu6Bbkx9IIBujdzkzCjSXZl162paWhr+/v58+OGHd5ypZWZmEhwczKeffsrDDz8MUGrI3e7ChQs0bNiQnTt30rt371K3yc3NJTf3n86uGRkZ+Pn5mVbIARQWQNjb//SnC+gJj39V6joSpu745TSW7bvA1ogEitokmno7MfWBBjwa7CNHVEhlMuuQA+jQoQN9+vRh3rx5JV4PDw+nTZs2Jc7GDH+vXK/Vajlz5gwNGzYsdZ/u7u7MnTuXp59+ukI1qNqFpCJO/gA/Pw/52cpl7MivVRn3agyXU7JZsf8i649eLp4ZxdvZlsndAxjV0Q97a/M5U5Wqh1l1IbldVlYW0dHReHt73/FeUFAQERERhIeHFz8GDRpEr169CA8Px8/Pr9R9xsXFkZycXOo+zVaLYTB5J9SuD+mxsKIfnPhe7aruiZ+rPbMHNefgaw8ys18T3BxsSEjP4Z3NUXSdv4v/7TxHWvadq5hJUkWofiY3Y8YMBg4ciL+/P/Hx8YSEhBAeHk5UVBTu7u6MGzcOX1/fO87qitx+uZqVlcXbb7/NsGHD8PLyIjo6mn//+99kZmYSERFR4a4kJn8mV+RmqjKTyfm/u+D83zToO8es7tPdLie/kB//vMLnv0Vz6e+ZjWtZ6xjdqR6Tuwfg6SQ7F9d0ZnUmFxcXx6hRo2jSpAkjRoygTp06HD58GHd3ZW612NhYEhISKrw/nU7HiRMnGDRoEI0bN2bSpEm0a9eOffv2mV9fuYqwqw2j10P3fynPDy+BNUPgxnVVy7oftlZKoIW90oNFo9rQ1NuJG3mFLNt3ke7/3c1rP5zg4vUbapcpmQnVz+RMldmcyd0qahNsfFYZJeFUV7lP59NG7arumxCCPWeu8eme8/wRkwoo42QfbunNsz0a0sLXNCYxkKqP2Tc8mAKzDDmApNOwbrQyP53OBgb+D1qPUrsqo/kjJoXP9kSz63RS8Wu9gzz410NNaOZjRv+dpPsiQ84IzDbk4O/1Xp+Gs6HK845PQ793lREUFuJUQgaf7Ylm84n44u4nA4K9eaVvYxq6O6hbnFTlZMgZgVmHHCjrve6dD3v/qzz37wrDV4KDh6plGduFa1l8tPMcvxyPB5TL2Mfa1mV670D8XO1Vrk6qKjLkjMDsQ67I6S3w49PKOhJOvvDEGvBtp3ZVRncqIYMPtp9l56mrgLKo9qiO9Xi+VyM8ZGusxZEhZwQWE3IA187Cd2Pg+lnlPt2AD6DtWLWrqhJ/xabywfaz7D+vtC7b6LWM71KfZ3o0xLVW1S/7KFUPGXJGYFEhB5CTAT89A2e2KM+rcb1XNRyKTmbh9jMcu6S0xjrY6JnUrQFTHgjAwcZ8+xBKChlyRmBxIQfKfbp9C2H3e4AAv04wYjU4mvmUVGUo6nry/rYzRCUoM524Odjwct9Anmjvh16OjTVbMuSMwCJDrsjZ7cooidx0cPBSgq5eJ7WrqjIGg+DXk4m8v+00MX+PoGjk4cBr/YPo3dRDznpihmTIGYFFhxxAcrQyrXpSFGit4OH/QvuJYMF/8HkFBr49con/hZ0jNVtZfKdTA1dmDWhKcF0XdYuTKkWGnBFYfMgB5GbBz9MgaqPyvM2T8MgHYGXZrZEZOfl8ujuaLw9cJK9AmcVmcGsfZjzURHY7MRMy5IygRoQcgBDK3HQ7Z4MwgE9bpZuJc121K6tyV9Ju8sG2M/z41xUArHVanupan+d6NcLZznI6TlsiGXJGUGNCrkj0LtgwUZnVxN4NRqyC+t3UrqpanLySzntbT3EwWpki38Xeium9Axn7f/6yccJEyZAzghoXcgCpMcp9usQI0Oig33vQ6WmLvk9XpKgl9r2tpziXlAVAY08HZg9qTpeGbipXJ91OhpwR1MiQA8jLhs0vwYnvlOetRsOjH1n8fboiBYUG1v1xmQ+2nylunBjQ0ps3BjTF10WuLGYqZMgZQY0NOVDu0x3+DLa/CaJQma7pia9rxH26ImnZeXyw/SzfHLmEQYCtlZZpPRsx5YEAbK1KXwxJqj4y5IygRodckQt74Pun4GYK1HKH4augfle1q6pWUfEZzN4Uye8xyspifq52vDWgGX2becr+dSqSIWcEMuT+lnoJ1o2BqxGg1StDwTpMrhH36YoIIdh0PJ55W0+TmJEDwAON3QkZ2ExO66QSGXJGIEPuFnnZsOl5ZYUwqDH96W53I7eAJbvPs3zfRfIKDei1GiZ2a8D03oHUkuNhq5UMOSOQIXcbIeDgYtgZovSn822n3Kdz8lG7smp38foN3tkcVTw7sY+zLXMGt6BPM0+VK6s5ZMgZgQy5MkTvUu7T5aRBLQ9l3Kt/Z7WrUsWu01cJ2RTJ5ZSbAPRv7sXsQc3xcq5ZZ7hqMKvVuiQz0/BBmLoHPJrDjSRY9Sj8sVw506thHgzyZPtLPXi6RwA6rYbQyET6fLiX1YdiKDTUvONhquSZXBnkmdxd5N2Ajc/VuHGvZTmVkMHrP0YQfjkNgNZ+Lrw3tKVcXKeKyMtVI5AhVwFCwIGPIWzO3+Ne28CINeDip3Zlqig0CL49cokFoWfIzC1Ap9UwuVsDpvcJxN5aNkwYkww5I5AhVwklxr3WURbMafCA2lWp5mpGDrM3RfLryUQA6ta2Y+6QFvRsYlmLCKlJ3pOTqlfDB2HqXvAKhuxkWD0EDn5SI+/TAXg62fLZk+1YPq49Ps62xKXeZMJXf/DC2r9IzspVu7waR4acZBy1/WHSdggeqQwF2z4Lfpik3Lurofo082THKz2Y1K0BWg38cjyevh/9xi/H45EXUNVHXq6WQV6u3iMh4PcvYNsbYChQWmFHfg2uAWpXpqoTcWn8e8MJTidmAtCvuSfvDGmBh2PNbKi5X/JyVVKPRqNMzzT+F6UfXVIkfNETzu1QuzJVBdd1YdPz3XixdyB6rYZtkVfp++Fv/PhnnDyrq2Iy5KSq4d8Fnt4LdTtATjp8Mxz2LlBWDKuhrPVaXunbmE3Pd6O5jxPpN/N5Zf1xJq06SmJ6jtrlWSwZclLVcfKBCVuUBXIQsPtdZVLOnAy1K1NVMx8nNk7rysx+TbDWadl1Oom+H+7luz9i5VldFZD35Mog78kZ2Z9rYMsrUJgHbk1g1Fqo01DtqlR39momMzec4PjfnYi7B7ox77GW1K0tF9Qpj7wnJ5metmPhqV/B0Ruun4EvetX4+3QAjT0d+fHZLrzxSBA2ei37zl2n30e/sfZ3eVZnLDLkpOpTt73Sn86vk7Kw9TfDYd+HNbY/XRGdVsPUBxry6/TutPevzY28Ql7/MYIpq49xXfaru28y5KTq5eiptLy2HQ8ICHsbNjxVo/vTFQlwd2D9052Z9UhTrHVadp66Sr+PfmNn1FW1SzNrMuSk6qe3gUGLlAVytHqI/AlWPKSsFlbDabUapjwQwM/Pd6WJpyPJN/KYvPoor/94ghu5BWqXZ5ZkyEnqaT8Rxm9W+tNdPan0p7uwR+2qTEJTbyd+fr4rU7o3QKOBtb9fZsCiffwZm6p2aWZHhpykLv/Oyvx0Pm2UAf5rhsKhJTX+Ph2ArZWOWQOa8c3kTvg42xKTnM3jnx3kw+1nyC+suf0NK0uGnKQ+Z194KlRZ41UYlCFhPz0N+TfVrswkdGnoxq8vPcCQ1j4YBCzadZ7HPztI9LUstUszCzLkJNNgZQtDPoX+/wWNTlnc+sv+kB6ndmUmwdnOio9HtmHxqDY42eo5HpfOgEX7WHP4kuxqchcy5CTTodHA/z0D4zaCnSskhCv36S4dUrkw0zGwlQ/bXn6Aro3qkJNv4K2NJ3l6zTHSsvPULs1kyZCTTE+DB5T7dJ4t4MY1WDUQjn6pdlUmw9vZjjUTO/HmgKZY6TRsj7rKI//bx9G/F8CWSpIhJ5mmovnpmg8FQz5sfll5FMgzFlC6mkzuHsCPz3alfh174tNzeOKLw3yy65xcROc2MuQk02VdCx7/Cnr/B9AoZ3OrB0FWktqVmYyWdZ3Z/GJ3hrT2odAgWLj9LGNXHCEpQ85qUkSGnGTaNBro/i8Y/R3YOEHsIeU+3ZU/1a7MZDjY6Pnoida8/3gwdlY6DkYn8/D/9rHnjPzHACo4C8mmTZsqveO+fftiZ2d3T0WZAjkLiQm6fg7WjoLkc6C3hUGLIXiE2lWZlPNJWbyw9i9OJSjTWU19IIAZDzXBWm9Z5zNGX61Lq63cAdJoNJw7d46AAPOd8lqGnInKSYcfpsC5bcrzLi9An7dBq1O3LhOSk1/IvK2nWHXoEgCt/FxYPLIN9epYzvRNVTLVUmJiIgaDoUIPe3vLOZiSibF1Vuai6/4v5fnBxfDN48poCQlQRkq8PbgFn49th7OdFccvpzFg0T62nEhQuzRVVCjkxo8fX6lLzyeffFKe/UhVR6tTGiOGrwQre2Xd12UPQtJptSszKf2ae7H17+mbMnMLmPbtn4T8fJLcgkK1S6tWcmbgMsjLVTORGAFrR0N6LFg7wrBl0ORhtasyKQWFBj7ccZZP90QDEFzXmSWj2+Lnar5XXNU+M3BhYSEbN25k8ODBxtidJFWcV0uYuhvqd4e8TKVh4rf35QD/W+h1Wv7dP4ivJnTAxd6KE38PCdsemah2adXivkLuxIkTvPzyy/j4+DBq1Cg0Go2x6pKkiqvlBmN/gg5TAAG75sL3E+REnLfpFeTBlhe706aeCxk5BUxdc4y5m6MsfkaTu4bcjRs3KCz85xo+JSWFxYsX065dO9q1a8eiRYuYNWsW165dY+PGjZUuYPbs2Wg0mhKPoKCgCn123bp1aDQahgwZUuJ1IQT/+c9/8Pb2xs7Ojj59+nDu3LlK1yaZEZ0VDFgIA/8HWiuI2vj3RJyX1K7MpPi62PHd1M5M7tYAgOX7L/LE54eIT7PcGV/KDblly5YRGBhImzZt+Oijj3j88cfx8fFh9erVjBs3jri4OLRaLX369MHBweGei2jevDkJCQnFj/3799/1MzExMcyYMYPu3bvf8d6CBQtYtGgRS5cu5ciRI9SqVYt+/fqRkyN7gVu8dhNgwmao5f7PRJwX96ldlUmx1mt589FmLH2yHY62ev6MVVpfd1tq52FRDj8/P3H48GGRlJQkdDqdeOmll8Tp06dLbKPX60VkZGR5uylXSEiIaNWqVaU+U1BQILp06SKWL18uxo8fLwYPHlz8nsFgEF5eXuL9998vfi0tLU3Y2NiItWvXVvg70tPTBSDS09MrVZtkItIuC7H0ASFCnISYXVuII18IYTCoXZXJuXT9hhiw6Dfh/+pm4f/qZrEg9JTILyhUu6y7qszfZ7lnco0bN+b777/n66+/xsXFhWXLljFnzhxCQ0ONOofVuXPn8PHxISAggDFjxhAbG1vu9nPmzMHDw4NJkybd8d7FixdJTEykT58+xa85OzvTqVMnDh0qe8qe3NxcMjIySjwkM+ZcFyaGQsvhIAph6wz45UUokKtf3apeHXs2PNOFsf/nD8CS3dGMWX6Ea5mWc5zKDblvv/0WvV5Pbm4uV65c4cCBA3h4eDBu3Di8vb2ZPn06wH01OHTq1ImVK1cSGhrKZ599xsWLF+nevTuZmZmlbr9//35WrFjBsmXLSn0/MVFpMfL09CzxuqenZ/F7pZk3bx7Ozs7FDz8/v3v8iSSTYWUHjy2DvnMADfy5GlbJAf63s7XS8c6QFiwe1YZa1jqOXExh4OL9xQtem717OVXMz88XP/zwgxg4cKCwsrISgYGB4vXXXxfHjh27l92VkJqaKpycnMTy5cvveC8jI0PUr19fbN26tfi12y9XDxw4IAARHx9f4rPDhw8XI0aMKPN7c3JyRHp6evHj8uXL8nLVkpzdLsR7fsrl6wfNhLjyl9oVmaTzSZniwYW7hf+rm0XgrK1i/R+xapdUKqNdrpZFr9fz2GOPsWnTJi5fvsyUKVP4+eef6dChw32HrouLC40bN+b8+fN3vBcdHU1MTAwDBw5Er9ej1+tZvXo1mzZtQq/XEx0djZeXFwBXr5Zcq/Lq1avF75XGxsYGJyenEg/JggT2hSlhUKcRZMQpU6uf/FHtqkxOQ3cHNk7rSp+mnuQVGJi54QSzN0WadTeT++4M7OnpycyZM4mMjCz3nldFZWVlER0djbe39x3vBQUFERERQXh4ePFj0KBB9OrVi/DwcPz8/GjQoAFeXl6EhYUVfy4jI4MjR47QuXPn+65PMmNugTA5DBr2hoKbyqLWYe+AwXz/gKuCo60VX4xtx0t9AgFYeTCGMcuPcD3LPO/TVSjkTpw4gaECvwgdO3YEIDIykoKCii2EO2PGDPbu3UtMTAwHDx5k6NCh6HQ6Ro0aBcC4ceN4/fXXAbC1taVFixYlHi4uLjg6OtKiRQusra3RaDS89NJLzJ07l02bNhEREcG4cePw8fG5oz+dVAPZucCY76Hz88rzfQvhuycht/R7wDWVVqvhpT6N+WJsOxxs9Px+MYVBi/cTEZeudmmVVqGQa9OmDcnJyRXeaefOne/aQlokLi6OUaNG0aRJE0aMGEGdOnU4fPgw7u7uAMTGxpKQULnZE/7973/zwgsvMHXqVDp06EBWVhahoaHY2tpWaj+ShdLqoN+7MGQp6GzgzBZY3hdSLqhdmcl5qLkXG6d1JcCtFvHpOQxbepAfjpnXCmoVnk9u6tSpFZ5C6dNPPyUqKkrOJyeZvrijsG4MZCWCXW0YvgoCeqhdlcnJyMnn5XXhhJ1WWqaf6lqfNx5pipVOnck4jT5pZs+ePSvdTeTbb78t9b6auZAhV4NkJMC60RD/p7Lma//50HGKMvW6VMxgEHy88yyLdimNgp0D6vDJ6DbUcbCp9lqMHnI1kQy5Gib/JvwyXVnUGqDteHhkIeit1a3LBIWeTORf68O5kVdI3dp2rBjfgSZejtVaQ7VPtSRJZs/KDoZ+fkvH4VWwejDcuK52ZSanfwvlPl39OvbEpd5k2GcH2X3adDtYy5CTpCIaDXSdDqPX/70y2EH4ohcknlS7MpMT6OnIT8915f8CXMnKLWDSqj9Yvu+CUYd7GosMOUm6XeOHYPJOcA1QZhxe8RCc+kXtqkxO7VrWrJ7YiVEd/TAImLvlFG/8FEFegWn1O5QhJ0mlcW+idBwO6An5N5S+dHvljMO3s9ZreW9oS956tBlaDaz9/TLjvjxC6o08tUsrJkNOkspi7wpjfoCOTyvPd8+FDRMhL1vdukyMRqNhUrcGLB/fHgcbPYcvpDD00wOcT8pSuzTgPlpXo6KiiI2NJS+vZGIPGjTIKIWpTbauSiUcWwlb/gWGAvBuDSO/BWdftasyOWcSM5m06g/iUm/iaKvn0zFt6R7obvTvqdIuJBcuXGDo0KFERESg0WiKbzQW9aO7dap0cyZDTrpDzAFYPxayk6GWhxJ0fvc/KYWluZ6VyzNrjnH0Uio6rYaQgc0Y17m+Ub+jSruQTJ8+nQYNGpCUlIS9vT2RkZH89ttvtG/fnj179txrzZJk+up3hSm7waM53EiClY9A+Fq1qzI5bg42fDOlE4+19aXQIPjPz5H85+eTFKg0k0mlQ+7QoUPMmTMHNzc3tFotWq2Wbt26MW/ePF588cWqqFGSTEdtf5i0HYIehcI82PgMbH8LDJZxBWMsNnodHwxvxav9g9BoYPWhS0xadZSs3IpN3GFMlQ65wsJCHB2V3s1ubm7Ex8cD4O/vz5kzZ4xbnSSZIhsHGLEGHpipPD+4SBkWJmcyKUGj0fBsz4Z8NqYdtlZa9p69xoilh0hMr94FpSodci1atOD48eOAMnX5ggULOHDgAHPmzDHrAfmSVClaLTz4Jgxbocxkcjb07yUQY9SuzOT0b+HFuqmdcXOwJiohg6GfHuBUQvWtoVLpkHvzzTeL55abM2dO8ZoMW7duZdGiRUYvUJJMWsvH4alfwcETkqJg2YNw6aDaVZmc1n4u/PRcVxq61yIhPYfhSw/x29lr1fLdRhmgn5KSQu3ate9rQRtTI1tXpUpJvwLrRkHCcWVx60c/grZj1a7K5KRn5zN1zVGOXExBp9Xw3tAWPNGhXqX3U+0D9F1dXS0q4CSp0px94alQaDYEDPmw6XnYNks2SNzG2d6K1ZM6MrSN0vL66g8RLNx2pkrHvMoRD5JkLNb28PhX0OM15fmhT2DtSMiRa/jeykav48MRrXjxwUYAfLL7PC99F05uQdX8gyBDTpKMSauFXq8rYae3hXPbYUVfSLmodmUmRaPR8MpDTVjweDB6rYafw+MZu/x30rKNP+ZVhpwkVYUWjykNEo7ecO200iARs1/tqkzOiPZ+rHyqI442en6PSeGxzw4Sm2zcscEy5CSpqvi2VUZI+LSFmynKJJx/rla7KpPTLdCNDc92wcfZlgvXbjD00wNcvH7DaPuXISdJVcnJG57aCs0fUwb3b3pBNkiUoomXIz9N60oLXyda+DrjV9vOaPuWazyUQXYhkYxKCNgzH/bOV5437g/DloNN9a6NYOpu5BYgAAcbfbnbyTUeJMnUaDRKg8SwFUqDxNlQWNEP0iq2PnFNUctGf9eAqywZcpJUnVo+DhO2KFM1JUUqDRKXf1e7KosmQ06Sqlvd9jB1N3i2hBvXYOWjcGK92lVZLBlykqQG57owMRSaDIDCXPhxCuyaCwbTWgTGEsiQkyS12DjAE18ryyAC/PY+bJgg15AwMhlykqQmrVZZ0Hrwp8rA/qif4auHISNB7coshgw5STIFbcbA+E1g5woJ4UqDRMJxtauyCDLkJMlU+HeBKbvArTFkxsOXD8PprWpXZfZkyEmSKXFtAJN2/LOo9brRcPATuaj1fZAhJ0mmxs4FxmyAdhMAAdtnweaXoDBf3brMlAw5STJFOit49GN46F1Aoyxu/c3jcDNN3brMkAw5STJVGg10eV5ZxNqqFlzY8/fcdBfUrsysyAH6ZajoAODCwkLy8+VlhARWVlbodLqq2XnCCfj2CaVBws5VCT7/zlXzXWagMgP0ZciV4W4HUQhBYmIiaWlp1V+cZLJcXFzw8vKqmjVPMhKU6dQTwkFnDYM+gVZPGP97zEBlQs64w/1rkKKA8/DwwN7eXi7kU8MJIcjOziYpKQkAb29v439J0dx0Pz0Np36Bn6ZC8jnoNUu5tJVKJUPuHhQWFhYHXJ06ddQuRzIRdnbKRI9JSUl4eHhUzaWrdS0YvhrC3oYDHytDwVIuwuAlYGVr/O+zALLh4R4U3YOzt7dXuRLJ1BT9TlTpfVqtFvq+rVyuavVwcoMytfqN61X3nWZMhtx9kJeo0u2q9Xei7Vh48gewcYbLh2F5H7h+rvq+30zIkJMkcxbQEybvABd/SL2oBN3FfWpXZVJkyEmSuXNvApPDoG4HyEmDNUMhfK3aVZkMGXKSZAkc3GH8L9B8KBjyYeMzsOtdOeYVGXI1UmJiIi+88AIBAQHY2Njg5+fHwIEDCQsLU7u0YteuXePZZ5+lXr162NjY4OXlRb9+/Thw4IDRvuPy5ctMnDgRHx8frK2t8ff3Z/r06SQnJxvtO6qVlR0M+xK6vaI8/20B/DAZ8nPUrUtlsgtJDRMTE0PXrl1xcXHh/fffp2XLluTn57Nt2zamTZvG6dOn1S4RgGHDhpGXl8eqVasICAjg6tWrhIWFGS2ALly4QOfOnWncuDFr166lQYMGREZGMnPmTH799VcOHz6Mq6urUb6rWmm10CcEXAOUQf0nN0B6nDJColYN7e4kpFKlp6cLQKSnp9/x3s2bN0VUVJS4efOmCpXdn4cfflj4+vqKrKysO95LTU0VQghRWFgo3nvvPVG/fn1ha2srgoODxffff19tNaampgpA7Nmzp9ztevToIV544QUxc+ZMUbt2beHp6SlCQkIq9B39+/cXdevWFdnZ2SVeT0hIEPb29uKZZ565p9pN6ncjercQ7/kJEeIkxMethLh2Vu2KjKa8v8/byctVIxBCkJ1XoMpDVOKeS0pKCqGhoUybNo1atWrd8b6LiwsA8+bNY/Xq1SxdupTIyEhefvllnnzySfbu3WusQ1YuBwcHHBwc2LhxI7m5ueVuu2rVKmrVqsWRI0dYsGABc+bMYceOHeV+JiUlhW3btvHcc88Vd+At4uXlxZgxY/juu+8qdWxNUnHLa71/Wl5jjHe5by7k5aoR3MwvpNl/tqny3VFz+mFvXbH/jOfPn0cIQVBQUJnb5Obm8t5777Fz5046d1YGgAcEBLB//34+//xzevToAcDQoUPZs2cPvXv3ZsOGDff/g9xCr9ezcuVKpkyZwtKlS2nbti09evRg5MiRBAcHl9g2ODiYkJAQAAIDA/nkk08ICwujb9++Ze7/3LlzCCFo2rRpqe83bdqU1NRUrl27Rm5uLmPHjiUpKQm9Xs9bb73F8OHDjffDVjX3JjB5lzLm9cpRWDNEWU8i2Ix+hvskz+RqkIqcmZw/f57s7Gz69u1bfEbl4ODA6tWriY6OLt5u+vTprF69uspqHTZsGPHx8WzatIn+/fuzZ88e2rZty8qVK0tsd3voeXt7F48fvZuKHA+9Xs/HH39MVFQU27dv56WXXuLGjRsV/jlMgoM7TNgMTQdCYR78OFkZDmbuZ6oVJM/kjMDOSkfUnH6qfXdFBQYGotFoym1cyMrKAmDLli34+vqWeM/Gxqb4//fs2ZM9e/ZUrthSdO3alQ8//JBOnToxadIkWrRowcsvvwyAra0tffv2pW/fvrz11ltMnjyZkJAQJkyYUPx5KyurEvvTaDQYblm7tLT9jxs3Do1Gw6lTpxg6dOgdNZ06dYratWvj7u6ORqMpHmzv5eWFm5sbKSkppV7umzQrO2XM64634NAnyhqvqTHKxJw6q7t92qypfiY3e/ZsNBpNiUd5l1M//vgj7du3x8XFhVq1atG6dWvWrFlTYpsJEybcsc/+/ftX2c+g0Wiwt9ar8qjMMCJXV1f69evHkiVLSj0bSUtLo1mzZtjY2BAbG0ujRo1KPPz8/Ix52AB46623mD9/Ph9++CFarbY44ErTrFmzSp9Flbb/OnXq0LdvXz799FNu3rxZYvvExES++eYbnnjiiTuO7bFjxygsLKyS41AttFro9y48shA0Wvjra/hmOOSkq11ZlTKJM7nmzZuzc+fO4ud6fdllubq6MmvWLIKCgrC2tmbz5s089dRTeHh40K/fP2dT/fv356uvvip+futZSE22ZMkSunbtSseOHZkzZw7BwcEUFBSwY8cOPvvsM06dOsWMGTN4+eWXMRgMdOvWjfT0dA4cOICTkxPjx4+v0Pe0bt2agoKCO17fvn07Pj4+xc/79+/PrFmz2LJlC6GhoQAkJyczfPhwJk6cSHBwMI6Ojhw9epQFCxYwePDgSv28pe0f4JNPPqFLly7069ePuXPnluhC4uvry7vvvltiPykpKYwbN45ly5ZV6vtNUscp4OwHGybChd3wZX8YvR5czDS878IkQk6v1+Pl5VWhbXv27Fni+fTp01m1ahX79+8vEXJFHUilkgICAvjzzz959913+de//kVCQgLu7u60a9eOzz77DIB33nkHd3d35s2bx4ULF3BxcaFt27a88cYbFf6e8PDwCm33xx9/kJKSgr+/f/Glp4ODA506deKjjz4iOjqa/Px8/Pz8mDJlSqVqKGv/oFy6Hz16lJCQEEaMGEFKSgpeXl4MGTKEkJCQEn3kcnNzGTJkCK+99hpdunSp1PebrCb9lbnpvn0CkqKUltfR34FPa7UrM76q7MtSESEhIcLe3l54e3uLBg0aiNGjR4tLly5V6LMGg0Hs3LlT2Nvbi+3btxe/Pn78eOHs7Czc3d1F48aNxTPPPCOuX79e7r5ycnJEenp68ePy5csW2U/OmHbv3i2GDRt2z5+Pi4sTwcHB4sKFC6JNmzYiIiLCiNUZZ/8Gg0GMHDmywv3vzO53IzVWiCX/p/Slm+stxOlf1a6oQirTT071kNu6datYv369OH78uAgNDRWdO3cW9erVExkZGWV+Ji0tTdSqVUvo9XphY2MjVqxYUeL9tWvXip9//lmcOHFC/PTTT6Jp06aiQ4cOoqCgoMx9hoSECOCOhwy50vXu3Vu4ubkJOzs74evrKw4ePFipz2dnZ4vOnTuLvXv3CiGEWL9+vRgxYoTR6jPW/vft2yc0Go1o1apV8ePEiRNlbm+Wvxs304RYNUgJutkuQhz5Qu2K7sqsQu52qampwsnJSSxfvrzMbQoLC8W5c+fEX3/9JRYuXCicnZ3F7t27y9w+OjpaAGLnzp1lbiPP5CRjMNvfjYI8ITY+pwRdiJMQ22YJUViodlVlqkzImcQ9uVu5uLjQuHFjzp8/X+Y2Wq2WRo0aAcoN7lOnTjFv3rw77tcVCQgIwM3NjfPnz9O7d+9St7GxsZGNE1LNpbNSZhquXV/pXnJwMaRfgSGfmf206qp3IbldVlYW0dHRlVoIxGAwlDv8Jy4ujuTk5KpZXESSLIVGAw/MhKGfK9OqR/4IXz8G2SlqV3ZfVA+5GTNmsHfvXmJiYjh48CBDhw5Fp9MxatQoAMaNG8frr79evP28efPYsWMHFy5c4NSpU3zwwQesWbOGJ598ElBCcubMmRw+fJiYmBjCwsIYPHgwjRo1KtH6KklSGVqN/HtadSe4dAC+7Aepl9Su6p6pfrkaFxfHqFGjSE5Oxt3dnW7dunH48GHc3d0BiI2NRav9J4tv3LjBc889R1xcHHZ2dgQFBfH111/zxBPK+pM6nY4TJ06watUq0tLS8PHx4aGHHuKdd96Rl6OSVFEBPWFiqNJZ+PpZWNH37y4mbdSurNLk4tJlKG/x2pycHC5evEiDBg2wtTXv+xWScVnc70ZGvBJ0V0+CVS0YvhIaP6R2VZVaXFr1y1VJkkyYkw889SsE9IL8G8psJsdWql1VpciQkySpfLZOMOZ7aDUaRCH8Mh3C3jGbWUxkyEmSdHc6KxjyKfR4VXm+byH89AwU5KlbVwXIkJMkqWI0Guj1htKfTqODE+vgm8chJ0PtysolQ06SpMppOxbGrAdrB7i4F756BDIS1K6qTDLkaiBTX5Jwz549d8wHeOujV69e97V/i1uKUA2N+sCELVDLA65GKF1Mrp1Ru6pSyZCrYWJiYmjXrh27du3i/fffJyIigtDQUHr16sW0adPULg+ALl26kJCQcMfj888/R6PR8Nxzz5X52by88u8RXbhwgfbt23Pu3DnWrl3L+fPnWbp0KWFhYXTu3JmUFPPu3V+tfForC+XUaQTpl2HFQxB7WO2q7lTlI2nNlFySUL0lCUsTFRUlHB0dxaxZs0q83qNHDzFt2jQxffp0UadOHdGzZ89y91NVSxEWMeffjXuWdV2IZb2Vgf3veAgRtanKv9KsZyExFZUKOYNBiNwsdR4GQ4V/puTkZKHRaMR7771X7nZz584VQUFBIjQ0VERHR4uvvvpK2NjY3HUd1KqSmpoqAgMDxcCBA4Xhtp+3R48ewsHBQcycOVOcPn1anD59usz93O3nnzJliqhdu/Yd31EZNTLkhBAi94YQ3478exYT5yqfrsmsZyExS/nZ8J7P3berCm/Eg3XFFlUxlyUJb2UwGBg9ejR6vZ5vvvmm1DUtAgMDWbBgwV33VaOWIqxu1vYwYg1snQHHvlL+NyMeev9HaZVVkQy5GkRUcknCW+Xl5dGmzT/jFqdPn87EiRNZtWqV0eu81RtvvMGhQ4f4/fffcXR0LHWbdu3aVWqfFTkORUsRtm7dmsTERNq1a8cjjzxifqt0VSedHh79CJx8Yfdc2P8hZCbAwEWgt1atLBlyxmBlr5xRqfXdFWRuSxKuW7eOhQsXsmXLFgIDA8vcx+3BU9Y+GzVqVPOWIqxuGg30mAlO3rDpRTi+FrKuwojVYFP6P1JVTYacMWg0Fb5kVNOtSxK++OKLd/zB3r4kYdGlaVUqWjKwe/fuJZYkDA8PZ9KkScyfP7/SU2SVtc9blyJ8+eWXsbOzK/5M0VKERWuy3srslyJUQ5snwcET1o+H6F1KX7oxG8DRs/prqdK7g2bMUltXo6OjhZeXl2jWrJnYsGGDOHv2rIiKihL/+9//RFBQkBBCiFmzZok6deqIlStXivPnz4tjx46JRYsWiZUrV5bYV3kL2bRq1Uo0b978jseVK1fu2LZt27biwQcfFHl5eUIIIa5duyb8/f3FI488IhISEu54JCUlFX+2R48eYvr06XfdZ5GzZ88KNzc30b17d7F3714RGxsrfv31V9GiRQsRGBgokpOTS2yfnJwsmjVrJg4cOHD3gyvM+3ejSsQdFeK/AUqDxEcthbh+3ii7lQ0PUpnMYUnCLVu2cOnSJS5dulTqbM7+/v7ExMRUap9FavxShNXNt53Sl27NUEiNUfrSPbmheuelM0qsWiBLPZMzJlNcktBY+6zsUoRF5O9GGTKvCrG0u3JG966PEOfD7mt3sp+cEciQK58pLklozH1WdinCIvJ3oxw5GUKsHKgE3dt1hDhx7x3MKxNycmbgMsiZgaV7IX837qIgV5miKfJH5Xm/edC57GF6ZZEzA0uSZJr0NjBsBXR6Rnm+7XXYEVKlE3DKkJMkqXpptdB/PvQOUZ4f+Bg2PgeF+VXzdVWyV0mSpPJoNND9FRi8RJmA8/i3sG405N0w+lfJkJMkST1tnoSR34LeDs5th1WDjL6YtQy5+yDbbKTbyd+Je9CkP4zfBLYucOWosph1+hWj7V6G3D0o6mCanZ2tciWSqSn6nbi9E7J0F34dYeI2ZXC/3tao41zliId7oNPpcHFxISkpCQB7e/tSpwCSag4hBNnZ2SQlJeHi4oJOp1O7JPPjEQSTdoBWryyDaCQy5O6Rl5cXQHHQSRKAi4tL8e+GdA+cfe++TSXJkLtHRVPxeHh4kJ9fNU3fknmxsrKSZ3AmSIbcfdLpdPIXW5JMmGx4kCTJosmQkyTJosmQkyTJosl7cmUo6tSZkZGhciWSJN2u6O+yIp2vZciVITMzE0DO6y9JJiwzMxNnZ+dyt5HzyZXBYDAQHx+Po6PjXTv6ZmRk4Ofnx+XLl+86t5V0d/J4Gp+lHVMhBJmZmfj4+KDVln/XTZ7JlUGr1VK3bt1KfcbJyckifoFMhTyexmdJx/RuZ3BFZMODJEkWTYacJEkWTYacEdjY2BASElJihXnp3snjaXw1+ZjKhgdJkiyaPJOTJMmiyZCTJMmiyZCTJMmiyZCTJMmiyZCroCVLllC/fn1sbW3p1KkTv//+e7nbf//99wQFBWFra0vLli3ZunVrNVVqHipzPJctW0b37t2pXbs2tWvXpk+fPnc9/jVRZX9Hi6xbtw6NRsOQIUOqtkC1COmu1q1bJ6ytrcWXX34pIiMjxZQpU4SLi4u4evVqqdsfOHBA6HQ6sWDBAhEVFSXefPNNYWVlJSIiIqq5ctNU2eM5evRosWTJEvHXX3+JU6dOiQkTJghnZ2cRFxdXzZWbrsoe0yIXL14Uvr6+onv37mLw4MHVU2w1kyFXAR07dhTTpk0rfl5YWCh8fHzEvHnzSt1+xIgRYsCAASVe69Spk3j66aertE5zUdnjebuCggLh6OgoVq1aVVUlmp17OaYFBQWiS5cuYvny5WL8+PEWG3LycvUu8vLyOHbsGH369Cl+TavV0qdPHw4dOlTqZw4dOlRie4B+/fqVuX1Nci/H83bZ2dnk5+fj6upaVWWalXs9pnPmzMHDw4NJkyZVR5mqkQP07+L69esUFhbi6elZ4nVPT09Onz5d6mcSExNL3T4xMbHK6jQX93I8b/fqq6/i4+Nzxz8kNdW9HNP9+/ezYsUKwsPDq6FCdcmQk8zK/PnzWbduHXv27MHW1lbtcsxSZmYmY8eOZdmyZbi5ualdTpWTIXcXbm5u6HQ6rl69WuL1q1evlrm+ppeXV6W2r0nu5XgWWbhwIfPnz2fnzp0EBwdXZZlmpbLHNDo6mpiYGAYOHFj8msFgAECv13PmzBkaNmxYtUVXI3lP7i6sra1p164dYWFhxa8ZDAbCwsLo3LlzqZ/p3Llzie0BduzYUeb2Ncm9HE+ABQsW8M477xAaGkr79u2ro1SzUdljGhQUREREBOHh4cWPQYMG0atXL8LDwy1vNmy1Wz7Mwbp164SNjY1YuXKliIqKElOnThUuLi4iMTFRCCHE2LFjxWuvvVa8/YEDB4RerxcLFy4Up06dEiEhIbILyS0qezznz58vrK2txYYNG0RCQkLxIzMzU60fweRU9pjezpJbV2XIVdDixYtFvXr1hLW1tejYsaM4fPhw8Xs9evQQ48ePL7H9+vXrRePGjYW1tbVo3ry52LJlSzVXbNoqczz9/f0FcMcjJCSk+gs3YZX9Hb2VJYecnGpJkiSLJu/JSZJk0WTISZJk0WTISZJk0WTISZJk0WTISZJk0WTISZJk0WTISZJk0WTISZJk0WTISZJk0WTISZJk0WTISZJk0WTISRZt7dq12NnZkZCQUPzaU089RXBwMOnp6SpWJlUXOUBfsmhCCFq3bs0DDzzA4sWLCQkJ4csvv+Tw4cP4+vqqXZ5UDeTMwJJF02g0vPvuuzz++ON4eXmxePFi9u3bJwOuBpFnclKN0LZtWyIjI9m+fTs9evRQuxypGsl7cpLFCw0N5fTp06WuaCVZPnkmJ1m0P//8k549e/L555+zcuVKnJyc+P7779UuS6pG8p6cZLFiYmIYMGAAb7zxBqNGjSIgIIDOnTvz559/0rZtW7XLk6qJPJOTLFJKSgpdunShZ8+eLF26tPj1AQMGUFhYSGhoqIrVSdVJhpwkSRZNNjxIkmTRZMhJkmTRZMhJkmTRZMhJkmTRZMhJkmTRZMhJkmTRZMhJkmTRZMhJkmTRZMhJkmTRZMhJkmTRZMhJkmTR/h8b/i8NGq1c8AAAAABJRU5ErkJggg==",
+ "text/plain": [
+ "