Files
lammps/src/ML-IAP/mliap_model_nn.cpp
2024-01-18 23:59:33 -05:00

394 lines
12 KiB
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.
------------------------------------------------------------------------- */
/* ----------------------------------------------------------------------
Contributing author: Pedro Antonio Santos Flórez (UNLV)
------------------------------------------------------------------------- */
#include "mliap_model_nn.h"
#include "mliap_data.h"
#include "comm.h"
#include "error.h"
#include "memory.h"
#include "tokenizer.h"
#include <cstring>
using namespace LAMMPS_NS;
static constexpr int MAXLINE = 1024;
/* ---------------------------------------------------------------------- */
MLIAPModelNN::MLIAPModelNN(LAMMPS *_lmp, char *coefffilename) : MLIAPModel(_lmp, coefffilename)
{
nnodes = nullptr;
activation = nullptr;
scale = nullptr;
if (coefffilename) MLIAPModelNN::read_coeffs(coefffilename);
nonlinearflag = 1;
}
/* ---------------------------------------------------------------------- */
MLIAPModelNN::~MLIAPModelNN()
{
memory->destroy(nnodes);
memory->destroy(activation);
memory->destroy(scale);
}
/* ----------------------------------------------------------------------
get number of parameters
---------------------------------------------------------------------- */
int MLIAPModelNN::get_nparams()
{
if (nparams == 0)
if (ndescriptors == 0) error->all(FLERR, "ndescriptors not defined");
return nparams;
}
void MLIAPModelNN::read_coeffs(char *coefffilename)
{
// open coefficient file on proc 0
FILE *fpcoeff;
if (comm->me == 0) {
fpcoeff = utils::open_potential(coefffilename, lmp, nullptr);
if (fpcoeff == nullptr)
error->one(FLERR, "Cannot open MLIAPModel coeff file {}: {}", coefffilename,
utils::getsyserror());
}
char line[MAXLINE], *ptr;
int n, eof = 0, nwords = 0;
while (nwords == 0) {
if (comm->me == 0) {
ptr = fgets(line, MAXLINE, fpcoeff);
if (ptr == nullptr) {
eof = 1;
fclose(fpcoeff);
} else
n = strlen(line) + 1;
}
MPI_Bcast(&eof, 1, MPI_INT, 0, world);
if (eof) break;
MPI_Bcast(&n, 1, MPI_INT, 0, world);
MPI_Bcast(line, n, MPI_CHAR, 0, world);
// strip comment, skip line if blank
if ((ptr = strchr(line, '#'))) *ptr = '\0';
nwords = utils::count_words(line);
}
if (nwords != 2) error->all(FLERR, "Incorrect format in MLIAPModel coefficient file");
// words = ptrs to all words in line
// strip single and double quotes from words
try {
ValueTokenizer coeffs(line);
nelements = coeffs.next_int();
nparams = coeffs.next_int();
} catch (TokenizerException &e) {
error->all(FLERR, "Incorrect format in MLIAPModel coefficient file: {}", e.what());
}
// set up coeff lists
memory->destroy(coeffelem);
memory->create(coeffelem, nelements, nparams, "mliap_snap_model:coeffelem");
int stats = 0;
int ielem = 0;
int l = 0;
while (true) {
if (comm->me == 0) {
ptr = fgets(line, MAXLINE, fpcoeff);
if (ptr == nullptr) {
eof = 1;
fclose(fpcoeff);
} else
n = strlen(line) + 1;
}
MPI_Bcast(&eof, 1, MPI_INT, 0, world);
if (eof) break;
MPI_Bcast(&n, 1, MPI_INT, 0, world);
MPI_Bcast(line, n, MPI_CHAR, 0, world);
// strip comment, skip line if blank
if ((ptr = strchr(line, '#'))) *ptr = '\0';
nwords = utils::trim_and_count_words(line);
if (nwords == 0) continue;
ValueTokenizer values(line, "\"' \t\n\t\f");
if (stats == 0) { // Header NET
auto tstr = values.next_string();
if (tstr.substr(0, 3) != "NET")
error->all(FLERR, "Incorrect format in MLIAPModel coefficient file");
ndescriptors = values.next_int();
nlayers = values.next_int();
memory->create(activation, nlayers, "mliap_model:activation");
memory->create(nnodes, nlayers, "mliap_model:nnodes");
memory->create(scale, nelements, 2, ndescriptors, "mliap_model:scale");
for (int ilayer = 0; ilayer < nlayers; ilayer++) {
tstr = values.next_string();
nnodes[ilayer] = values.next_int();
if (tstr == "linear")
activation[ilayer] = 0;
else if (tstr == "sigmoid")
activation[ilayer] = 1;
else if (tstr == "tanh")
activation[ilayer] = 2;
else if (tstr == "relu")
activation[ilayer] = 3;
else
activation[ilayer] = 4;
}
stats = 1;
} else if (stats == 1) {
scale[ielem][0][l] = values.next_double();
for (int icoeff = 1; icoeff < nwords; icoeff++) {
scale[ielem][0][l + icoeff] = values.next_double();
}
l += nwords;
if (l == ndescriptors) {
stats = 2;
l = 0;
}
} else if (stats == 2) {
scale[ielem][1][l] = values.next_double();
for (int icoeff = 1; icoeff < nwords; icoeff++) {
scale[ielem][1][l + icoeff] = values.next_double();
}
l += nwords;
if (l == ndescriptors) {
stats = 3;
l = 0;
}
// set up coeff lists
} else if (stats == 3) {
if (nwords > 30) error->all(FLERR, "Incorrect format in MLIAPModel coefficient file");
coeffelem[ielem][l] = values.next_double();
for (int icoeff = 1; icoeff < nwords; icoeff++) {
coeffelem[ielem][l + icoeff] = values.next_double();
}
l += nwords;
if (l == nparams) {
stats = 1;
l = 0;
ielem++;
}
}
}
}
/* ----------------------------------------------------------------------
Calculate model gradients w.r.t descriptors
for each atom beta_i = dE(B_i)/dB_i
---------------------------------------------------------------------- */
void MLIAPModelNN::compute_gradients(MLIAPData *data)
{
data->energy = 0.0;
for (int ii = 0; ii < data->nlistatoms; ii++) {
const int ielem = data->ielems[ii];
const int nl = nlayers;
double *coeffi = coeffelem[ielem];
double **scalei = scale[ielem];
double **nodes, **dnodes, **bnodes;
nodes = new double *[nl];
dnodes = new double *[nl];
bnodes = new double *[nl];
for (int l = 0; l < nl; ++l) {
nodes[l] = new double[nnodes[l]];
dnodes[l] = new double[nnodes[l]];
bnodes[l] = new double[nnodes[l]];
}
// forwardprop
// input - hidden1
for (int n = 0; n < nnodes[0]; n++) {
nodes[0][n] = 0;
for (int icoeff = 0; icoeff < data->ndescriptors; icoeff++) {
nodes[0][n] += coeffi[n * ((data->ndescriptors) + 1) + icoeff + 1] *
(data->descriptors[ii][icoeff] - scalei[0][icoeff]) / scalei[1][icoeff];
}
if (activation[0] == 1) {
nodes[0][n] = sigm(nodes[0][n] + coeffi[n * ((data->ndescriptors) + 1)], dnodes[0][n]);
} else if (activation[0] == 2) {
nodes[0][n] = tanh(nodes[0][n] + coeffi[n * ((data->ndescriptors) + 1)], dnodes[0][n]);
} else if (activation[0] == 3) {
nodes[0][n] = relu(nodes[0][n] + coeffi[n * ((data->ndescriptors) + 1)], dnodes[0][n]);
} else {
nodes[0][n] += coeffi[n * ((data->ndescriptors) + 1)];
dnodes[0][n] = 1;
}
}
// hidden~output
int k = 0;
if (nl > 1) {
k += ((data->ndescriptors) + 1) * nnodes[0];
for (int l = 1; l < nl; l++) {
for (int n = 0; n < nnodes[l]; n++) {
nodes[l][n] = 0;
for (int j = 0; j < nnodes[l - 1]; j++) {
nodes[l][n] += coeffi[k + n * (nnodes[l - 1] + 1) + j + 1] * nodes[l - 1][j];
}
if (activation[l] == 1) {
nodes[l][n] = sigm(nodes[l][n] + coeffi[k + n * (nnodes[l - 1] + 1)], dnodes[l][n]);
} else if (activation[l] == 2) {
nodes[l][n] = tanh(nodes[l][n] + coeffi[k + n * (nnodes[l - 1] + 1)], dnodes[l][n]);
} else if (activation[l] == 3) {
nodes[l][n] = relu(nodes[l][n] + coeffi[k + n * (nnodes[l - 1] + 1)], dnodes[l][n]);
} else {
nodes[l][n] += coeffi[k + n * (nnodes[l - 1] + 1)];
dnodes[l][n] = 1;
}
}
k += (nnodes[l - 1] + 1) * nnodes[l];
}
}
// backwardprop
// output layer dnode initialized to 1.
for (int n = 0; n < nnodes[nl - 1]; n++) {
if (activation[nl - 1] == 0) {
bnodes[nl - 1][n] = 1;
} else {
bnodes[nl - 1][n] = dnodes[nl - 1][n];
}
}
if (nl > 1) {
for (int l = nl - 1; l > 0; l--) {
k -= (nnodes[l - 1] + 1) * nnodes[l];
for (int n = 0; n < nnodes[l - 1]; n++) {
bnodes[l - 1][n] = 0;
for (int j = 0; j < nnodes[l]; j++) {
bnodes[l - 1][n] += coeffi[k + j * (nnodes[l - 1] + 1) + n + 1] * bnodes[l][j];
}
if (activation[l - 1] >= 1) { bnodes[l - 1][n] *= dnodes[l - 1][n]; }
}
}
}
for (int icoeff = 0; icoeff < data->ndescriptors; icoeff++) {
data->betas[ii][icoeff] = 0;
for (int j = 0; j < nnodes[0]; j++) {
data->betas[ii][icoeff] +=
coeffi[j * ((data->ndescriptors) + 1) + icoeff + 1] * bnodes[0][j];
}
data->betas[ii][icoeff] = data->betas[ii][icoeff] / scalei[1][icoeff];
}
if (data->eflag) {
// energy of atom I (E_i)
double etmp = nodes[nl - 1][0];
data->energy += etmp;
data->eatoms[ii] = etmp;
}
// Deleting the variables
for (int n = 0; n < nl; n++) {
delete[] nodes[n];
delete[] dnodes[n];
delete[] bnodes[n];
}
delete[] nodes;
delete[] dnodes;
delete[] bnodes;
}
}
/* ----------------------------------------------------------------------
Calculate model double gradients w.r.t descriptors and parameters
for each atom energy gamma_lk = d2E(B)/dB_k/dsigma_l,
where sigma_l is a parameter, B_k a descriptor,
and atom subscript i is omitted
gamma is in CSR format:
nnz = number of non-zero values
gamma_row_index[inz] = l indices, 0 <= l < nparams
gamma_col_indexiinz] = k indices, 0 <= k < ndescriptors
gamma[i][inz] = non-zero values, 0 <= inz < nnz
egradient is derivative of energy w.r.t. parameters
---------------------------------------------------------------------- */
void MLIAPModelNN::compute_gradgrads(class MLIAPData * /*data*/)
{
error->all(FLERR, "compute_gradgrads not implemented");
}
/* ----------------------------------------------------------------------
calculate gradients of forces w.r.t. parameters
egradient is derivative of energy w.r.t. parameters
---------------------------------------------------------------------- */
void MLIAPModelNN::compute_force_gradients(class MLIAPData * /*data*/)
{
error->all(FLERR, "compute_force_gradients not implemented");
}
/* ----------------------------------------------------------------------
count the number of non-zero entries in gamma matrix
---------------------------------------------------------------------- */
int MLIAPModelNN::get_gamma_nnz(class MLIAPData * /*data*/)
{
// todo: get_gamma_nnz
return 0;
}
double MLIAPModelNN::memory_usage()
{
double bytes = 0;
bytes += (double) nelements * nparams * sizeof(double); // coeffelem
bytes += (double) nelements * 2 * ndescriptors * sizeof(double); // scale
bytes += (int) nlayers * sizeof(int); // nnodes
bytes += (int) nlayers * sizeof(int); // activation
return bytes;
}