Merge pull request #199 from mir-group/develop

v0.5.4

.flake8

@@ -0,0 +1,4 @@
+[flake8]
+ignore = E226, E501, E741, E743, C901, W503, E203
+max-line-length = 127
+exclude = .git,__pycache__,docs/source/conf.py,old,build,dist,examples,tmp

.github/workflows/lint.yaml

@@ -0,0 +1,35 @@
+name: Check coding style
+
+on:
+  push:
+    branches:
+      - main
+      - develop
+  pull_request:
+    branches:
+      - main
+      - develop
+
+jobs:
+  black:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v2
+      - name: Black Check
+        uses: psf/black@stable
+        with:
+          version: "22.3.0"
+
+  flake8:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v2
+      - uses: actions/setup-python@v2
+        with:
+          python-version: '3.x'
+      - name: Install flake8
+        run: |
+          pip install flake8==4.0.1
+      - name: run flake8
+        run: |
+          flake8 . --count --show-source --statistics

.github/workflows/release.yaml

@@ -0,0 +1,29 @@
+name: Upload Python Package
+
+on:
+  release:
+    types: [created]
+
+jobs:
+  deploy:
+
+    runs-on: ubuntu-latest
+
+    steps:
+    - uses: actions/checkout@v2
+    - name: Set up Python
+      uses: actions/setup-python@v2
+      with:
+        python-version: 3.8
+    - name: Install dependencies
+      run: |
+        python -m pip install --upgrade pip
+        pip install setuptools wheel twine
+    - name: Build and publish
+      env:
+        TWINE_USERNAME: ${{ secrets.PYPI_USERNAME }}
+        TWINE_PASSWORD: ${{ secrets.PYPI_PASSWORD }}
+        GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+      run: |
+        python setup.py sdist bdist_wheel
+        twine upload dist/*

.github/workflows/tests.yml

@@ -1,4 +1,4 @@
-name: Check Syntax and Run Tests
+name: Run Tests
 
 on:
   push:
@@ -15,29 +15,24 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: [3.6, 3.9]
-        torch-version: [1.8.0, 1.10.0]
+        python-version: [3.7, 3.9]
+        torch-version: [1.8.0, 1.11.0]
 
     steps:
     - uses: actions/checkout@v2
     - name: Set up Python ${{ matrix.python-version }}
       uses: actions/setup-python@v2
       with:
         python-version: ${{ matrix.python-version }}
-    - name: Install flake8
-      run: |
-        pip install flake8
-    - name: Lint with flake8
-      run: |
-        flake8 . --count --show-source --statistics
     - name: Install dependencies
       env:
         TORCH: "${{ matrix.torch-version }}"
         GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
       run: |
         python -m pip install --upgrade pip
+        pip install setuptools wheel
         pip install torch==${TORCH} -f https://download.pytorch.org/whl/cpu/torch_stable.html
-        pip install .
+        pip install  --upgrade-strategy only-if-needed .
     - name: Install pytest
       run: |
         pip install pytest

.github/workflows/tests_develop.yml

@@ -1,4 +1,4 @@
-name: Check Syntax and Run Tests
+name: Run Tests
 
 on:
   push:
@@ -16,28 +16,23 @@ jobs:
     strategy:
       matrix:
         python-version: [3.9]
-        torch-version: [1.10.0]
+        torch-version: [1.11.0]
 
     steps:
     - uses: actions/checkout@v2
     - name: Set up Python ${{ matrix.python-version }}
       uses: actions/setup-python@v2
       with:
         python-version: ${{ matrix.python-version }}
-    - name: Install flake8
-      run: |
-        pip install flake8
-    - name: Lint with flake8
-      run: |
-        flake8 . --count --show-source --statistics
     - name: Install dependencies
       env:
         TORCH: "${{ matrix.torch-version }}"
         GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
       run: |
         python -m pip install --upgrade pip
+        pip install setuptools wheel
         pip install torch==${TORCH} -f https://download.pytorch.org/whl/cpu/torch_stable.html
-        pip install .
+        pip install --upgrade-strategy only-if-needed .
     - name: Install pytest
       run: |
         pip install pytest

.pre-commit-config.yaml

@@ -0,0 +1,14 @@
+exclude: '.git|.tox'
+default_stages: [commit]
+fail_fast: true
+
+repos:
+  - repo: https://github.com/psf/black
+    rev: 22.3.0
+    hooks:
+      - id: black
+
+  - repo: https://gitlab.com/pycqa/flake8
+    rev: 4.0.1
+    hooks:
+      - id: flake8

CHANGELOG.md

@@ -7,7 +7,29 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 Most recent change on the bottom.
 
 
-## [Unreleased]
+## [Unreleased] - 0.5.4
+
+## [0.5.4] - 2022-04-12
+### Added
+- `NequIPCalculator` now handles per-atom energies
+- Added `initial_model_state_strict` YAML option
+- `load_model_state` builder
+- fusion strategy support
+- `cumulative_wall` for early stopping
+- Deploy model from YAML file directly
+
+### Changed
+- Disallow PyTorch 1.9, which has some JIT bugs.
+- `nequip-deploy build` now requires `--train-dir` option when specifying the training session
+- Minimum Python version is now 3.7
+
+### Fixed
+- Better error in `Dataset.statistics` when field is missing
+- `NequIPCalculator` now outputs energy as scalar rather than `(1, 1)` array
+- `dataset: ase` now treats automatically adds `key_mapping` keys to `include_keys`, which is consistant with the npz dataset
+- fixed reloading models with `per_species_rescale_scales/shifts` set to `null`/`None`
+- graceful exit for `-n 0` in `nequip-benchmark`
+- Strictly correct CSV headers for metrics (#198)
 
 ## [0.5.3] - 2022-02-23
 ### Added

CONTRIBUTING.md

@@ -19,6 +19,7 @@ We use the [`black`](https://black.readthedocs.io/en/stable/index.html) code for
 ```
 
 Please run the formatter before you commit and certainly before you make a PR. The formatter can be easily set up to run automatically on file save in various editors.
+You can also use ``pre-commit install`` to install a [pre-commit](https://pre-commit.com/) hook.
 
 ## CUDA support
 

README.md

@@ -13,7 +13,7 @@ NequIP is an open-source code for building E(3)-equivariant interatomic potentia
 NequIP requires:
 
 * Python >= 3.6
-* PyTorch >= 1.8, <=1.11.*. PyTorch can be installed following the [instructions from their documentation](https://pytorch.org/get-started/locally/). Note that neither `torchvision` nor `torchaudio`, included in the default install command, are needed for NequIP.
+* PyTorch >= 1.8, !=1.9, <=1.11.*. PyTorch can be installed following the [instructions from their documentation](https://pytorch.org/get-started/locally/). Note that neither `torchvision` nor `torchaudio`, included in the default install command, are needed for NequIP.
 
 To install:
 
@@ -96,7 +96,7 @@ The `nequip-deploy` command is used to deploy the result of a training session i
 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/ where/to/put/deployed_model.pth
+nequip-deploy build --train-dir path/to/training/session/ where/to/put/deployed_model.pth
 ```
 For more details on this command, please run `nequip-deploy --help`.
 

configs/full.yaml

@@ -189,7 +189,7 @@ early_stopping_lower_bounds:
   LR: 1.0e-5
 
 early_stopping_upper_bounds:                                                       # stop early if a metric value is higher than the bound
-  wall: 1.0e+100
+  cumulative_wall: 1.0e+100
 
 # loss function
 loss_coeffs:                                                                       # different weights to use in a weighted loss functions

examples/lj/README.md

@@ -0,0 +1,10 @@
+Run commands with
+```
+PYTHONPATH=`pwd`:$PYTHONPATH nequip-* ...
+```
+so that the model from `lj.py` can be imported.
+
+For example, to create a deployed LJ model `lj.pth`:
+```bash
+PYTHONPATH=`pwd`:$PYTHONPATH nequip-deploy build --model lj.yaml lj.pth
+```

examples/lj/lj.py

@@ -0,0 +1,109 @@
+"""Example implementation of a Lennard-Jones potential in the NequIP framework.
+
+This serves as a basic example of how to write a NequIP framework model from scratch.
+"""
+
+from typing import Union
+
+import torch
+
+from torch_runstats.scatter import scatter
+
+from nequip.data import AtomicDataDict
+from nequip.nn import GraphModuleMixin, SequentialGraphNetwork, AtomwiseReduce
+
+
+# First, we define a model module to do the actual computation:
+class LennardJonesModule(GraphModuleMixin, torch.nn.Module):
+    """NequIP model module implementing a Lennard-Jones potential term.
+
+    See, for example, `lj/cut` in LAMMPS:
+    https://docs.lammps.org/pair_lj.html
+
+    Args:
+        initial_epsilon: initial value of the epsilon parameter.
+        initial_sigma: initial value of the sigma parameter.
+        trainable: whether epsilon and sigma should be trainable.
+            Default False.
+    """
+
+    def __init__(
+        self,
+        initial_epsilon: Union[float, torch.Tensor],
+        initial_sigma: Union[float, torch.Tensor],
+        trainable: bool = False,
+        irreps_in=None,
+    ) -> None:
+        super().__init__()
+        # We have to tell `GraphModuleMixin` what fields we expect in the input and output
+        # and what their irreps will be. Having basic geometry information (positions and edges)
+        # in the input is assumed.
+        # Per-atom energy is a scalar, so 0e.
+        self._init_irreps(irreps_out={AtomicDataDict.PER_ATOM_ENERGY_KEY: "0e"})
+        self.trainable = trainable
+        eps = torch.as_tensor(initial_epsilon)
+        sigma = torch.as_tensor(initial_sigma)
+        assert eps.ndim == sigma.ndim == 0, "epsilon and sigma must be scalars"
+        if self.trainable:
+            self.epsilon = torch.nn.Parameter(eps)
+            self.sigma = torch.nn.Parameter(sigma)
+        else:
+            # buffers act like parameters, but are not trainable
+            self.register_buffer("epsilon", eps)
+            self.register_buffer("sigma", sigma)
+
+    def forward(self, data: AtomicDataDict.Type) -> AtomicDataDict.Type:
+        """Run the module.
+
+        The module both takes and returns an `AtomicDataDict.Type` = `Dict[str, torch.Tensor]`.
+        Keys that the module does not modify/add are expected to be propagated to the output unchanged.
+        """
+        # If they are not already present, compute and add the edge vectors and lengths to `data`:
+        data = AtomicDataDict.with_edge_vectors(data, with_lengths=True)
+        # compute the LJ energy:
+        lj_eng = (self.sigma / data[AtomicDataDict.EDGE_LENGTH_KEY]) ** 6.0
+        lj_eng = torch.neg(lj_eng)
+        lj_eng = lj_eng + lj_eng.square()
+        # 2.0 because we do the slightly wastefull symmetric thing and let
+        # ij and ji each contribute half
+        # this avoids indexing out certain edges in the general case where
+        # the edges are not ordered.
+        lj_eng = (2.0 * self.epsilon) * lj_eng
+        # assign halves to centers
+        atomic_eng = scatter(
+            lj_eng,
+            # the edge indexes are of shape [2, n_edge];
+            # edge_index[0] is the index of the central atom of each edge
+            data[AtomicDataDict.EDGE_INDEX_KEY][0],
+            dim=0,
+            # dim_size says that even if some atoms have no edges, we still
+            # want an output energy for them (it will be zero)
+            dim_size=len(data[AtomicDataDict.POSITIONS_KEY]),
+        )
+        # NequIP defines standardized keys for typical fields:
+        data[AtomicDataDict.PER_ATOM_ENERGY_KEY] = atomic_eng
+        return data
+
+
+# then, we define a *model builder* function that builds an LJ energy model
+# from this and other modules:
+def LennardJonesPotential(config) -> SequentialGraphNetwork:
+    # `from_parameters` builds a model containing each of these modules in sequence
+    # from a configuration `config`
+    return SequentialGraphNetwork.from_parameters(
+        shared_params=config,
+        layers={
+            # LennardJonesModule will be built using options from `config`
+            "lj": LennardJonesModule,
+            # AtomwiseReduce will be built using the provided default parameters,
+            # and also those from `config`.
+            "total_energy_sum": (
+                AtomwiseReduce,
+                dict(
+                    reduce="sum",
+                    field=AtomicDataDict.PER_ATOM_ENERGY_KEY,
+                    out_field=AtomicDataDict.TOTAL_ENERGY_KEY,
+                ),
+            ),
+        },
+    )

examples/lj/lj.yaml

@@ -0,0 +1,15 @@
+# model
+model_builders:
+ - lj.LennardJonesPotential
+ # LennardJonesPotential gives an energy model
+ # ForceOutput takes an energy model and wraps it with an
+ # autodifferentiation call to get forces too:
+ - ForceOutput
+
+initial_epsilon: 1
+initial_sigma: 1
+
+# system
+r_max: 4.0
+chemical_symbols:
+  - H

examples/monkeypatch.py

@@ -5,18 +5,17 @@
 convolution for later analysis.
 """
 
-import torch
-
 from nequip.utils import Config, find_first_of_type
 from nequip.data import AtomicDataDict, AtomicData, dataset_from_config
 from nequip.nn import SequentialGraphNetwork, SaveForOutput
+from nequip.train import Trainer
 
 # The path to the original training session
 path = "../results/aspirin/minimal"
 
 # Load the model
-model = torch.load(path + "/best_model.pth")
-
+# there are other ways to do this, such as model_from_config etc.
+model = Trainer.load_model_from_training_session(traindir=path)
 
 # Find the SequentialGraphNetwork, which contains the
 # sequential bulk of the NequIP GNN model. To see the

nequip/_version.py

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