Merge pull request #134 from mir-group/develop

v0.5.1
mir-group · Jan 13, 2022 · 4e6a091 · 4e6a091
2 parents a5d160e + d4ccae2
commit 4e6a091
Show file tree

Hide file tree

Showing 42 changed files with 943 additions and 317 deletions.
diff --git a/.github/PULL_REQUEST_TEMPLATE.md b/.github/PULL_REQUEST_TEMPLATE.md
@@ -24,7 +24,7 @@ Resolves: #???
 <!-- Put an `x` in all the boxes that apply. If you're unsure about any of
      these, don't hesitate to ask. We're here to help! -->
 - [ ] My code follows the code style of this project and has been formatted using `black`.
-- [ ] All new and existing tests passed.
+- [ ] All new and existing tests passed, including on GPU (if relevant).
 - [ ] I have added tests that cover my changes (if relevant).
 - [ ] The option documentation (`docs/options`) has been updated with new or changed options.
 - [ ] I have updated `CHANGELOG.md`.

diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
@@ -16,7 +16,7 @@ jobs:
     strategy:
       matrix:
         python-version: [3.6, 3.9]
-        torch-version: [1.8.0, 1.9.0]
+        torch-version: [1.8.0, 1.10.0]
 
     steps:
     - uses: actions/checkout@v2
@@ -42,6 +42,10 @@ jobs:
       run: |
         pip install pytest
         pip install pytest-xdist[psutil]
+    - name: Download test data
+      run: |
+        mkdir benchmark_data
+        cd benchmark_data; wget "http://quantum-machine.org/gdml/data/npz/aspirin_ccsd.zip"; cd ..
     - name: Test with pytest
       run: |
         # See https://github.com/pytest-dev/pytest/issues/1075

diff --git a/.github/workflows/tests_develop.yml b/.github/workflows/tests_develop.yml
@@ -16,7 +16,7 @@ jobs:
     strategy:
       matrix:
         python-version: [3.9]
-        torch-version: [1.9.0]
+        torch-version: [1.10.0]
 
     steps:
     - uses: actions/checkout@v2
@@ -42,6 +42,10 @@ jobs:
       run: |
         pip install pytest
         pip install pytest-xdist[psutil]
+    - name: Download test data
+      run: |
+        mkdir benchmark_data
+        cd benchmark_data; wget "http://quantum-machine.org/gdml/data/npz/aspirin_ccsd.zip"; cd ..
     - name: Test with pytest
       run: |
         # See https://github.com/pytest-dev/pytest/issues/1075

diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -6,9 +6,29 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 Most recent change on the bottom.
 
+
 ## [Unreleased]
