Implement the Alchemical Model from torch_alchemical repo (#66)

---------

Co-authored-by: frostedoyster <[email protected]>

.github/workflows/alchemical-model-tests.yml

@@ -0,0 +1,36 @@
+name: Alchemical Model tests
+
+on:
+  push:
+    branches: [main]
+  pull_request:
+    # Check all PR
+
+jobs:
+  tests:
+    runs-on: ${{ matrix.os }}
+    strategy:
+      matrix:
+        include:
+          - os: ubuntu-22.04
+            python-version: "3.11"
+
+    steps:
+    - uses: actions/checkout@v3
+
+    - name: Set up Python ${{ matrix.python-version }}
+      uses: actions/setup-python@v4
+      with:
+        python-version: ${{ matrix.python-version }}
+    - run: pip install tox
+
+    - name: run Alchemical Model tests
+      run: tox -e alchemical-model-tests
+      env:
+        # Use the CPU only version of torch when building/running the code
+        PIP_EXTRA_INDEX_URL: https://download.pytorch.org/whl/cpu
+
+    - name: Upload codecoverage
+      uses: codecov/codecov-action@v3
+      with:
+        files: ./tests/coverage.xml

README.rst

@@ -65,6 +65,10 @@ atomistic model.
    * - SOAP BPNN
      - A Behler-Parrinello neural network with SOAP features
 
+   * - Alchemical Model
+     - A Behler-Parrinello neural network with SOAP features
+       and Alchemical Compression of the composition space
+
 .. marker-documentation
 
 Documentation

docs/src/architectures/alchemical-model.rst

@@ -0,0 +1,37 @@
+.. _architecture-alchemical-model:
+
+Alchemical Model
+================
+
+This is an implementation of Alchemical Model: a Behler-Parrinello neural network
+:footcite:p:`behler_generalized_2007` with Smooth overlab of atomic positions (SOAP)
+features :footcite:p:`bartok_representing_2013` and Alchemical Compression of the
+composition space :footcite:p:`willatt_feature_2018, lopanitsyna_modeling_2023,
+mazitov_surface_2024`. This model is extremely useful for simulating systems with
+large amount of chemical elements.
+
+
+Installation
+------------
+
+To install the package, you can run the following command in the root
+directory of the repository:
+
+.. code-block:: bash
+
+    pip install .[alchemical-model]
+
+This will install the package with the Alchemical Model dependencies.
+
+
+Hyperparameters
+---------------
+
+The hyperparameters (and relative default values) for the Alchemical Model model are:
+
+.. literalinclude:: ../../../src/metatensor/models/cli/conf/architecture/experimental.alchemical_model.yaml
+   :language: yaml
+
+Any of these hyperparameters can be overridden with the training parameter file.
+
+

docs/src/dev-docs/utils/index.rst

@@ -12,4 +12,5 @@ This is the API for the ``utils`` module of ``metatensor-models``.
    model-io
    omegaconf
    combine_dataloaders
-   neighbor_lists
+   neighbors_lists
+   normalize

docs/src/dev-docs/utils/neighbor_lists.rst → docs/src/dev-docs/utils/neighbors_lists.rst

@@ -3,7 +3,7 @@ Neighbor lists
 
 Utilities to attach neighbor lists to a ``metatensor.torch.atomistic.System`` object.
 
-.. automodule:: metatensor.models.utils.neighbor_list
+.. automodule:: metatensor.models.utils.neighbors_lists
     :members:
     :undoc-members:
     :show-inheritance:

docs/src/dev-docs/utils/normalize.rst

@@ -0,0 +1,9 @@
+Normalization
+=============
+
+Utilities to normalize the target values used for training.
+
+.. automodule:: metatensor.models.utils.normalize
+    :members:
+    :undoc-members:
+    :show-inheritance:

docs/static/refs.bib

@@ -41,3 +41,28 @@ @article{willatt_feature_2018
   urldate = {2024-02-19},
   langid = {english}
 }
+
+@article{mazitov_surface_2024,
+	author={Mazitov, Arslan and Springer, Maximilian A. and Lopanitsyna, Nataliya and Fraux, Guillaume and De, Sandip and Ceriotti, Michele},
+	title={Surface segregation in high-entropy alloys from alchemical machine learning},
+	journal={Journal of Physics: Materials},
+	url={http://iopscience.iop.org/article/10.1088/2515-7639/ad2983},
+	year={2024},
+	abstract={High-entropy alloys (HEAs), containing several metallic elements in near-equimolar proportions, have long been of interest for their unique mechanical properties. 
More recently, they have emerged as a promising platform for the development of novel heterogeneous catalysts, because of the large design space, and the synergistic effects between their components.
In this work we use a machine-learning potential that can model simultaneously up to 25 transition metals to study the tendency of different elements to segregate at the surface of a HEA.
We use as a starting point a potential that was previously developed using exclusively crystalline bulk phases, and show that, thanks to the physically-inspired functional form of the model, adding a much smaller number of defective configurations makes it capable of describing surface phenomena.
We then present several computational studies of surface segregation, including both a simulation of a 25-element alloy, that provides a rough estimate of the relative surface propensity of the various elements, and targeted studies of CoCrFeMnNi and IrFeCoNiCu, which provide further validation of the model, and insights to guide the modeling and design of alloys for heterogeneous catalysis.}
+}
+
+@article{lopanitsyna_modeling_2023,
+  title = {Modeling high-entropy transition metal alloys with alchemical compression},
+  author = {Lopanitsyna, Nataliya and Fraux, Guillaume and Springer, Maximilian A. and De, Sandip and Ceriotti, Michele},
+  journal = {Phys. Rev. Mater.},
+  volume = {7},
+  issue = {4},
+  pages = {045802},
+  numpages = {15},
+  year = {2023},
+  month = {Apr},
+  publisher = {American Physical Society},
+  doi = {10.1103/PhysRevMaterials.7.045802},
+  url = {https://link.aps.org/doi/10.1103/PhysRevMaterials.7.045802}
+}
+

examples/alchemical_model/README.rst

@@ -0,0 +1,2 @@
+Basic usage of the Alchemical Model CLI
+========================================

examples/alchemical_model/alchemical_reduced_10.xyz

@@ -0,0 +1 @@
+../../tests/resources/alchemical_reduced_10.xyz

examples/alchemical_model/eval.yaml

@@ -0,0 +1,6 @@
+structures: "alchemical_reduced_10.xyz" # file where the positions are stored
+targets:
+  energy:
+    key: "energy" # name of the target value
+    forces: false
+    stress: false

examples/alchemical_model/options.yaml

@@ -0,0 +1,18 @@
+# architecture used to train the model
+architecture:
+  name: experimental.alchemical_model
+  training:
+    num_epochs: 10
+
+# Mandatory section defining the parameters for structure and target data of the
+# training set
+training_set:
+  structures: "alchemical_reduced_10.xyz" # file where the positions are stored
+  targets:
+    energy:
+      key: "energy" # name of the target value
+      forces: false
+      stress: false
+
+test_set: 0.1  # 10 % of the training_set are randomly split and taken for test set
+validation_set: 0.1 # 10 % of the training_set are randomly split and for validation

pyproject.toml

@@ -57,6 +57,9 @@ build-backend = "setuptools.build_meta"
 
 [project.optional-dependencies]
 soap-bpnn = []
+alchemical-model = [
+  "torch_alchemical @ git+https://github.com/abmazitov/torch_alchemical.git@fafb0bd",
+]
 
 [tool.setuptools.packages.find]
 where = ["src"]

src/metatensor/models/cli/conf/architecture/experimental.alchemical_model.yaml

@@ -0,0 +1,23 @@
+# default hyperparameters for the Alchemical Model
+name: alchemical_model
+
+model:
+  soap:
+    num_pseudo_species: 4
+    cutoff_radius: 5.0
+    basis_cutoff: 400  # controls how large the radial-angular basis is
+    radial_basis_type: 'physical'  # 'physical' or 'le'
+    basis_scale: 3.0  # controls the initial scale of the physical basis (in Angstroms, does not affect the le basis)
+    trainable_basis: true  # whether the radial basis is trainable (i.e. contains a small NN)
+
+  bpnn:
+    num_hidden_layers: 2
+    num_neurons_per_layer: 32
+    activation_function: SiLU  # only SiLU is supported for the moment
+
+training:
+  batch_size: 8
+  num_epochs: 100
+  learning_rate: 0.001
+  log_interval: 10
+  checkpoint_interval: 25

src/metatensor/models/cli/eval_model.py

@@ -14,7 +14,7 @@
 from ..utils.info import finalize_aggregated_info, update_aggregated_info
 from ..utils.loss import TensorMapDictLoss
 from ..utils.model_io import load_exported_model
-from ..utils.neighbor_list import attach_neighbor_lists
+from ..utils.neighbors_lists import get_system_with_neighbors_lists
 from ..utils.omegaconf import expand_dataset_config
 from .formatter import CustomHelpFormatter
 
@@ -74,7 +74,7 @@ def _eval_targets(model, dataset: Union[_BaseDataset, torch.utils.data.Subset])
         dataset, batch_size=1, collate_fn=collate_fn
     )
     for (structure,), _ in dataloader:
-        attach_neighbor_lists(structure, requested_neighbor_lists)
+        get_system_with_neighbors_lists(structure, requested_neighbor_lists)
 
     # Extract all the possible outputs and their gradients from the dataset:
     outputs_dict = get_outputs_dict([dataset])
@@ -165,17 +165,15 @@ def eval_model(
 
     logging.basicConfig(level=logging.INFO, format="%(message)s")
     logger.info("Setting up evaluation set.")
-    dtype = next(model.parameters()).dtype
 
     options = expand_dataset_config(options)
     eval_structures = read_structures(
         filename=options["structures"]["read_from"],
         fileformat=options["structures"]["file_format"],
-        dtype=dtype,
     )
     # Predict targets
     if hasattr(options, "targets"):
-        eval_targets = read_targets(conf=options["targets"], dtype=dtype)
+        eval_targets = read_targets(options["targets"])
         eval_dataset = Dataset(structure=eval_structures, energy=eval_targets["energy"])
         _eval_targets(model, eval_dataset)
 
@@ -185,7 +183,9 @@ def eval_model(
     if not hasattr(options, "targets"):
         # otherwise, the NLs will have been computed for the RMSE calculations above
         eval_structures = [
-            attach_neighbor_lists(structure, model.requested_neighbors_lists())
+            get_system_with_neighbors_lists(
+                structure, model.requested_neighbors_lists()
+            )
             for structure in eval_structures
         ]
     eval_options = ModelEvaluationOptions(

src/metatensor/models/experimental/alchemical_model/__init__.py

@@ -0,0 +1,2 @@
+from .model import Model, DEFAULT_HYPERS  # noqa: F401
+from .train import train  # noqa: F401