Add files via upload

mace-bench/3rdparty/SevenNet/sevenn/pair_e3gnn/pair_d3.h

+/* -*- c++ -*- ----------------------------------------------------------
+   LAMMPS - Large-scale Atomic/Molecular Massively Parallel Simulator
+   https://www.lammps.org/, Sandia National Laboratories
+   LAMMPS development team: developers@lammps.org
+
+   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.
+------------------------------------------------------------------------- */
+
+#ifdef PAIR_CLASS
+PairStyle(d3, PairD3)
+
+#else
+
+#ifndef LMP_PAIR_D3
+#define LMP_PAIR_D3
+
+#include <cmath>
+#include <string>
+#include <vector>
+#include <algorithm>
+#include <map>
+#include <unordered_map>
+#include <cuda_runtime.h>
+
+#include "pair.h"
+#include "utils.h"
+#include "atom.h"
+#include "domain.h"
+#include "error.h"
+#include "comm.h"
+#include "neighbor.h"
+#include "neigh_list.h"
+#include "math_extra.h"
+
+#include "pair_d3_pars.h"
+
+// Removed dependencies to STL
+// #include <stdlib.h> -> no more C style functions
+// #define _USE_MATH_DEFINES -> no predefined constants
+
+// Removed dependencies to LAMMPS
+// #include "potential_file_reader.h" -> removed, PotentialFileReader
+// #include "memory.h"     -> already no dependency for CUDA version
+
+namespace LAMMPS_NS {
+    class PairD3 : public Pair {
+    public:
+        PairD3(class LAMMPS*);
+        ~PairD3() override;
+
+        void compute(int, int) override;
+        void settings(int, char**) override;
+        void coeff(int, char**) override;
+        double init_one(int i, int j) override;
+        void init_style() override;
+
+        void write_restart(FILE*) override;
+        void read_restart(FILE*) override;
+        void write_restart_settings(FILE*) override;
+        void read_restart_settings(FILE*) override;
+
+    protected:
+        virtual void allocate();
+
+        /* ------- Read parameters ------- */
+        int find_atomic_number(std::string&);
+        int is_int_in_array(int*, int, int);
+        void read_r0ab(int*, int);
+        void get_limit_in_pars_array(int&, int&, int&, int&);
+        void read_c6ab(int*, int);
+
+        void setfuncpar_zero();
+        void setfuncpar_bj();
+        void setfuncpar_zerom();
+        void setfuncpar_bjm();
+        void setfuncpar();
+        /* ------- Read parameters ------- */
+
+        /* ------- Lattice information ------- */
+        void set_lattice_repetition_criteria(float, int*);
+        void set_lattice_vectors();
+        /* ------- Lattice information ------- */
+
+        /* ------- Initialize & Precalculate ------- */
+        void load_atom_info();
+        void precalculate_tau_array();
+        /* ------- Initialize & Precalculate ------- */
+
+        /* ------- Reallocate (when number of atoms changed) ------- */
+        void reallocate_arrays();
+        /* ------- Reallocate (when number of atoms changed) ------- */
+
+        /* ------- Coordination number ------- */
+        void get_coordination_number();
+        void get_dC6_dCNij();
+        /* ------- Coordination number ------- */
+
+        /* ------- Main workers ------- */
+        void get_forces_without_dC6_zero();
+        void get_forces_without_dC6_bj();
+        void get_forces_without_dC6_zerom();
+        void get_forces_without_dC6_bjm();
+        void get_forces_without_dC6();
+        void get_forces_with_dC6();
+        void update(int, int);
+        /* ------- Main workers ------- */
+
+        /*--------- Constants ---------*/
+        static constexpr int MAX_ELEM = 94;          // maximum of the element number
+        static constexpr int MAXC = 5;               // maximum coordination number references per element
+
+        static constexpr double AU_TO_ANG = 0.52917726; // conversion factors (atomic unit --> angstrom)
+        static constexpr double AU_TO_EV = 27.21138505; // conversion factors (atomic unit --> eV)
+
+        static constexpr float K1 = 16.0;              // global ad hoc parameters
+        static constexpr float K3 = -4.0;              // global ad hoc parameters
+        /*--------- Constants ---------*/
+
+        /*--------- Parameters to read ---------*/
+        int damping;
+        std::string functional;
+        float* r2r4 = nullptr;             // scale r4/r2 values of the atoms by sqrt(Z)
+        float* rcov = nullptr;             // covalent radii
+        int* mxc = nullptr;                // How large the grid for c6 interpolation
+        float** r0ab = nullptr;            // cut-off radii for all element pairs
+        float***** c6ab = nullptr;         // C6 for all element pairs
+        float rthr;                        // R^2 distance to cutoff for C calculation
+        float cnthr;                       // R^2 distance to cutoff for CN_calculation
+        float s6, s8, s18, rs6, rs8, rs18, alp, alp6, alp8, a1, a2; // parameters for D3
+        /*--------- Parameters to read ---------*/
+
+        /*--------- Lattice related values ---------*/
+        double* lat_v_1 = nullptr;           // lattice coordination vector
+        double* lat_v_2 = nullptr;           // lattice coordination vector
+        double* lat_v_3 = nullptr;           // lattice coordination vector
+        int* rep_vdw = nullptr;              // repetition of cell for calculating D3
+        int* rep_cn = nullptr;               // repetition of cell for calculating
+        double** sigma = nullptr;            // virial pressure on cell
+        /*--------- Lattice related values ---------*/
+
+        /*--------- Per-atom values/arrays ---------*/
+        double* cn = nullptr;               // Coordination numbers
+        float** x = nullptr;                // Positions
+        double** f = nullptr;               // Forces
+        double* dc6i = nullptr;             // dC6i(iat) saves dE_dsp/dCN(iat)
+        /*--------- Per-atom values/arrays ---------*/
+
+        /*--------- Per-pair values/arrays ---------*/
+        float* c6_ij_tot = nullptr;
+        float* dc6_iji_tot = nullptr;
+        float* dc6_ijj_tot = nullptr;
+        /*--------- Per-pair values/arrays ---------*/
+
+        /*---------- Global values ---------*/
+        int n_save;                         // to check whether the number of atoms has changed
+        float disp_total;                   // Dispersion energy
+        /*---------- Global values ---------*/
+
+        /*--------- For loop over tau (translation of cell) ---------*/
+        float**** tau_vdw = nullptr;
+        float**** tau_cn = nullptr;
+        int* tau_idx_vdw = nullptr;
+        int* tau_idx_cn = nullptr;
+        int tau_idx_vdw_total_size;
+        int tau_idx_cn_total_size;
+        /*--------- For loop over tau (translation of cell) ---------*/
+
+        /*--------- For cuda memory transfer (pointerized) ---------*/
+        int *atomtype;
+        double *disp;
+        /*--------- For cuda memory transfer (pointerized) ---------*/
+    };
+}
+
+#endif // LMP_PAIR_D3
+#endif // PAIR_CLASS

