Merge pull request #149 from mir-group/develop

0.5.2

CHANGELOG.md

@@ -7,7 +7,33 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 Most recent change on the bottom.
 
 
-## [Unreleased]
+## [Unreleased] 
+
+## [0.5.2] - 2022-02-04
+### Added
+- Model builders may now process only the configuration
+- Allow irreps to optionally be specified through the simplified keys `l_max`, `parity`, and `num_features`
+- `wandb.watch` via `wandb_watch` option
+- Allow polynomial cutoff _p_ values besides 6.0
+- `nequip-evaluate` now sets a default `r_max` taken from the model for the dataset config
+- Support multiple rescale layers in trainer
+- `AtomicData.to_ase` supports arbitrary fields
+- `nequip-evaluate` can now output arbitrary fields to an XYZ file
+- `nequip-evaluate` reports which frame in the original dataset was used as input for each output frame
+
+### Changed
+- `minimal.yaml`, `minimal_eng.yaml`, and `example.yaml` now use the simplified irreps options `l_max`, `parity`, and `num_features`
+- Default value for `resnet` is now `False`
+
+### Fixed
+- Handle one of `per_species_shifts`/`scales` being `null` when the other is a dataset statistc
+- `include_frames` now works with ASE datasets
+- no training data labels in input_data
+- Average number of neighbors no longer crashes sometimes when not all nodes have neighbors (small cutoffs)
+- Handle field registrations correctly in `nequip-evaluate`
+
+### Removed
+- `compile_model`
 
 ## [0.5.1] - 2022-01-13
 ### Added

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.10.*. 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. NequIP is also not currently compatible with PyTorch 1.10; PyTorch 1.9 can be specified with `pytorch==1.9` in the install command.
+* 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.
 
 To install:
 
@@ -33,7 +33,7 @@ pip install .
 
 ### Installation Issues
 
-The easiest way to check if your installation is working is to train a toy model:
+The easiest way to check if your installation is working is to train a **toy** model:
 ```bash
 $ nequip-train configs/minimal.yaml
 ```
@@ -50,15 +50,15 @@ To run the full tests, including a set of longer/more intensive integration test
 pytest tests/
 ```
 
-Note: the integration tests have hung in the past 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)!
+If a GPU is present, the unit tests will use it.
 
 ## Tutorial 
 
-The best way to learn how to use NequIP is through the [Colab Tutorial](https://bit.ly/mrs-nequip). This will run entirely on Google Hardware, you will not need to install anything, but can instead simply run it in your browser. 
+The best way to learn how to use NequIP is through the [Colab Tutorial](https://bit.ly/mrs-nequip). This will run entirely on Google's cloud virtual machine; you do not need to install or run anything locally. 
 
 ## 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.
+**! PLEASE NOTE:** the first few 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. (The `nequip-benchmark` script accounts for this.)
 
 ### Basic network training
 
@@ -69,12 +69,14 @@ $ nequip-train configs/example.yaml
 ```
 
 A number of example configuration files are provided:
- - [`configs/minimal.yaml`](configs/minimal.yaml): A minimal example of training a toy model on force data.
- - [`configs/minimal_eng.yaml`](configs/minimal_eng.yaml): The same, but for a toy model that predicts and trains on only energy labels.
- - [`configs/example.yaml`](configs/example.yaml): Training a more realistic model on forces and energies. Start here for real models.
- - [`configs/full.yaml`](configs/full.yaml): A complete configuration file containing all available options along with documenting comments.
+ - [`configs/minimal.yaml`](configs/minimal.yaml): A minimal example of training a **toy** model on force data.
+ - [`configs/minimal_eng.yaml`](configs/minimal_eng.yaml): The same, but for a **toy** model that predicts and trains on only energy labels.
+ - [`configs/example.yaml`](configs/example.yaml): Training a more realistic model on forces and energies. **Start here for real models!**
+ - [`configs/full.yaml`](configs/full.yaml): A complete configuration file containing all available options along with documenting comments. This file is **for reference**, `example.yaml` is the right starting point for a project.
 
-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.
+Training runs can also be restarted by running the same `nequip-train` command if the `append: True` option is specified in the original YAML. (Otherwise, a new training run with a different name can be started from the loaded state of the previous run.)
+
+All `nequip-*` commands accept the `--help` option to show their call signatures and options.
 
 ### Evaluating trained models (and their error)
 

