Merge pull request #61 from mir-group/develop

0.3.3

CHANGELOG.md

@@ -7,12 +7,27 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 Most recent change on the bottom.
 
 ## [Unreleased]
+
+## [0.3.3] - 2021-08-11
+### Added
+- `to_ase` method in `AtomicData.py` to convert `AtomicData` object to (list of) `ase.Atoms` object(s)
+- `SequentialGraphNetwork` now has insertion methods
+- `nn.SaveForOutput`
+- `nequip-evaluate` command for evaluating (metrics on) trained models
+- `AtomicData.from_ase` now catches `energy`/`energies` arrays
+
 ### Changed
 - Nonlinearities now specified with `e` and `o` instead of `1` and `-1`
+- Update interfaces for `torch_geometric` 1.7.1 and `e3nn` 0.3.3
+- `nonlinearity_scalars` now also affects the nonlinearity used in the radial net of `InteractionBlock`
+- Cleaned up naming of initializers
 
 ### Fixed
 - Fix specifying nonlinearities when wandb enabled
+- `Final` backport for <3.8 compatability
 - Fixed `nequip-*` commands when using `pip install`
+- Default models rescale per-atom energies, and not just total
+- Fixed Python <3.8 backward compatability with `atomic_save`
 
 ## [0.3.2] - 2021-06-09
 ### Added

README.md

@@ -2,6 +2,7 @@
 
 NequIP is an open-source code for building E(3)-equivariant interatomic potentials.
 
