Add stress, precision support, compute

.github/workflows/tests.yml

@@ -47,7 +47,7 @@ jobs:
       run: |
         mkdir lammps_dir/
         cd lammps_dir/
-        git clone -b stable_29Sep2021_update2 --depth 1 "https://github.com/lammps/lammps"
+        git clone --depth 1 "https://github.com/lammps/lammps"
         cd ..
         ./patch_lammps.sh lammps_dir/lammps/
         cd lammps_dir/lammps/

.gitignore

@@ -32,8 +32,6 @@
 *.out
 *.app
 
-.vscode
-
 
 
 # ---------- Python .gigignores-----------

.vscode/settings.json

@@ -0,0 +1,7 @@
+{
+    "editor.formatOnSave": false,
+    "[python]": {
+        "editor.formatOnSave": true
+    },
+    "python.formatting.provider": "black"
+}

README.md

@@ -2,17 +2,20 @@
 
 This pair style allows you to use Allegro models from the [`allegro`](https://github.com/mir-group/allegro) package in LAMMPS simulations. Allegro is designed to enable parallelism, and so `pair_allegro` **supports MPI in LAMMPS**. It also supports OpenMP (better performance) or Kokkos (best performance) for accelerating the pair style.
 
-For more details on Allegro itself, background, and the LAMMPS pair style please see the [`allegro`](https://github.com/mir-group/allegro) package and our pre-print:
+For more details on Allegro itself, background, and the LAMMPS pair style please see the [`allegro`](https://github.com/mir-group/allegro) package and our paper:
 > *Learning Local Equivariant Representations for Large-Scale Atomistic Dynamics* <br/>
 > Albert Musaelian, Simon Batzner, Anders Johansson, Lixin Sun, Cameron J. Owen, Mordechai Kornbluth, Boris Kozinsky <br/>
-> https://arxiv.org/abs/2204.05249 <br/>
-> https://doi.org/10.48550/arXiv.2204.05249
+> https://www.nature.com/articles/s41467-023-36329-y <br/>
+and
+> *Scaling the leading accuracy of deep equivariant models to biomolecular simulations of realistic size* <br/>
+> Albert Musaelian, Anders Johansson, Simon Batzner, Boris Kozinsky <br/>
+> https://arxiv.org/abs/2304.10061 <br/>
 
 `pair_allegro` authors: **Anders Johansson**, Albert Musaelian.
 
 ## Pre-requisites
 
-* PyTorch or LibTorch >= 1.11.0
+* PyTorch or LibTorch >= 1.11.0;  please note that at present we **only recommend 1.11** on CUDA systems.
 
 ## Usage in LAMMPS
 
@@ -23,36 +26,53 @@ pair_coeff	* * deployed.pth <Allegro type name for LAMMPS type 1> <Allegro type
 where `deployed.pth` is the filename of your trained, **deployed** model.
 
 The names after the model path `deployed.pth` indicate, in order, the names of the Allegro model's atom types to use for LAMMPS atom types 1, 2, and so on. The number of names given must be equal to the number of atom types in the LAMMPS configuration (not the Allegro model!).
-The given names must be consistent with the names specified in the Allegro training YAML in `chemical_symbol_to_type` or `type_names`.
+The given names must be consistent with the names specified in the Allegro training YAML in `chemical_symbol_to_type` or `type_names`. Typically, this will be the chemical symbol for each LAMMPS type.
 
 To run with Kokkos, please see the [LAMMPS Kokkos documentation](https://docs.lammps.org/Speed_kokkos.html#running-on-gpus). Example:
 ```bash
 mpirun -np 8 lmp -sf kk -k on g 4 -pk kokkos newton on neigh full -in in.script
 ```
 to run on 2 nodes with 4 GPUs each.
 
+### Compute
+We provide an experimental "compute" that allows you to extract custom quantities from Allegro models, such as [polarization](https://arxiv.org/abs/2403.17207). You can extract either global or per-atom properties with syntax along the lines of
+```
+compute polarization all allegro polarization 3
+compute polarizability all allegro polarizability 9
+compute borncharges all allegro/atom born_charge 9 1
+```
+
+The name after `allegro[/atom]` is attempted extracted from the dictionary that the Allegro model returns. The following number is the number of elements after flattening the output. In the examples above, polarization is a 3-element global vector, while polarizability and Born charges are global and per-atom 3x3 matrices, respectively. For per-atom quantities, the second number is a flag indicating whether the properties should be reverse-communicated "Newton-style" like forces, which will depend on your property and the specifics of your implementation.
+
+*Note: For extracting multiple quantities, simply use multiple commands. The properties will be extracted from the same dictionary, without any recomputation.*
+
+*Note: The quantities will be attempted extracted at every timestep. In the future, we may add support for passing a flag to the model indicating that the "custom" output should be computed.*
+
+*Note: The group flag shoul generally be `all`.*
+
 ## Building LAMMPS with this pair style
 
 ### Download LAMMPS
 ```bash
-git clone -b stable_29Sep2021_update2 --depth 1 git@github.com:lammps/lammps
+git clone --depth 1 https://github.com/lammps/lammps
 ```
 or your preferred method.
 (`--depth 1` prevents the entire history of the LAMMPS repository from being downloaded.)
 
 ### Download this repository
 ```bash
-git clone git@github.com:mir-group/pair_allegro
+git clone https://github.com/mir-group/pair_allegro
 ```
 or by downloading a ZIP of the source.
 
 ### Patch LAMMPS
-#### Automatically
 From the `pair_allegro` directory, run:
 ```bash
 ./patch_lammps.sh /path/to/lammps/
 ```
 
+### Libraries
+
 #### Libtorch
 If you have PyTorch installed and are **NOT** using Kokkos:
 ```bash
@@ -61,7 +81,7 @@ mkdir build
 cd build
 cmake ../cmake -DCMAKE_PREFIX_PATH=`python -c 'import torch;print(torch.utils.cmake_prefix_path)'`
 ```
-If you don't have PyTorch installed **OR** are using Kokkos, you need to download LibTorch from the [PyTorch download page](https://pytorch.org/get-started/locally/). **Ensure you download the cxx11 ABI version.** Unzip the downloaded file, then configure LAMMPS:
+If you don't have PyTorch installed **OR** are using Kokkos, you need to download LibTorch from the [PyTorch download page](https://pytorch.org/get-started/locally/). **Ensure you download the cxx11 ABI version if using Kokkos.** Unzip the downloaded file, then configure LAMMPS:
 ```bash
 cd lammps
 mkdir build
@@ -86,15 +106,15 @@ CMake will look for CUDA and cuDNN. You may have to explicitly provide the path
 Note that the CUDA that comes with PyTorch when installed with `conda` (the `cudatoolkit` package) is usually insufficient (see [here](https://github.com/pytorch/extension-cpp/issues/26), for example) and you may have to install full CUDA seperately. A minor version mismatch between the available full CUDA version and the version of `cudatoolkit` is usually *not* a problem, as long as the system CUDA is equal or newer. (For example, PyTorch's requested `cudatoolkit==11.3` with a system CUDA of 11.4 works, but a system CUDA 11.1 will likely fail.) cuDNN is also required by PyTorch.
 
 #### With OpenMP (optional, better performance)
-`pair_allegro` supports the use of OpenMP to accelerate certain parts of the pair style.
+`pair_allegro` supports the use of OpenMP to accelerate certain parts of the pair style, by setting `OMP_NUM_THREADS` and using the [LAMMPS OpenMP package](https://docs.lammps.org/Speed_omp.html).
 
 #### With Kokkos (GPU, optional, best performance)
 `pair_allegro` supports the use of Kokkos to accelerate certain parts of the pair style on the GPU to avoid host-GPU transfers.
 `pair_allegro` supports two setups for Kokkos: pair_style and model both on CPU, or both on GPU. Please ensure you build LAMMPS with the appropriate Kokkos backends enabled for your usecase. For example, to use CUDA GPUs, add:
 ```
 -DPKG_KOKKOS=ON -DKokkos_ENABLE_CUDA=ON
 ```
-to your `cmake` command.
+to your `cmake` command. See the [LAMMPS documentation](https://docs.lammps.org/Speed_kokkos.html) for more build options and how to correctly run LAMMPS with Kokkos.
 
 ### Building LAMMPS
 ```bash
@@ -106,14 +126,24 @@ This gives `lammps/build/lmp`, which can be run as usual with `/path/to/lmp -in
 
 1. Q: My simulation is immediately or bizzarely unstable
 
-   A: Please ensure that your mapping from LAMMPS atom types to NequIP atom types, specified in the `pair_coeff` line, is correct.
+   A: Please ensure that your mapping from LAMMPS atom types to NequIP atom types, specified in the `pair_coeff` line, is correct, and that the units are consistent between your training data and your LAMMPS simulation.
 2. Q: I get the following error:
    ```
     instance of 'c10::Error'
         what():  PytorchStreamReader failed locating file constants.pkl: file not found
    ```
 
    A: Make sure you remembered to deploy (compile) your model using `nequip-deploy`, and that the path to the model given with `pair_coeff` points to a deployed model `.pth` file, **not** a file containing only weights like `best_model.pth`.
-3. Q: The output pressures and stresses seem wrong / my NPT simulation is broken
+3. Q: I get the following error:
+   ```
+    instance of 'c10::Error'
+        what():  isTuple()INTERNAL ASSERT FAILED
+   ```
+
+   A: We've seen this error occur when you try to load a TorchScript model deployed from PyTorch>1.11 in LAMMPS built against 1.11. Try redeploying your model (retraining not necessary) in a PyTorch 1.11 install.
+4. Q: I get the following error:
+    ```
+    Exception: Argument passed to at() was not in the map
+    ```
 
-    A: NPT/stress support in LAMMPS for `pair_allegro` is in-progress and not yet available.
+    A: We now require models to have been trained with stress support, which is achieved by replacing `ForceOutput` with `StressForceOutput` in the training configuration. Note that you do not need to train on stress (though it may improve your potential, assuming your stress data is correct and converged). If you desperately wish to keep using a model without stress output, you can remove lines that look like [these](https://github.com/mir-group/pair_allegro/blob/99036043e74376ac52993b5323f193dee3f4f401/pair_allegro_kokkos.cpp#L332-L343) in your version of `pair_allegro[_kokkos].cpp`.

compute_allegro.cpp

@@ -0,0 +1,179 @@
+/* ----------------------------------------------------------------------
+   LAMMPS - Large-scale Atomic/Molecular Massively Parallel Simulator
+   https://lammps.sandia.gov/, Sandia National Laboratories
+   Steve Plimpton, sjplimp@sandia.gov
+
+   Copyright (2003) Sandia Corporation.  Under the terms of Contract
+   DE-AC04-94AL85000 with Sandia Corporation, the U.S. Government retains
+   certain rights in this software.  This software is distributed under
+   the GNU General Public License.
+
+   See the README file in the top-level LAMMPS directory.
+------------------------------------------------------------------------- */
+
+/* ----------------------------------------------------------------------
+   Contributing author: Anders Johansson (Harvard)
+------------------------------------------------------------------------- */
+
+#include "compute_allegro.h"
+#include "atom.h"
+#include "comm.h"
+#include "error.h"
+#include "force.h"
+#include "memory.h"
+#include "pair_allegro.h"
+#include "update.h"
+
+#include <cassert>
+#include <cmath>
+#include <cstring>
+#include <iostream>
+#include <numeric>
+#include <sstream>
+#include <string>
+#include <torch/script.h>
+#include <torch/torch.h>
+
+using namespace LAMMPS_NS;
+
+template<int peratom>
+ComputeAllegro<peratom>::ComputeAllegro(LAMMPS *lmp, int narg, char **arg) : Compute(lmp, narg, arg)
+{
+
+  if constexpr (!peratom) {
+    // compute 1 all allegro quantity length
+    if (narg != 5) error->all(FLERR, "Incorrect args for compute allegro");
+  } else {
+    // compute 1 all allegro/atom quantity length newton(1/0)
+    if (narg != 6) error->all(FLERR, "Incorrect args for compute allegro/atom");
+  }
+
+  if (strcmp(arg[1], "all") != 0)
+    error->all(FLERR, "compute allegro can only operate on group 'all'");
+
+  quantity = arg[3];
+  if constexpr (peratom) {
+    peratom_flag = 1;
+    nperatom = std::atoi(arg[4]);
+    newton = std::atoi(arg[5]);
+    if (newton) comm_reverse = nperatom;
+    size_peratom_cols = nperatom==1 ? 0 : nperatom;
+    nmax = -12;
+    if (comm->me == 0)
+      error->message(FLERR, "compute allegro/atom will evaluate the quantity {} of length {} with newton {}", quantity,
+                     size_peratom_cols, newton);
+  } else {
+    vector_flag = 1;
+    size_vector = std::atoi(arg[4]);
+    if (size_vector <= 0) error->all(FLERR, "Incorrect vector length!");
+    memory->create(vector, size_vector, "ComputeAllegro:vector");
+    if (comm->me == 0)
+      error->message(FLERR, "compute allegro will evaluate the quantity {} of length {}", quantity,
+                     size_vector);
+  }
+
+  if (force->pair == nullptr) {
+    error->all(FLERR, "no pair style; compute allegro must be defined after pair style");
+  }
+
+  ((PairAllegro<lowhigh> *) force->pair)->add_custom_output(quantity);
+}
+
+template<int peratom>
+void ComputeAllegro<peratom>::init()
+{
+  ;
+}
+
+template<int peratom>
+ComputeAllegro<peratom>::~ComputeAllegro()
+{
+  if (copymode) return;
+  if constexpr (peratom) {
+    memory->destroy(vector_atom);
+  } else {
+    memory->destroy(vector);
+  }
+}
+
+template<int peratom>
+void ComputeAllegro<peratom>::compute_vector()
+{
+  invoked_vector = update->ntimestep;
+
+  const torch::Tensor &quantity_tensor =
+      ((PairAllegro<lowhigh> *) force->pair)->custom_output.at(quantity).cpu().ravel();
+
+  auto quantity = quantity_tensor.data_ptr<double>();
+
+  if (quantity_tensor.size(0) != size_vector) {
+    error->one(FLERR, "size {} of quantity tensor {} does not match expected {} on rank {}",
+               quantity_tensor.size(0), this->quantity, size_vector, comm->me);
+  }
+
+  for (int i = 0; i < size_vector; i++) { vector[i] = quantity[i]; }
+
+  MPI_Allreduce(MPI_IN_PLACE, vector, size_vector, MPI_DOUBLE, MPI_SUM, world);
+}
+
+template<int peratom>
+void ComputeAllegro<peratom>::compute_peratom()
+{
+  invoked_peratom = update->ntimestep;
+
+  if (atom->nmax > nmax) {
+    nmax = atom->nmax;
+    memory->destroy(array_atom);
+    memory->create(array_atom, nmax, nperatom, "allegro/atom:array");
+    if (nperatom==1) vector_atom = &array_atom[0][0];
+  }
+
+  const torch::Tensor &quantity_tensor =
+      ((PairAllegro<lowhigh> *) force->pair)->custom_output.at(quantity).cpu().contiguous().reshape({-1,nperatom});
+
+  auto quantity = quantity_tensor.accessor<double,2>();
+  quantityptr = quantity_tensor.data_ptr<double>();
+
+  int nlocal = atom->nlocal;
+  for (int i = 0; i < nlocal; i++) {
+    for (int j = 0; j < nperatom; j++) {
+      array_atom[i][j] = quantity[i][j];
+    }
+  }
+  if (newton) comm->reverse_comm(this);
+}
+
+template<int peratom>
+int ComputeAllegro<peratom>::pack_reverse_comm(int n, int first, double *buf)
+{
+  int i, m, last;
+
+  m = 0;
+  last = first + n;
+  for (i = first; i < last; i++) {
+    for (int j = 0; j < nperatom; j++) {
+      buf[m++] = quantityptr[i*nperatom + j];
+    }
+  }
+  return m;
+}
+
+template<int peratom>
+void ComputeAllegro<peratom>::unpack_reverse_comm(int n, int *list, double *buf)
+{
+  int i, j, m;
+
+  m = 0;
+  for (i = 0; i < n; i++) {
+    j = list[i];
+    for (int k = 0; k < nperatom; k++) {
+      array_atom[j][k] += buf[m++];
+    }
+  }
+}
+
+
+namespace LAMMPS_NS {
+  template class ComputeAllegro<0>;
+  template class ComputeAllegro<1>;
+}