configs/example.yaml

@@ -14,12 +14,10 @@ default_dtype: float32
 # network
 r_max: 4.0                                                                        # cutoff radius in length units, here Angstrom, this is an important hyperparamter to scan
 num_layers: 4                                                                     # number of interaction blocks, we find 4-6 to work best
-chemical_embedding_irreps_out: 32x0e                                              # irreps for the chemical embedding of species
-feature_irreps_hidden: 32x0o + 32x0e + 32x1o + 32x1e                              # irreps used for hidden features, here we go up to lmax=1, with even and odd parities; for more accurate but slower networks, use l=2 or higher, smaller number of features is faster
-irreps_edge_sh: 0e + 1o                                                           # irreps of the spherical harmonics used for edges. If a single integer, indicates the full SH up to L_max=that_integer
-conv_to_output_hidden_irreps_out: 16x0e                                           # irreps used in hidden layer of output block
+l_max: 1                                                                          # the maximum irrep order (rotation order) for the network's features
+parity: true                                                                      # whether to include features with odd mirror parity
+num_features: 32                                                                  # the multiplicity of the features
 nonlinearity_type: gate                                                           # may be 'gate' or 'norm', 'gate' is recommended
-resnet: false                                                                     # set true to make interaction block a resnet-style update
 
 # scalar nonlinearities to use — available options are silu, ssp (shifted softplus), tanh, and abs.
 # Different nonlinearities are specified for e (even) and o (odd) parity;
@@ -67,8 +65,7 @@ chemical_symbols:
 # logging
 wandb: true                                                                       # we recommend using wandb for logging, we'll turn it off here as it's optional
 wandb_project: toluene-example                                                     # project name used in wandb
-wandb_resume: true                                                                 # if true and restart is true, wandb run data will be restarted and updated.
-                                                                                   # if false, a new wandb run will be generated
+
 verbose: info                                                                      # the same as python logging, e.g. warning, info, debug, error. case insensitive
 log_batch_freq: 1000000                                                                  # batch frequency, how often to print training errors withinin the same epoch
 log_epoch_freq: 1                                                                  # epoch frequency, how often to print and save the model

configs/full.yaml

@@ -20,12 +20,23 @@ allow_tf32: false
 # network
 r_max: 4.0                                                                        # cutoff radius in length units, here Angstrom, this is an important hyperparamter to scan
 num_layers: 4                                                                     # number of interaction blocks, we find 4-6 to work best
-chemical_embedding_irreps_out: 32x0e                                              # irreps for the chemical embedding of species
-feature_irreps_hidden: 32x0o + 32x0e + 32x1o + 32x1e                              # irreps used for hidden features, here we go up to lmax=1, with even and odd parities; for more accurate but slower networks, use l=2 or higher, smaller number of features is faster
-irreps_edge_sh: 0e + 1o                                                           # irreps of the spherical harmonics used for edges. If a single integer, indicates the full SH up to L_max=that_integer
-conv_to_output_hidden_irreps_out: 16x0e                                           # irreps used in hidden layer of output block
+
+l_max: 1                                                                          # the maximum irrep order (rotation order) for the network's features
+parity: true                                                                      # whether to include features with odd mirror parity
+num_features: 32                                                                  # the multiplicity of the features
+
+# alternatively, the irreps of the features in various parts of the network can be specified directly:
+# the following options use e3nn irreps notation
+# either these four options, or the above three options, should be provided--- they cannot be mixed.
+# chemical_embedding_irreps_out: 32x0e                                              # irreps for the chemical embedding of species
+# feature_irreps_hidden: 32x0o + 32x0e + 32x1o + 32x1e                              # irreps used for hidden features, here we go up to lmax=1, with even and odd parities; for more accurate but slower networks, use l=2 or higher, smaller number of features is faster
+# irreps_edge_sh: 0e + 1o                                                           # irreps of the spherical harmonics used for edges. If a single integer, indicates the full SH up to L_max=that_integer
+# conv_to_output_hidden_irreps_out: 16x0e                                           # irreps used in hidden layer of output block
+
+
 nonlinearity_type: gate                                                           # may be 'gate' or 'norm', 'gate' is recommended
 resnet: false                                                                     # set true to make interaction block a resnet-style update