+[![Documentation Status](https://readthedocs.org/projects/nequip/badge/?version=latest)](https://nequip.readthedocs.io/en/latest/?badge=latest)
 
 ![nequip](./nequip.png)
 
@@ -12,35 +13,18 @@ NequIP is an open-source code for building E(3)-equivariant interatomic potentia
 NequIP requires:
 
 * Python >= 3.6
-* PyTorch = 1.8
+* PyTorch >= 1.8
 
 To install:
 
-* Install [PyTorch Geometric](https://github.com/rusty1s/pytorch_geometric), make sure to install this with your correct version of CUDA/CPU and to use PyTorch Geometric version 1.7.0:
-
-```
-pip install torch-scatter -f https://pytorch-geometric.com/whl/torch-1.8.0+${CUDA}.html
-pip install torch-sparse -f https://pytorch-geometric.com/whl/torch-1.8.0+${CUDA}.html
-pip install torch-cluster -f https://pytorch-geometric.com/whl/torch-1.8.0+${CUDA}.html
-pip install torch-spline-conv -f https://pytorch-geometric.com/whl/torch-1.8.0+${CUDA}.html
-pip install torch-geometric==1.7.0
-pip install e3nn==0.2.9
-```
-
-where ```${CUDA}``` should be replaced by either ```cpu```, ```cu101```, ```cu102```, or ```cu111``` depending on your PyTorch installation, for details see [here](https://github.com/rusty1s/pytorch_geometric). 
-
-* Install [e3nn](https://github.com/e3nn/e3nn), version 0.2.9: 
-
-```
-pip install e3nn==0.2.9
-```
+* Install [PyTorch Geometric](https://github.com/rusty1s/pytorch_geometric), following [their installation instructions](https://pytorch-geometric.readthedocs.io/en/latest/notes/installation.html) and making sure to install with the correct version of CUDA. (Please note that `torch_geometric>=1.7.1)` is required.)
 
 * Install our fork of [`pytorch_ema`](https://github.com/Linux-cpp-lisp/pytorch_ema) for using an Exponential Moving Average on the weights: 
 ```bash
-$ pip install git+https://github.com/Linux-cpp-lisp/pytorch_ema@context_manager#egg=torch_ema
+$ pip install "git+https://github.com/Linux-cpp-lisp/pytorch_ema@context_manager#egg=torch_ema"
 ```
 
-* We use [Weights&Biases](https://wandb.ai) to keep track of experiments. This is not a strict requirement, you can use our package without this, but it may make your life easier. If you want to use it, create an account [here](https://wandb.ai) and install it: 
+* We use [Weights&Biases](https://wandb.ai) to keep track of experiments. This is not a strict requirement — you can use our package without it — but it may make your life easier. If you want to use it, create an account [here](https://wandb.ai) and install the Python package:
 
 ```
 pip install wandb
@@ -56,13 +40,15 @@ pip install .
 
 ### Installation Issues
 
-We recommend running the tests using ```pytest``` on a CPU: 
+We recommend running the tests using ```pytest```: 
 
 ```
 pip install pytest
 pytest ./tests/
 ```
 
+While the tests are somewhat compute intensive, we've known them to hang on certain systems that have GPUs. If this happens to you, please report it along with information on your software environment in the [Issues](https://github.com/mir-group/nequip/issues)!
+
 ## Usage
 
 **! PLEASE NOTE:** the first few training epochs/calls to a NequIP model can be painfully slow. This is expected behaviour as the [profile-guided optimization of TorchScript models](https://program-transformations.github.io/slides/pytorch_neurips.pdf) takes a number of calls to warm up before optimizing the model. This occurs regardless of whether the entire model is compiled because many core components from e3nn are compiled and optimized through TorchScript.
@@ -83,22 +69,45 @@ A number of example configuration files are provided:
 
 Training runs can be restarted using `nequip-restart`; training that starts fresh or restarts depending on the existance of the working directory can be launched using `nequip-requeue`. All `nequip-*` commands accept the `--help` option to show their call signatures and options.
 
-### In-depth tutorial 
+### Evaluating trained models (and their error)
+
+The `nequip-evaluate` command can be used to evaluate a trained model on a specified dataset, optionally computing error metrics or writing the results to an XYZ file for further processing.
 
-A more in-depth introduction to the internals of NequIP can be found in the [tutorial notebook](https://deepnote.com/project/2412ca93-7ad1-4458-972c-5d5add5a667e).
+The simplest command is:
+```bash
+$ nequip-evaluate --train-dir /path/to/training/session/
+```
+which will evaluate the original training error metrics over any part of the original dataset not used in the training or validation sets.
+
+For more details on this command, please run `nequip-evaluate --help`.
+
+### Deploying models
+
+The `nequip-deploy` command is used to deploy the result of a training session into a model that can be stored and used for inference.
+It compiles a NequIP model trained in Python to [TorchScript](https://pytorch.org/docs/stable/jit.html).
+The result is an optimized model file that has no dependency on the `nequip` Python library, or even on Python itself:
+```bash
+nequip-deploy build path/to/training/session/ path/to/deployed.pth
+```
+For more details on this command, please run `nequip-deploy --help`.
+
+### Using models in Python
+
+Both deployed and undeployed models can be used in Python code; for details, see the end of the [Developer's tutorial](https://deepnote.com/project/2412ca93-7ad1-4458-972c-5d5add5a667e) mentioned again below.
+
+An ASE calculator is also provided in `nequip.dynamics`.
 
 ### LAMMPS Integration 
 
 NequIP is integrated with the popular Molecular Dynamics code [LAMMPS](https://www.lammps.org/) which allows for MD simulations over large time- and length-scales and gives users access to the full suite of LAMMPS features. 
 
 The interface is implemented as `pair_style nequip`. Using it requires two simple steps: 
 
-1. Deploy a trained NequIP model. This step compiles a NequIP model trained in Python to [TorchScript](https://pytorch.org/docs/stable/jit.html). 
-The result is an optimized model file that has no Python dependency and can be used by standalone C++ programs such as LAMMPS: 
-
+1. Deploy a trained NequIP model, as discussed above.
 ```
 nequip-deploy build path/to/training/session/ path/to/deployed.pth
 ```
+The result is an optimized model file that has no Python dependency and can be used by standalone C++ programs such as LAMMPS.
 
 2. Change the LAMMPS input file to the nequip `pair_style` and point it to the deployed NequIP model:
 
@@ -107,40 +116,39 @@ pair_style	nequip
 pair_coeff	* * deployed.pth
 ```
 
-For installation instructions, please see the NequIP `pair_style` repo at https://github.com/mir-group/pair_nequip.
+For installation instructions, please see the [`pair_nequip` repository](https://github.com/mir-group/pair_nequip).
 
 
-## References
+## Developer's tutorial 
 
-The theory behind NequIP is described in our preprint [1]. NequIP's backend builds on e3nn, a general framework for building E(3)-equivariant neural networks [2]. 
+A more in-depth introduction to the internals of NequIP can be found in the [tutorial notebook](https://deepnote.com/project/2412ca93-7ad1-4458-972c-5d5add5a667e). This notebook discusses theoretical background as well as the Python interfaces that can be used to train and call models.
 
-    [1] https://arxiv.org/abs/2101.03164
-    [2] https://github.com/e3nn/e3nn
+Please note that for most common usecases, including customized models, the `nequip-*` commands should be prefered for training models.
+
+## References & citing
+
+The theory behind NequIP is described in our preprint (1). NequIP's backend builds on e3nn, a general framework for building E(3)-equivariant neural networks (2). If you use this repository in your work, please consider citing NequIP (1) and e3nn (3):
+
+ 1. https://arxiv.org/abs/2101.03164
+ 2. https://e3nn.org
+ 3. https://doi.org/10.5281/zenodo.3724963
 
 ## Authors
 
 NequIP is being developed by:
 
-    - Simon Batzner
-    - Albert Musaelian
-    - Lixin Sun
-    - Anders Johansson
-    - Mario Geiger
-    - Tess Smidt
-
-under the guidance of Boris Kozinsky at Harvard.
+ - Simon Batzner
+ - Albert Musaelian
+ - Lixin Sun
+ - Anders Johansson
+ - Mario Geiger
+ - Tess Smidt
 
+under the guidance of [Boris Kozinsky at Harvard](https://bkoz.seas.harvard.edu/).
 
 ## Contact & questions
 
 If you have questions, please don't hesitate to reach out at batzner[at]g[dot]harvard[dot]edu. 
 
 If you find a bug or have a proposal for a feature, please post it in the [Issues](https://github.com/mir-group/nequip/issues).
 If you have a question, topic, or issue that isn't obviously one of those, try our [GitHub Disucssions](https://github.com/mir-group/nequip/discussions).
-
-## Citing
-
-If you use this repository in your work, please consider citing NequIP (1) and e3nn (2): 
-
-    [1] https://arxiv.org/abs/2101.03164
-    [2] https://doi.org/10.5281/zenodo.3724963

configs/full.yaml

@@ -12,6 +12,7 @@ seed: 0
 restart: false                                                                    # set True for a restarted run
 append: false                                                                     # set True if a restarted run should append to the previous log file
 default_dtype: float32                                                            # type of float to use, e.g. float32 and float64
+allow_tf32: True                                                                  # whether to use TensorFloat32 if it is available
 
 # network
 r_max: 4.0                                                                        # cutoff radius in length units
@@ -170,11 +171,15 @@ optimizer_weight_decay: 0
 
 # weight initialization
 # this can be the importable name of any function that can be `model.apply`ed to initialize some weights in the model. NequIP provides a number of useful initializers:
+# For more details please see the docstrings of the individual initializers
 #model_initializers:
 #  - nequip.utils.initialization.uniform_initialize_fcs
-#  - nequip.utils.initialization.uniform_initialize_tps
-#  - nequip.utils.initialization.orthogonal_initialize_linears
-#  - nequip.utils.initialization.uniform_initialize_linears
+#  - nequip.utils.initialization.uniform_initialize_equivariant_linears
+#  - nequip.utils.initialization.uniform_initialize_tp_internal_weights
+#  - nequip.utils.initialization.xavier_initialize_fcs
+#  - nequip.utils.initialization.(unit_)orthogonal_initialize_equivariant_linears
+#  - nequip.utils.initialization.(unit_)orthogonal_initialize_fcs
+#  - nequip.utils.initialization.(unit_)orthogonal_initialize_e3nn_fcs
 
 # lr scheduler, currently only supports the two options listed below, if you need more please file an issue
 # first: on-plateau, reduce lr by factory of lr_scheduler_factor if metrics_key hasn't improved for lr_scheduler_patience epoch

configs/minimal.yaml

@@ -11,11 +11,6 @@ conv_to_output_hidden_irreps_out: 16x0e
 feature_irreps_hidden: 16x0o + 16x0e + 16x1o + 16x1e + 16x2o + 16x2e
 model_uniform_init: false
 
-model_initializers:
-  - nequip.utils.initialization.uniform_initialize_fcs
-  - nequip.utils.initialization.uniform_initialize_tps
-  - nequip.utils.initialization.orthogonal_initialize_linears
-
 # data
 dataset: aspirin
 dataset_file_name: benchmark_data/aspirin_ccsd-train.npz

nequip/_version.py

@@ -2,4 +2,4 @@
 # See Python packaging guide
 # https://packaging.python.org/guides/single-sourcing-package-version/
 
-__version__ = "0.3.2"
+__version__ = "0.3.3"

nequip/data/AtomicData.py

@@ -5,11 +5,12 @@
 
 import warnings
 from copy import deepcopy
-from typing import Union, Tuple, Dict, Optional
+from typing import Union, Tuple, Dict, Optional, List
 from collections.abc import Mapping
 
 import numpy as np
 import ase.neighborlist
+import ase
 from ase.calculators.singlepoint import SinglePointCalculator, SinglePointDFTCalculator
 
 import torch
@@ -200,9 +201,10 @@ def from_ase(cls, atoms, r_max, **kwargs):
 
         Respects ``atoms``'s ``pbc`` and ``cell``.
 
-        Automatically recognize force, energy (overridden by free energy tag)
-        get_atomic_numbers() will be stored as the atomic_numbers attributes
+        First tries to extract energies and forces from a single-point calculator associated with the ``Atoms`` if one is present and has those fields.
+        If either is not found, the method will look for ``energy``/``energies`` and ``force``/``forces`` in ``atoms.arrays``.
 
+        `get_atomic_numbers()` will be stored as the atomic_numbers attribute.
 
         Args:
             atoms (ase.Atoms): the input.
@@ -233,10 +235,19 @@ def from_ase(cls, atoms, r_max, **kwargs):
                         "energy"
                     )
 
-        elif "forces" in atoms.arrays:
-            add_fields[AtomicDataDict.FORCE_KEY] = atoms.arrays["forces"]
-        elif "force" in atoms.arrays:
-            add_fields[AtomicDataDict.FORCE_KEY] = atoms.arrays["force"]
+        if AtomicDataDict.FORCE_KEY not in add_fields:
+            # Get it from arrays
+            for k in ("force", "forces"):
+                if k in atoms.arrays:
+                    add_fields[AtomicDataDict.FORCE_KEY] = atoms.arrays[k]
+                    break
+
+        if AtomicDataDict.TOTAL_ENERGY_KEY not in add_fields:
+            # Get it from arrays
+            for k in ("energy", "energies"):
+                if k in atoms.arrays:
+                    add_fields[AtomicDataDict.TOTAL_ENERGY_KEY] = atoms.arrays[k]
+                    break
 
         add_fields[AtomicDataDict.ATOMIC_NUMBERS_KEY] = atoms.get_atomic_numbers()
 
@@ -249,6 +260,72 @@ def from_ase(cls, atoms, r_max, **kwargs):
             **add_fields,
         )
 
+    def to_ase(self) -> Union[List[ase.Atoms], ase.Atoms]:
+        """Build a (list of) ``ase.Atoms`` object(s) from an ``AtomicData`` object.
+
+        For each unique batch number provided in ``AtomicDataDict.BATCH_KEY``,
+        an ``ase.Atoms`` object is created. If ``AtomicDataDict.BATCH_KEY`` does not
+        exist in self, a single ``ase.Atoms`` object is created.
+
+        Returns:
+            A list of ``ase.Atoms`` objects if ``AtomicDataDict.BATCH_KEY`` is in self
+            and is not None. Otherwise, a single ``ase.Atoms`` object is returned.
+        """
+        positions = self.pos
+        if positions.device != torch.device("cpu"):
+            raise TypeError(
+                "Explicitly move this `AtomicData` to CPU using `.to()` before calling `to_ase()`."
+            )
+        atomic_nums = self.atomic_numbers
+        pbc = getattr(self, AtomicDataDict.PBC_KEY, None)
+        cell = getattr(self, AtomicDataDict.CELL_KEY, None)
+        batch = getattr(self, AtomicDataDict.BATCH_KEY, None)
+        energy = getattr(self, AtomicDataDict.TOTAL_ENERGY_KEY, None)
+        force = getattr(self, AtomicDataDict.FORCE_KEY, None)
+        do_calc = energy is not None or force is not None
+
+        if cell is not None:
+            cell = cell.view(-1, 3, 3)
+        if pbc is not None:
+            pbc = pbc.view(-1, 3)
+
+        if batch is not None:
+            n_batches = batch.max() + 1
+            cell = cell.expand(n_batches, 3, 3) if cell is not None else None
+            pbc = pbc.expand(n_batches, 3) if pbc is not None else None
+        else:
+            n_batches = 1
+
+        batch_atoms = []
+        for batch_idx in range(n_batches):
+            if batch is not None:
+                mask = batch == batch_idx
+            else:
+                mask = slice(None)
+
+            mol = ase.Atoms(
+                numbers=atomic_nums[mask],
+                positions=positions[mask],
+                cell=cell[batch_idx] if cell is not None else None,
+                pbc=pbc[batch_idx] if pbc is not None else None,
+            )
+
+            if do_calc:
+                fields = {}
+                if energy is not None:
+                    fields["energy"] = energy[batch_idx].cpu().numpy()
+                if force is not None:
+                    fields["forces"] = force[mask].cpu().numpy()
+                mol.calc = SinglePointCalculator(mol, **fields)
+
+            batch_atoms.append(mol)
+
+        if batch is not None:
+            return batch_atoms
+        else:
+            assert len(batch_atoms) == 1
+            return batch_atoms[0]
+
     def get_edge_vectors(data: Data) -> torch.Tensor:
         data = AtomicDataDict.with_edge_vectors(AtomicData.to_AtomicDataDict(data))
         return data[AtomicDataDict.EDGE_VECTORS_KEY]
@@ -263,8 +340,15 @@ def to_AtomicDataDict(
             keys = data.keys()
         else:
             raise ValueError(f"Invalid data `{repr(data)}`")
+
         return {
-            k: data[k] for k in keys if (k not in exclude_keys and data[k] is not None)
+            k: data[k]
+            for k in keys
+            if (
+                k not in exclude_keys
+                and data[k] is not None
+                and isinstance(data[k], torch.Tensor)
+            )
         }
 
     @classmethod
@@ -322,7 +406,7 @@ def without_nodes(self, which_nodes):
             elif k == AtomicDataDict.CELL_KEY:
                 new_dict[k] = self[k]
             else:
-                if len(self[k]) == self.num_nodes:
+                if isinstance(self[k], torch.Tensor) and len(self[k]) == self.num_nodes:
                     new_dict[k] = self[k][mask]
                 else:
                     new_dict[k] = self[k]