{

this->inp_ = &inp;

nep_potential = 0.0;

nep_virial.create(3, 3);

if (basecell.kind() == BaseCell::Kind::mdcell)

{

MDCell& mdcell = static_cast<MDCell&>(basecell);

#ifdef __NEP

const double cutoff = std::max(nep.paramb.rc_radial_max, nep.paramb.rc_angular_max)

* ModuleBase::ANGSTROM_AU;

mdcell.initialize_neighbors(cutoff);

initialize_type_map_(mdcell.type_labels());

#else

ModuleBase::WARNING_QUIT("ESolver_NEP", "Please recompile with -D__NEP");

#endif

return;

}

UnitCell& ucell = static_cast<UnitCell&>(basecell);

nep_force.create(ucell.nat, 3);

atype.resize(ucell.nat);

_e.resize(ucell.nat);

_f.resize(3 * ucell.nat);

_v.resize(9 * ucell.nat);

ModuleIO::CifParser::write(PARAM.globalv.global_out_dir + "STRU.cif",

ucell,

"# Generated by ABACUS ModuleIO::CifParser",

"data_?");

#ifdef __NEP

/// determine the type map from STRU to NEP model

type_map(ucell);

#endif

{

ModuleBase::TITLE("ESolver_NEP", "runner");

ModuleBase::timer::start("ESolver_NEP", "runner");

if (basecell.kind() == BaseCell::Kind::mdcell)

{

#ifndef __NEP

ModuleBase::WARNING_QUIT("ESolver_NEP", "Please recompile with -D__NEP");

#else

static_cast<void>(istep);

MDCell& mdcell = static_cast<MDCell&>(basecell);

if (!mdcell.has_neighbor_search())

{

mdcell.prepare_neighbors();

}

const int nowned_atoms = mdcell.nowned_atoms();

const int nghost = mdcell.nghost();

const int natom = nowned_atoms + nghost;

if (natom == 0)

{

ModuleBase::WARNING_QUIT("ESolver_NEP", "MDCell contains no atoms.");

}

const std::vector<LocalAtom>& owned_atoms = mdcell.owned_atoms();

const std::vector<LocalAtom>& ghost_atoms = mdcell.ghost_atoms();

std::vector<int> local_type(static_cast<std::size_t>(natom), 0);

std::vector<std::array<double, 3> > position(static_cast<std::size_t>(natom));

std::vector<std::array<double, 3> > force(static_cast<std::size_t>(natom));

std::vector<double*> position_ptrs(static_cast<std::size_t>(natom), NULL);

std::vector<double*> force_ptrs(static_cast<std::size_t>(natom), NULL);

for (int iat = 0; iat < natom; ++iat)

{

const LocalAtom& atom = iat < nowned_atoms ? owned_atoms[static_cast<std::size_t>(iat)]

: ghost_atoms[static_cast<std::size_t>(iat - nowned_atoms)];

if (atom.type < 0 || static_cast<std::size_t>(atom.type) >= md_type_to_nep_type_.size())

{

ModuleBase::WARNING_QUIT("ESolver_NEP", "MDCell atom type is outside the NEP type map.");

}

local_type[static_cast<std::size_t>(iat)] = atom.type;

position[static_cast<std::size_t>(iat)][0] = atom.cart.x * mdcell.lat0() * ModuleBase::BOHR_TO_A;

position[static_cast<std::size_t>(iat)][1] = atom.cart.y * mdcell.lat0() * ModuleBase::BOHR_TO_A;

position[static_cast<std::size_t>(iat)][2] = atom.cart.z * mdcell.lat0() * ModuleBase::BOHR_TO_A;

force[static_cast<std::size_t>(iat)].fill(0.0);

position_ptrs[static_cast<std::size_t>(iat)] = position[static_cast<std::size_t>(iat)].data();

force_ptrs[static_cast<std::size_t>(iat)] = force[static_cast<std::size_t>(iat)].data();

}

const NeighborList& neighbor_list = mdcell.neighbor_search().get_neighbor_list();

std::vector<int> ilist(static_cast<std::size_t>(nowned_atoms), 0);

std::vector<int> numneigh(static_cast<std::size_t>(natom), 0);

std::vector<int*> firstneigh(static_cast<std::size_t>(natom), NULL);

for (int iat = 0; iat < nowned_atoms; ++iat)

{

ilist[static_cast<std::size_t>(iat)] = iat;

numneigh[static_cast<std::size_t>(iat)] = neighbor_list.get_numneigh(iat);

firstneigh[static_cast<std::size_t>(iat)] = const_cast<int*>(neighbor_list.get_firstneigh(iat));

}

double local_energy = 0.0;

double local_virial[6] = {0.0, 0.0, 0.0, 0.0, 0.0, 0.0};

ModuleBase::timer::start("ESolver_NEP", "compute");

nep.compute_for_lammps(nowned_atoms,

nowned_atoms,

nowned_atoms > 0 ? ilist.data() : NULL,

numneigh.data(),

firstneigh.data(),

local_type.data(),

md_type_to_nep_type_.data(),

position_ptrs.data(),

local_energy,

local_virial,

NULL,

force_ptrs.data(),

NULL);

ModuleBase::timer::end("ESolver_NEP", "compute");

std::vector<LocalAtom>& mutable_owned_atoms = mdcell.mutable_owned_atoms();

std::vector<LocalAtom>& mutable_ghost_atoms = mdcell.mutable_ghost_atoms();

for (int iat = 0; iat < nowned_atoms; ++iat)

{

mutable_owned_atoms[static_cast<std::size_t>(iat)].force.set(force[static_cast<std::size_t>(iat)][0],

force[static_cast<std::size_t>(iat)][1],

force[static_cast<std::size_t>(iat)][2]);

}

for (int iat = 0; iat < nghost; ++iat)

{

mutable_ghost_atoms[static_cast<std::size_t>(iat)].force.set(force[static_cast<std::size_t>(nowned_atoms + iat)][0],

force[static_cast<std::size_t>(nowned_atoms + iat)][1],

force[static_cast<std::size_t>(nowned_atoms + iat)][2]);

}

mdcell.accumulate_ghost_forces();

#ifdef __MPI

MPI_Allreduce(MPI_IN_PLACE, &local_energy, 1, MPI_DOUBLE, MPI_SUM, mdcell.communicator());

MPI_Allreduce(MPI_IN_PLACE, local_virial, 6, MPI_DOUBLE, MPI_SUM, mdcell.communicator());

#endif

const double fact_e = 1.0 / ModuleBase::Ry_to_eV;

const double fact_f = 1.0 / (ModuleBase::Ry_to_eV * ModuleBase::ANGSTROM_AU);

const double fact_v = 1.0 / (mdcell.omega() * ModuleBase::Ry_to_eV);

nep_potential = local_energy * fact_e;

for (int iat = 0; iat < nowned_atoms; ++iat)

{

LocalAtom& atom = mutable_owned_atoms[static_cast<std::size_t>(iat)];

atom.force *= fact_f;

}

nep_virial(0, 0) = local_virial[0] * fact_v;

nep_virial(1, 1) = local_virial[1] * fact_v;

nep_virial(2, 2) = local_virial[2] * fact_v;

nep_virial(0, 1) = nep_virial(1, 0) = local_virial[3] * fact_v;

nep_virial(0, 2) = nep_virial(2, 0) = local_virial[4] * fact_v;

nep_virial(1, 2) = nep_virial(2, 1) = local_virial[5] * fact_v;

#endif

ModuleBase::timer::end("ESolver_NEP", "runner");

return;

}

UnitCell& ucell = static_cast<UnitCell&>(basecell);

// note that NEP are column major, thus a transpose is needed

// cell

std::vector<double> cell(9, 0.0);

cell[0] = ucell.latvec.e11 * ucell.lat0_angstrom;

cell[1] = ucell.latvec.e21 * ucell.lat0_angstrom;

cell[2] = ucell.latvec.e31 * ucell.lat0_angstrom;

cell[3] = ucell.latvec.e12 * ucell.lat0_angstrom;

cell[4] = ucell.latvec.e22 * ucell.lat0_angstrom;

cell[5] = ucell.latvec.e32 * ucell.lat0_angstrom;

cell[6] = ucell.latvec.e13 * ucell.lat0_angstrom;

cell[7] = ucell.latvec.e23 * ucell.lat0_angstrom;

cell[8] = ucell.latvec.e33 * ucell.lat0_angstrom;

// coord

std::vector<double> coord(3 * ucell.nat, 0.0);

int iat = 0;

const int nat = ucell.nat;

for (int it = 0; it < ucell.ntype; ++it)

{

for (int ia = 0; ia < ucell.atoms[it].na; ++ia)

{

coord[iat] = ucell.atoms[it].tau[ia].x * ucell.lat0_angstrom;

coord[iat + nat] = ucell.atoms[it].tau[ia].y * ucell.lat0_angstrom;

coord[iat + 2 * nat] = ucell.atoms[it].tau[ia].z * ucell.lat0_angstrom;

iat++;

}

}

assert(ucell.nat == iat);

#ifdef __NEP

nep_potential = 0.0;

nep_force.zero_out();

nep_virial.zero_out();

nep.compute(atype, cell, coord, _e, _f, _v);

// unit conversion

const double fact_e = 1.0 / ModuleBase::Ry_to_eV;

const double fact_f = 1.0 / (ModuleBase::Ry_to_eV * ModuleBase::ANGSTROM_AU);

const double fact_v = 1.0 / (ucell.omega * ModuleBase::Ry_to_eV);

// potential energy

nep_potential = fact_e * std::accumulate(_e.begin(), _e.end(), 0.0);

GlobalV::ofs_running << " #TOTAL ENERGY# " << std::setprecision(11) << nep_potential * ModuleBase::Ry_to_eV << " eV"

<< std::endl;

// forces

for (int i = 0; i < nat; ++i)

{

nep_force(i, 0) = _f[i] * fact_f;

nep_force(i, 1) = _f[i + nat] * fact_f;

nep_force(i, 2) = _f[i + 2 * nat] * fact_f;

}

// virial

std::vector<double> v_sum(9, 0.0);

for (int j = 0; j < 9; ++j)

{

for (int i = 0; i < nat; ++i)

{

int index = j * nat + i;

v_sum[j] += _v[index];

}

}

// virial -> stress

for (int i = 0; i < 3; ++i)

{

for (int j = 0; j < 3; ++j)

{

nep_virial(i, j) = v_sum[3 * i + j] * fact_v;

}

}

#else

ModuleBase::WARNING_QUIT("ESolver_NEP", "Please recompile with -D__NEP");

#endif

ModuleBase::timer::end("ESolver_NEP", "runner");

{

if (basecell.kind() == BaseCell::Kind::mdcell)

{

const MDCell& mdcell = static_cast<const MDCell&>(basecell);

force.create(mdcell.nowned_atoms(), 3);

for (int iat = 0; iat < mdcell.nowned_atoms(); ++iat)

{

const LocalAtom& atom = mdcell.owned_atoms()[static_cast<std::size_t>(iat)];

force(iat, 0) = atom.force.x;

force(iat, 1) = atom.force.y;

force(iat, 2) = atom.force.z;

}

return;

}

force = nep_force;

UnitCell& ucell = static_cast<UnitCell&>(basecell);

ModuleIO::print_force(GlobalV::ofs_running, ucell, "TOTAL-FORCE (eV/Angstrom)", force, false);

{

stress = nep_virial;

if (basecell.kind() == BaseCell::Kind::unitcell)

{

ModuleIO::print_stress("TOTAL-STRESS", stress, true, false, GlobalV::ofs_running);

}

else

{

#ifdef __MPI

if (static_cast<MDCell&>(basecell).mpi_rank() == 0)

#endif

{

ModuleIO::print_stress("TOTAL-STRESS", stress, true, false, GlobalV::ofs_running);

}

}

// external stress

double unit_transform = ModuleBase::RYDBERG_SI / pow(ModuleBase::BOHR_RADIUS_SI, 3) * 1.0e-8;

double external_stress[3] = {this->inp_->press1, this->inp_->press2, this->inp_->press3};

for (int i = 0; i < 3; i++)

{

stress(i, i) -= external_stress[i] / unit_transform;

}

{

static_cast<void>(basecell);

GlobalV::ofs_running << "\n --------------------------------------------" << std::endl;

GlobalV::ofs_running << std::setprecision(16);

GlobalV::ofs_running << " !FINAL_ETOT_IS " << nep_potential * ModuleBase::Ry_to_eV << " eV" << std::endl;

GlobalV::ofs_running << " --------------------------------------------\n\n" << std::endl;

{

std::vector<std::string> type_labels(static_cast<std::size_t>(ucell.ntype));

for (int it = 0; it < ucell.ntype; ++it)

{

type_labels[static_cast<std::size_t>(it)] = ucell.atoms[it].label;

}

initialize_type_map_(type_labels);

int iat = 0;

for (int it = 0; it < ucell.ntype; ++it)

{

for (int ia = 0; ia < ucell.atoms[it].na; ++ia)

{

atype[static_cast<std::size_t>(iat++)] = md_type_to_nep_type_[static_cast<std::size_t>(it)];

}

}

assert(ucell.nat == iat);

{

// parse the element list from NEP model file

std::unordered_map<std::string, int> label;

std::string temp;

for (int i = 0; i < nep.element_list.size(); ++i)

{

label[nep.element_list[i]] = i; //> label: map from element string to index int.

}

std::cout << "\n Element list of model file " << nep_file << " " << std::endl;

std::cout << " ----------------------------------------------------------------";

int count = 0;

for (auto it = label.begin(); it != label.end(); ++it)

{

if (count % 5 == 0)

{

std::cout << std::endl;

std::cout << " ";

}

count++;

temp = it->first + ": " + std::to_string(it->second);

std::cout << std::left << std::setw(10) << temp;

}

std::cout << "\n -----------------------------------------------------------------" << std::endl;

// parse the atype based on the element list

md_type_to_nep_type_.resize(type_labels.size());

for (std::size_t it = 0; it < type_labels.size(); ++it)

{

const std::unordered_map<std::string, int>::const_iterator found = label.find(type_labels[it]);

if (found == label.end())

{

ModuleBase::WARNING_QUIT("ESolver_NEP",

"The label " + type_labels[it] + " is not found in the type map.");

}

md_type_to_nep_type_[it] = found->second;

}