From: @zhangxinfeng3
Reviewed-by: @wang_zi_dong,@ljl0711
Signed-off-by: @ljl0711
pull/14144/MERGE
mindspore-ci-bot 4 years ago committed by Gitee
commit 0a0dc05d51

@ -35,6 +35,7 @@ class PMEEnergyGpuKernel : public GpuKernel {
bool Init(const CNodePtr &kernel_node) override {
kernel_node_ = kernel_node;
atom_numbers = static_cast<int>(GetAttr<int64_t>(kernel_node, "atom_numbers"));
excluded_numbers = static_cast<int>(GetAttr<int64_t>(kernel_node, "excluded_numbers"));
beta = static_cast<float>(GetAttr<float_t>(kernel_node, "beta"));
fftx = static_cast<int>(GetAttr<int64_t>(kernel_node, "fftx"));
ffty = static_cast<int>(GetAttr<int64_t>(kernel_node, "ffty"));
@ -90,7 +91,7 @@ class PMEEnergyGpuKernel : public GpuKernel {
input_size_list_.push_back(atom_numbers * sizeof(VECTOR));
input_size_list_.push_back(atom_numbers * sizeof(T1));
input_size_list_.push_back(max_nl_numbers * sizeof(T1));
input_size_list_.push_back(atom_numbers * sizeof(VECTOR));
input_size_list_.push_back(sizeof(VECTOR));
input_size_list_.push_back(atom_numbers * sizeof(T1));
input_size_list_.push_back(excluded_numbers * sizeof(T1));
@ -118,7 +119,7 @@ class PMEEnergyGpuKernel : public GpuKernel {
std::vector<size_t> output_size_list_;
std::vector<size_t> workspace_size_list_;
int atom_numbers;
int excluded_numbers = 2719;
int excluded_numbers;
int max_nl_numbers = 800;
int fftx;
int ffty;

@ -35,6 +35,7 @@ class PMEExcludedForceGpuKernel : public GpuKernel {
bool Init(const CNodePtr &kernel_node) override {
kernel_node_ = kernel_node;
atom_numbers = static_cast<int>(GetAttr<int64_t>(kernel_node, "atom_numbers"));
excluded_numbers = static_cast<int>(GetAttr<int64_t>(kernel_node, "excluded_numbers"));
beta = static_cast<float>(GetAttr<float_t>(kernel_node, "beta"));
InitSizeLists();
return true;
@ -62,7 +63,7 @@ class PMEExcludedForceGpuKernel : public GpuKernel {
protected:
void InitSizeLists() override {
input_size_list_.push_back(atom_numbers * sizeof(UNSIGNED_INT_VECTOR));
input_size_list_.push_back(atom_numbers * sizeof(VECTOR));
input_size_list_.push_back(sizeof(VECTOR));
input_size_list_.push_back(atom_numbers * sizeof(T));
input_size_list_.push_back(atom_numbers * sizeof(T1));
input_size_list_.push_back(excluded_numbers * sizeof(T1));
@ -77,7 +78,7 @@ class PMEExcludedForceGpuKernel : public GpuKernel {
std::vector<size_t> output_size_list_;
std::vector<size_t> workspace_size_list_;
int atom_numbers;
int excluded_numbers = 2719;
int excluded_numbers;
float beta;
struct VECTOR {
float x;

File diff suppressed because it is too large Load Diff

@ -119,18 +119,19 @@ class Particle_Mesh_Ewald(nn.Cell):
tempi += 4
def PME_Energy(self, uint_crd, charge, nl_atom_numbers, nl_atom_serial, uint_dr_to_dr_cof, excluded_list_start,
excluded_list, excluded_numbers):
excluded_list, excluded_numbers, excluded_atom_numbers):
"""PME_Energy"""
self.pmee = P.PMEEnergy(self.atom_numbers, self.beta, self.fftx, self.ffty, self.fftz)
self.pmee = P.PMEEnergy(self.atom_numbers, excluded_atom_numbers, self.beta, self.fftx, self.ffty, self.fftz)
self.reciprocal_energy, self.self_energy, self.direct_energy, self.correction_energy = \
self.pmee(self.box_length, uint_crd, charge, nl_atom_numbers, nl_atom_serial, uint_dr_to_dr_cof,
excluded_list_start, excluded_list, excluded_numbers)
return self.reciprocal_energy, self.self_energy, self.direct_energy, self.correction_energy
def PME_Excluded_Force(self, uint_crd, scaler, charge, excluded_list_start, excluded_list,
excluded_numbers):
excluded_numbers, excluded_atom_numbers):
"""PME Excluded Force"""
self.pmeef = P.PMEExcludedForce(atom_numbers=self.atom_numbers, beta=self.beta)
self.pmeef = P.PMEExcludedForce(atom_numbers=self.atom_numbers, excluded_numbers=excluded_atom_numbers,
beta=self.beta)
self.frc = self.pmeef(uint_crd, scaler, charge, excluded_list_start, excluded_list, excluded_numbers)
return self.frc

@ -124,7 +124,7 @@ class Simulation(nn.Cell):
pme_excluded_frc = pme_method.PME_Excluded_Force(
md_info.uint_crd, md_info.uint_dr_to_dr_cof, md_info.charge,
nb_info.excluded_list_start, nb_info.excluded_list,
nb_info.excluded_numbers)
nb_info.excluded_numbers, nb_info.excluded_atom_numbers)
frc_t += pme_excluded_frc.asnumpy()
pme_reciprocal_frc = pme_method.PME_Reciprocal_Force(md_info.uint_crd, md_info.charge)
@ -149,7 +149,7 @@ class Simulation(nn.Cell):
_ = self.pme_method.PME_Energy(
self.md_info.uint_crd, self.md_info.charge, self.nb_info.nl_atom_numbers, self.nb_info.nl_atom_serial,
self.md_info.uint_dr_to_dr_cof, self.nb_info.excluded_list_start, self.nb_info.excluded_list,
self.nb_info.excluded_numbers)
self.nb_info.excluded_numbers, self.nb_info.excluded_atom_numbers)
_ = self.pme_method.Energy_Device_To_Host()
def Main_After_Calculate_Energy(self):

Loading…
Cancel
Save