mace-bench/3rdparty/SevenNet/sevenn/pair_e3gnn/pair_e3gnn.h

+/* -*- c++ -*- ----------------------------------------------------------
+   LAMMPS - Large-scale Atomic/Molecular Massively Parallel Simulator
+http://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.
+------------------------------------------------------------------------- */
+
+#ifdef PAIR_CLASS
+PairStyle(e3gnn, PairE3GNN)
+
+#else
+
+#ifndef LMP_PAIR_E3GNN
+#define LMP_PAIR_E3GNN
+
+#include "pair.h"
+
+#include <torch/torch.h>
+
+namespace LAMMPS_NS {
+class PairE3GNN : public Pair {
+private:
+  double cutoff;
+  double cutoff_square;
+  torch::jit::Module model;
+  torch::Device device = torch::kCPU;
+  int nelements;
+  bool print_info = false;
+
+  int nedges_bound = -1;
+
+public:
+  PairE3GNN(class LAMMPS *);
+  ~PairE3GNN();
+  void compute(int, int);
+
+  void settings(int, char **);
+  // read Atom type string from input script & related coeff
+  void coeff(int, char **);
+  void allocate();
+
+  void init_style();
+  double init_one(int, int);
+};
+} // namespace LAMMPS_NS
+
+#endif
+#endif