+
+## [0.5.1] - 2022-01-13
+### Added
+- `NequIPCalculator` can now be built via a `nequip_calculator()` function. This adds a minimal compatibility with [vibes](https://gitlab.com/vibes-developers/vibes/)
+- Added `avg_num_neighbors: auto` option
+- Asynchronous IO: during training, models are written asynchronously. Enable this with environment variable `NEQUIP_ASYNC_IO=true`.
+- `dataset_seed` to separately control randomness used to select training data (and their order).
+- The types may now be specified with a simpler `chemical_symbols` option
+- Equivariance testing reports per-field errors
+- `--equivariance-test n` tests equivariance on `n` frames from the training dataset
+
+### Changed
+- All fields now have consistant [N, dim] shaping
+- Changed default `seed` and `dataset_seed` in example YAMLs
+- Equivariance testing can only use training frames now
+
 ### Fixed
+- Equivariance testing no longer unintentionally skips translation
+- Correct cat dim for all registered per-graph fields
 - `PerSpeciesScaleShift` now correctly outputs when scales, but not shifts, are enabled— previously it was broken and would only output updated values when both were enabled.
+- `nequip-evaluate` outputs correct species to the `extxyz` file when a chemical symbol <-> type mapping exists for the test dataset
 
 ## [0.5.0] - 2021-11-24
 ### Changed

diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
@@ -1,8 +1,27 @@
 # Contributing to NequIP
 
+Issues and pull requests are welcome!
+
+**!! If you want to make a major change, or one whose correct location/implementation is not obvious, please reach out to discuss with us first. !!**
+
+In general:
+ - Optional additions/alternatives to how the model is built or initialized should be implemented as model builders (see `nequip.model`)
+ - New model features should be implemented as new modules when possible
+ - Added options should be documented in the docs and changes in the CHANGELOG.md file
+
+Unless they fix a significant bug with immediate impact, **all PRs should be onto the `develop` branch!**
+
 ## Code style
 
 We use the [`black`](https://black.readthedocs.io/en/stable/index.html) code formatter with default settings and the flake8 linter with settings:
 ```
---ignore=E501,W503,E203
-```
+--ignore=E226,E501,E741,E743,C901,W503,E203 --max-line-length=127
+```
+
+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.
+
+## CUDA support
+
+All additions should support CUDA/GPU.
+
+If possible, please test your changes on a GPU— the CI tests on GitHub actions do not have GPU resources available.
diff --git a/README.md b/README.md
@@ -145,9 +145,11 @@ NequIP is being developed by:
 
 under the guidance of [Boris Kozinsky at Harvard](https://bkoz.seas.harvard.edu/).
 
-## Contact & questions
+## Contact, questions, and contributing
 
 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).
+
+If you want to contribute to the code, please read [`CONTRIBUTING.md`](CONTRIBUTING.md).
diff --git a/configs/example.yaml b/configs/example.yaml
@@ -6,7 +6,8 @@
 # if 'root'/'run_name' exists, 'root'/'run_name'_'year'-'month'-'day'-'hour'-'min'-'s' will be used instead.
 root: results/toluene
 run_name: example-run-toluene
-seed: 0                                                                           # random number seed for numpy and torch
+seed: 123
+dataset_seed: 456                                                                          # random number seed for numpy and torch
 append: true                                                                      # 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
 
@@ -39,7 +40,7 @@ PolynomialCutoff_p: 6
 # radial network
 invariant_layers: 2                                                               # number of radial layers, usually 1-3 works best, smaller is faster
 invariant_neurons: 64                                                             # number of hidden neurons in radial function, smaller is faster
-avg_num_neighbors: null                                                           # number of neighbors to divide by, null => no normalization.
+avg_num_neighbors: auto                                                           # number of neighbors to divide by, null => no normalization.
 use_sc: true                                                                      # use self-connection or not, usually gives big improvement
 
 # data set
@@ -58,10 +59,10 @@ key_mapping:
 npz_fixed_field_keys:                                                              # fields that are repeated across different examples
   - atomic_numbers
 
-# A mapping of chemical species to type indexes is necessary if the dataset is provided with atomic numbers instead of type indexes.
-chemical_symbol_to_type:
-  H: 0
-  C: 1
+# A list of atomic types to be found in the data. The NequIP types will be named with the chemical symbols, and inputs with the correct atomic numbers will be mapped to the corresponding types.
+chemical_symbols:
+  - H
+  - C
 
 # logging
 wandb: true                                                                       # we recommend using wandb for logging, we'll turn it off here as it's optional

diff --git a/configs/full.yaml b/configs/full.yaml
@@ -10,11 +10,12 @@
 # if 'root'/'run_name' exists, 'root'/'run_name'_'year'-'month'-'day'-'hour'-'min'-'s' will be used instead.
 root: results/toluene
 run_name: example-run-toluene
-seed: 0                                                                           # random number seed for numpy and torch
+seed: 123
+dataset_seed: 456                                                                          # random number seed for numpy and torch
 append: true                                                                      # 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: false                                                                 # whether to use TensorFloat32 if it is available
-# device: cuda                                                                     # which device to use. Default: automatically detected cuda or "cpu"
+# device:  cuda                                                                     # which device to use. Default: automatically detected cuda or "cpu"
 
 # network
 r_max: 4.0                                                                        # cutoff radius in length units, here Angstrom, this is an important hyperparamter to scan
@@ -45,7 +46,7 @@ PolynomialCutoff_p: 6
 # radial network
 invariant_layers: 2                                                               # number of radial layers, usually 1-3 works best, smaller is faster
 invariant_neurons: 64                                                             # number of hidden neurons in radial function, smaller is faster
-avg_num_neighbors: null                                                           # number of neighbors to divide by, null => no normalization.
+avg_num_neighbors: auto                                                           # number of neighbors to divide by, null => no normalization.
 use_sc: true                                                                      # use self-connection or not, usually gives big improvement
 compile_model: false                                                              # whether to compile the constructed model to TorchScript
 
@@ -88,16 +89,21 @@ npz_fixed_field_keys:
 # key_mapping:
 #   free_energy: total_energy
 
-# A mapping of chemical species to type indexes is necessary if the dataset is provided with atomic numbers instead of type indexes.
-chemical_symbol_to_type:
-  H: 0
-  C: 1
-
-# Alternatively, if the dataset has type indicess, the total number of types is all that is required:
+# A list of atomic types to be found in the data. The NequIP types will be named with the chemical symbols, and inputs with the correct atomic numbers will be mapped to the corresponding types.
+chemical_symbols:
+ - H
+ - C
+# Alternatively, you may explicitly specify which chemical species maps to which type in NequIP (type index; the name is still taken from the chemical symbol)
+# chemical_symbol_to_type:
+#   H: 0
+#   C: 1
+
+# Alternatively, if the dataset has type indices, you may give the names for the types in order:
+# (this also sets the number of types)
 # type_names:
-#   0: my_type
-#   1: atom
-#   2: thing
+#   - my_type
+#   - atom
+#   - thing
 
 # As an alternative option to npz, you can also pass data ase ASE Atoms-objects
 # This can often be easier to work with, simply make sure the ASE Atoms object

diff --git a/configs/minimal.yaml b/configs/minimal.yaml
@@ -1,7 +1,8 @@
 # general
 root: results/aspirin
 run_name: minimal
-seed: 0
+seed: 123
+dataset_seed: 456
 
 # network
 num_basis: 8
@@ -25,16 +26,11 @@ key_mapping:
   R: pos                                                                           # raw atomic positions
 npz_fixed_field_keys:                                                              # fields that are repeated across different examples
   - atomic_numbers
-# A mapping of chemical species to type indexes is necessary if the dataset is provided with atomic numbers instead of type indexes.
-chemical_symbol_to_type:
-  H: 0
-  C: 1
-  O: 2
-# Alternatively, if the dataset has type indexes, the total number of types is all that is required:
-# type_names:
-#   0: my_type
-#   1: atom
-#   2: thing
+
+chemical_symbols:
+  - H
+  - O
+  - C
 
 # logging
 wandb: false

diff --git a/configs/minimal_eng.yaml b/configs/minimal_eng.yaml
@@ -29,16 +29,11 @@ key_mapping:
   R: pos                                                                           # raw atomic positions
 npz_fixed_field_keys:                                                              # fields that are repeated across different examples
   - atomic_numbers
-# A mapping of chemical species to type indexes is necessary if the dataset is provided with atomic numbers instead of type indexes.
-chemical_symbol_to_type:
-  H: 0
-  C: 1
-  O: 2
-# Alternatively, if the dataset has type indexes, the total number of types is all that is required:
-# type_names:
-#   0: my_type
-#   1: atom
-#   2: thing
+
+chemical_symbols:
+  - H
+  - O
+  - C
 
 # logging
 wandb: false

diff --git a/docs/options/dataset.rst b/docs/options/dataset.rst
@@ -13,6 +13,11 @@ type_names
     | Type: NoneType
     | Default: ``None``
 
+chemical_symbols
+^^^^^^^^^^^^^^^^
+    | Type: NoneType
+    | Default: ``None``
+
 chemical_symbol_to_type
 ^^^^^^^^^^^^^^^^^^^^^^^
     | Type: NoneType

diff --git a/nequip/_version.py b/nequip/_version.py
@@ -2,4 +2,4 @@
 # See Python packaging guide
 # https://packaging.python.org/guides/single-sourcing-package-version/
 
-__version__ = "0.5.0"
+__version__ = "0.5.1"
diff --git a/nequip/ase/nequip_calculator.py b/nequip/ase/nequip_calculator.py
@@ -10,6 +10,11 @@
 import nequip.scripts.deploy
 
 
+def nequip_calculator(model, **kwargs):
+    """Build ASE Calculator directly from deployed model."""
+    return NequIPCalculator.from_deployed_model(model, **kwargs)
+
+
 class NequIPCalculator(Calculator):
     """NequIP ASE Calculator.