+                                                                                  # the resnet update will only be applied when the input and output irreps of the layer are the same
 
 # scalar nonlinearities to use — available options are silu, ssp (shifted softplus), tanh, and abs.
 # Different nonlinearities are specified for e (even) and o (odd) parity;
@@ -48,7 +59,6 @@ invariant_layers: 2
 invariant_neurons: 64                                                             # number of hidden neurons in radial function, smaller is faster
 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
 
 # to specify different parameters for each convolutional layer, try examples below
 # layer1_use_sc: true                                                             # use "layer{i}_" prefix to specify parameters for only one of the layer,
@@ -121,8 +131,13 @@ chemical_symbols:
 # logging
 wandb: false                                                                       # we recommend using wandb for logging, we'll turn it off here as it's optional
 wandb_project: toluene-example                                                     # project name used in wandb
-wandb_resume: true                                                                 # if true and restart is true, wandb run data will be restarted and updated.
-                                                                                   # if false, a new wandb run will be generated
+wandb_watch: false
+# see https://docs.wandb.ai/ref/python/watch
+# wandb_watch_kwargs:
+#   log: all
+#   log_freq: 1
+#   log_graph: true
+
 verbose: info                                                                      # the same as python logging, e.g. warning, info, debug, error. case insensitive
 log_batch_freq: 1                                                                  # batch frequency, how often to print training errors withinin the same epoch
 log_epoch_freq: 1                                                                  # epoch frequency, how often to print and save the model

configs/minimal.yaml

@@ -1,3 +1,9 @@
+# !! PLEASE NOTE: `minimal.yaml` is meant as a _minimal_ example of a tiny, fast
+#                 training that can be used to verify your nequip install,
+#                 the syntax of your configuration edits, etc.
+#                 These are NOT recommended hyperparameters for real applications!
+#                 Please see `example.yaml` for a reasonable starting point.
+
 # general
 root: results/aspirin
 run_name: minimal
@@ -7,9 +13,9 @@ dataset_seed: 456
 # network
 num_basis: 8
 r_max: 4.0
-irreps_edge_sh: 0e + 1o
-conv_to_output_hidden_irreps_out: 16x0e
-feature_irreps_hidden: 16x0o + 16x0e + 16x1o + 16x1e + 16x2o + 16x2e
+l_max: 2
+parity: true
+num_features: 16
 
 # data set
 # the keys used need to be stated at least once in key_mapping, npz_fixed_field_keys or npz_keys

configs/minimal_eng.yaml

@@ -1,18 +1,31 @@
+# !! PLEASE NOTE: `minimal_eng.yaml` is meant as a _minimal_ example of a tiny, fast
+#                 training that can be used to verify your nequip install,
+#                 the syntax of your configuration edits, etc.
+#                 These are NOT recommended hyperparameters for real applications!
+#                 Please see `example.yaml` for a reasonable starting point.
+
 # general
 root: results/aspirin
 run_name: minimal_eng
-seed: 0
+seed: 123
+dataset_seed: 456
 
 # network
+# The default is to build a model with forces, so we need to specify
+# `model_builders` to get one without forces. This list is the default,
+# except without the `ForceOutput` builder that makes a force+energy
+# model out of an energy model:
 model_builders:
-  - EnergyModel
-  - PerSpeciesRescale
-  - RescaleEnergyEtc
+  - SimpleIrrepsConfig  # make configuration easier
+  - EnergyModel         # the core nequip model
+  - PerSpeciesRescale   # per-species/per-atom shift & scaling
+  - RescaleEnergyEtc    # global scaling
+# options for the model:
 num_basis: 8
 r_max: 4.0
-irreps_edge_sh: 0e + 1o
-conv_to_output_hidden_irreps_out: 16x0o + 16x0e + 16x1o + 16x1e + 16x2o + 16x2e
-feature_irreps_hidden: 16x0o + 16x0e 
+l_max: 2
+parity: true
+num_features: 16
 
 # data set
 # the keys used need to be stated at least once in key_mapping, npz_fixed_field_keys or npz_keys

docs/options/model.rst

@@ -99,7 +99,7 @@ invariant_neurons
 resnet
 ^^^^^^
     | Type: bool
-    | Default: ``True``
+    | Default: ``False``
 
 nonlinearity_type
 ^^^^^^^^^^^^^^^^^

nequip/_version.py

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