mace-bench/3rdparty/SevenNet/sevenn/pair_e3gnn/pair_e3gnn_parallel.h

+/* -*- c++ -*- ----------------------------------------------------------
+   LAMMPS - Large-scale Atomic/Molecular Massively Parallel Simulator
+http://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.
+------------------------------------------------------------------------- */
+
+#ifdef PAIR_CLASS
+PairStyle(e3gnn/parallel, PairE3GNNParallel)
+
+#else
+
+#ifndef LMP_PAIR_E3GNN_PARALLEL
+#define LMP_PAIR_E3GNN_PARALLEL
+
+#include "pair.h"
+
+#include <torch/torch.h>
+#include <vector>
+
+namespace LAMMPS_NS {
+class PairE3GNNParallel : public Pair {
+private:
+  double cutoff;
+  double cutoff_square;
+  std::vector<torch::jit::Module> model_list;
+  torch::Device device = torch::kCPU;
+  torch::Device device_comm = torch::kCPU;
+  torch::Device get_cuda_device();
+  bool use_cuda_mpi;
+
+  // for communication
+  // Most of these variables for communication is temporary and valid for only
+  // one MD step.
+  int x_dim; // to determine per atom data size
+  int graph_size;
+  torch::Tensor x_comm; // x_local + x_ghost + x_comm_extra
+
+  void comm_preprocess();
+  bool comm_preprocess_done = false;
+
+  // temporary variables holds for each compute step
+  std::unordered_map<int, long> extra_graph_idx_map;
+  // To use scatter, store long instead of int
+  // array of vector
+  std::vector<long> comm_index_pack_forward[6];
+  std::vector<long> comm_index_unpack_forward[6];
+  std::vector<long> comm_index_unpack_reverse[6];
+
+  // its size is 6 and initialized at comm_preprocess()
+  torch::Tensor comm_index_pack_forward_tensor[6];
+  torch::Tensor comm_index_unpack_forward_tensor[6];
+  torch::Tensor comm_index_unpack_reverse_tensor[6];
+
+  // to use tag_to_graph_idx inside comm methods
+  int *tag_to_graph_idx_ptr = nullptr;
+
+  int sendproc[6];
+  int recvproc[6];
+
+public:
+  PairE3GNNParallel(class LAMMPS *);
+  ~PairE3GNNParallel();
+
+  // TODO: keep encapsulation..
+  void compute(int, int) override;
+  void settings(int, char **) override;
+  // read Atom type string from input script & related coeff
+  void coeff(int, char **) override;
+  void allocate();
+
+  void pack_forward_init(int n, int *list, int comm_phase);
+  void unpack_forward_init(int n, int first, int comm_phase);
+
+  int pack_forward_comm_gnn(float *buf, int comm_phase);
+  void unpack_forward_comm_gnn(float *buf, int comm_phase);
+  int pack_reverse_comm_gnn(float *buf, int comm_phase);
+  void unpack_reverse_comm_gnn(float *buf, int comm_phase);
+
+  void init_style() override;
+  double init_one(int, int) override;
+
+  int get_x_dim();
+  bool use_cuda_mpi_();
+  bool is_comm_preprocess_done();
+  void notify_proc_ids(const int *sendproc, const int *recvproc);
+
+  bool print_info = false;
+  int world_rank;
+};
+
+class DeviceBuffManager {
+private:
+  DeviceBuffManager() {}
+  DeviceBuffManager(const DeviceBuffManager &);
+  DeviceBuffManager &operator=(const DeviceBuffManager &);
+
+  float *buf_send_device = nullptr;
+  float *buf_recv_device = nullptr;
+  int send_buf_size = 0;
+  int recv_buf_size = 0;
+
+public:
+  static DeviceBuffManager &getInstance();
+  void get_buffer(int, int, float *&, float *&);
+
+  ~DeviceBuffManager();
+};
+} // namespace LAMMPS_NS
+
+#endif
+#endif