p60975123
/
mindspore
forked from mindspore-Ecosystem/mindspore
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

master 03252 sponge ops zoo cu

This commit is contained in:
q00596439 2021-03-25 20:42:18 +08:00
parent a48785cdcc
commit aa7fdf53a7
33 changed files with 441 additions and 198 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/angle/angle_atom_energy_impl.cu
View File

@ -50,7 +50,7 @@ __global__ void AngleAtomEnergyKernel(int angle_numbers, const UNSIGNED_INT_VECT
void AngleAtomEnergy(int angle_numbers, int atom_numbers, const int *uint_crd_f, const float *scaler_f, void AngleAtomEnergy(int angle_numbers, int atom_numbers, const int *uint_crd_f, const float *scaler_f,
const int *atom_a, const int *atom_b, const int *atom_c, const float *angle_k, const int *atom_a, const int *atom_b, const int *atom_c, const float *angle_k,
const float *angle_theta0, float *ene, cudaStream_t stream) { const float *angle_theta0, float *ene, cudaStream_t stream) {
Reset_List<<<ceilf(static_cast<float>(atom_numbers) / 128), 128>>>(atom_numbers, ene, 0.); Reset_List<<<ceilf(static_cast<float>(atom_numbers) / 128), 128, 0, stream>>>(atom_numbers, ene, 0.);
size_t thread_per_block = 128; size_t thread_per_block = 128;
size_t block_per_grid = ceilf(static_cast<float>(angle_numbers) / 128); size_t block_per_grid = ceilf(static_cast<float>(angle_numbers) / 128);
UNSIGNED_INT_VECTOR *uint_crd = UNSIGNED_INT_VECTOR *uint_crd =

2
mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/angle/angle_force_impl.cu
View File

@ -69,7 +69,7 @@ __global__ void AngleForceKernel(int angle_numbers, const UNSIGNED_INT_VECTOR *u
void AngleForce(int angle_numbers, int atom_numbers, const int *uint_crd_f, const float *scaler_f, const int *atom_a, void AngleForce(int angle_numbers, int atom_numbers, const int *uint_crd_f, const float *scaler_f, const int *atom_a,
const int *atom_b, const int *atom_c, const float *angle_k, const float *angle_theta0, float *frc_f, const int *atom_b, const int *atom_c, const float *angle_k, const float *angle_theta0, float *frc_f,
cudaStream_t stream) { cudaStream_t stream) {
Reset_List<<<ceilf(static_cast<float>(3. * atom_numbers) / 128), 128>>>(3 * atom_numbers, frc_f, 0.); Reset_List<<<ceilf(static_cast<float>(3. * atom_numbers) / 128), 128, 0, stream>>>(3 * atom_numbers, frc_f, 0.);
size_t thread_per_block = 128; size_t thread_per_block = 128;
size_t block_per_grid = ceilf(static_cast<float>(angle_numbers) / 128); size_t block_per_grid = ceilf(static_cast<float>(angle_numbers) / 128);
UNSIGNED_INT_VECTOR *uint_crd = UNSIGNED_INT_VECTOR *uint_crd =

2
mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/angle/angle_force_with_atom_energy_impl.cu
View File

@ -73,7 +73,7 @@ __global__ void AngleForceWithAtomEnergyKernel(int angle_numbers, const UNSIGNED
void AngleForceWithAtomEnergy(int angle_numbers, int atom_numbers, const int *uint_crd_f, const float *scaler_f, void AngleForceWithAtomEnergy(int angle_numbers, int atom_numbers, const int *uint_crd_f, const float *scaler_f,
const int *atom_a, const int *atom_b, const int *atom_c, const float *angle_k, const int *atom_a, const int *atom_b, const int *atom_c, const float *angle_k,
const float *angle_theta0, float *frc_f, float *ene, cudaStream_t stream) { const float *angle_theta0, float *frc_f, float *ene, cudaStream_t stream) {
Reset_List<<<ceilf(static_cast<float>(3. * atom_numbers) / 128), 128>>>(3 * atom_numbers, frc_f, 0.); Reset_List<<<ceilf(static_cast<float>(3. * atom_numbers) / 128), 128, 0, stream>>>(3 * atom_numbers, frc_f, 0.);
size_t thread_per_block = 128; size_t thread_per_block = 128;
size_t block_per_grid = ceilf(static_cast<float>(angle_numbers) / 128); size_t block_per_grid = ceilf(static_cast<float>(angle_numbers) / 128);
UNSIGNED_INT_VECTOR *uint_crd = UNSIGNED_INT_VECTOR *uint_crd =

2
mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/bond/bond_atom_energy_cuda_gpu_impl.cu
View File

@ -41,7 +41,7 @@ __global__ void BondAtomEnergyCudaKernel(const int bond_numbers, const UNSIGNED_
void BondAtomEnergy(int bond_numbers, int atom_numbers, const int *uint_crd_f, const float *scaler_f, const int *atom_a, void BondAtomEnergy(int bond_numbers, int atom_numbers, const int *uint_crd_f, const float *scaler_f, const int *atom_a,
const int *atom_b, const float *bond_k, const float *bond_r0, float *atom_ene, const int *atom_b, const float *bond_k, const float *bond_r0, float *atom_ene,
cudaStream_t stream) { cudaStream_t stream) {
Reset_List<<<ceilf(static_cast<float>(atom_numbers) / 128), 128>>>(atom_numbers, atom_ene, 0.); Reset_List<<<ceilf(static_cast<float>(atom_numbers) / 128), 128, 0, stream>>>(atom_numbers, atom_ene, 0.);
size_t thread_per_block = 128; size_t thread_per_block = 128;
size_t block_per_grid = ceilf(static_cast<float>(bond_numbers) / 128); size_t block_per_grid = ceilf(static_cast<float>(bond_numbers) / 128);
UNSIGNED_INT_VECTOR *uint_crd = UNSIGNED_INT_VECTOR *uint_crd =

2
mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/bond/bond_force_cuda_gpu_impl.cu
View File

@ -45,7 +45,7 @@ __global__ void BondForceCudaKernel(int bond_numbers, const UNSIGNED_INT_VECTOR
void BondForce(int bond_numbers, int atom_numbers, const int *uint_crd_f, const float *scaler_f, const int *atom_a, void BondForce(int bond_numbers, int atom_numbers, const int *uint_crd_f, const float *scaler_f, const int *atom_a,
const int *atom_b, const float *bond_k, const float *bond_r0, float *frc_f, cudaStream_t stream) { const int *atom_b, const float *bond_k, const float *bond_r0, float *frc_f, cudaStream_t stream) {
Reset_List<<<ceilf(static_cast<float>(3. * atom_numbers) / 128), 128>>>(3 * atom_numbers, frc_f, 0.); Reset_List<<<ceilf(static_cast<float>(3. * atom_numbers) / 128), 128, 0, stream>>>(3 * atom_numbers, frc_f, 0.);
size_t thread_per_block = 128; size_t thread_per_block = 128;
size_t block_per_grid = ceilf(static_cast<float>(bond_numbers) / 128); size_t block_per_grid = ceilf(static_cast<float>(bond_numbers) / 128);
UNSIGNED_INT_VECTOR *uint_crd = UNSIGNED_INT_VECTOR *uint_crd =

4
mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/bond/bond_force_with_atom_energy_impl.cu
View File

@ -52,8 +52,8 @@ __global__ void BondForceWithAtomEnergyKernel(int bond_numbers, const UNSIGNED_I
void BondForceWithAtomEnergy(int bond_numbers, int atom_numbers, const int *uint_crd_f, const float *scaler_f, void BondForceWithAtomEnergy(int bond_numbers, int atom_numbers, const int *uint_crd_f, const float *scaler_f,
const int *atom_a, const int *atom_b, const float *bond_k, const float *bond_r0, const int *atom_a, const int *atom_b, const float *bond_k, const float *bond_r0,
float *frc_f, float *atom_e, cudaStream_t stream) { float *frc_f, float *atom_e, cudaStream_t stream) {
Reset_List<<<ceilf(static_cast<float>(3. * atom_numbers) / 128), 128>>>(3 * atom_numbers, frc_f, 0.); Reset_List<<<ceilf(static_cast<float>(3. * atom_numbers) / 128), 128, 0, stream>>>(3 * atom_numbers, frc_f, 0.);
Reset_List<<<ceilf(static_cast<float>(atom_numbers) / 128), 128>>>(atom_numbers, atom_e, 0.); Reset_List<<<ceilf(static_cast<float>(atom_numbers) / 128), 128, 0, stream>>>(atom_numbers, atom_e, 0.);
size_t thread_per_block = 128; size_t thread_per_block = 128;
size_t block_per_grid = ceilf(static_cast<float>(bond_numbers) / 128); size_t block_per_grid = ceilf(static_cast<float>(bond_numbers) / 128);
UNSIGNED_INT_VECTOR *uint_crd = UNSIGNED_INT_VECTOR *uint_crd =

4
mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/bond/bond_force_with_atom_virial_impl.cu
View File

@ -52,8 +52,8 @@ __global__ void BondForceWithAtomVirialKernel(int bond_numbers, const UNSIGNED_I
void BondForceWithAtomVirial(int bond_numbers, int atom_numbers, const int *uint_crd_f, const float *scaler_f, void BondForceWithAtomVirial(int bond_numbers, int atom_numbers, const int *uint_crd_f, const float *scaler_f,
const int *atom_a, const int *atom_b, const float *bond_k, const float *bond_r0, const int *atom_a, const int *atom_b, const float *bond_k, const float *bond_r0,
float *frc_f, float *atom_v, cudaStream_t stream) { float *frc_f, float *atom_v, cudaStream_t stream) {
Reset_List<<<ceilf(static_cast<float>(3. * atom_numbers) / 128), 128>>>(3 * atom_numbers, frc_f, 0.); Reset_List<<<ceilf(static_cast<float>(3. * atom_numbers) / 128), 128, 0, stream>>>(3 * atom_numbers, frc_f, 0.);
Reset_List<<<ceilf(static_cast<float>(atom_numbers) / 128), 128>>>(atom_numbers, atom_v, 0.); Reset_List<<<ceilf(static_cast<float>(atom_numbers) / 128), 128, 0, stream>>>(atom_numbers, atom_v, 0.);
size_t thread_per_block = 128; size_t thread_per_block = 128;
size_t block_per_grid = ceilf(static_cast<float>(bond_numbers) / 128); size_t block_per_grid = ceilf(static_cast<float>(bond_numbers) / 128);
UNSIGNED_INT_VECTOR *uint_crd = UNSIGNED_INT_VECTOR *uint_crd =

2
mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/dihedral/dihedral_atom_energy_impl.cu
View File

@ -67,7 +67,7 @@ void DihedralAtomEnergy(int dihedral_numbers, int atom_numbers, const int *uint_
const int *atom_a, const int *atom_b, const int *atom_c, const int *atom_d, const int *ipn, const int *atom_a, const int *atom_b, const int *atom_c, const int *atom_d, const int *ipn,
const float *pk, const float *gamc, const float *gams, const float *pn, float *ene, const float *pk, const float *gamc, const float *gams, const float *pn, float *ene,
cudaStream_t stream) { cudaStream_t stream) {
Reset_List<<<ceilf(static_cast<float>(atom_numbers) / 128), 128>>>(atom_numbers, ene, 0.); Reset_List<<<ceilf(static_cast<float>(atom_numbers) / 128), 128, 0, stream>>>(atom_numbers, ene, 0.);
size_t thread_per_block = 128; size_t thread_per_block = 128;
size_t block_per_grid = ceilf(static_cast<float>(dihedral_numbers) / 128); size_t block_per_grid = ceilf(static_cast<float>(dihedral_numbers) / 128);
UNSIGNED_INT_VECTOR *uint_crd = UNSIGNED_INT_VECTOR *uint_crd =

2
mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/dihedral/dihedral_force_impl.cu
View File

@ -103,7 +103,7 @@ void DihedralForce(int dihedral_numbers, int atom_numbers, const int *uint_crd_f
const int *atom_a, const int *atom_b, const int *atom_c, const int *atom_d, const int *ipn, const int *atom_a, const int *atom_b, const int *atom_c, const int *atom_d, const int *ipn,
const float *pk, const float *gamc, const float *gams, const float *pn, float *frc_f, const float *pk, const float *gamc, const float *gams, const float *pn, float *frc_f,
cudaStream_t stream) { cudaStream_t stream) {
Reset_List<<<ceilf(static_cast<float>(3. * atom_numbers) / 128), 128>>>(3 * atom_numbers, frc_f, 0.); Reset_List<<<ceilf(static_cast<float>(3. * atom_numbers) / 128), 128, 0, stream>>>(3 * atom_numbers, frc_f, 0.);
size_t thread_per_block = 128; size_t thread_per_block = 128;
size_t block_per_grid = ceilf(static_cast<float>(dihedral_numbers) / 128); size_t block_per_grid = ceilf(static_cast<float>(dihedral_numbers) / 128);
UNSIGNED_INT_VECTOR *uint_crd = UNSIGNED_INT_VECTOR *uint_crd =

2
mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/dihedral/dihedral_force_with_atom_energy_impl.cu
View File

@ -107,7 +107,7 @@ void DihedralForceWithAtomEnergy(int dihedral_numbers, int atom_numbers, const i
const int *atom_a, const int *atom_b, const int *atom_c, const int *atom_d, const int *atom_a, const int *atom_b, const int *atom_c, const int *atom_d,
const int *ipn, const float *pk, const float *gamc, const float *gams, const float *pn, const int *ipn, const float *pk, const float *gamc, const float *gams, const float *pn,
float *frc_f, float *ene, cudaStream_t stream) { float *frc_f, float *ene, cudaStream_t stream) {
Reset_List<<<ceilf(static_cast<float>(3. * atom_numbers) / 128), 128>>>(3 * atom_numbers, frc_f, 0.); Reset_List<<<ceilf(static_cast<float>(3. * atom_numbers) / 128), 128, 0, stream>>>(3 * atom_numbers, frc_f, 0.);
size_t thread_per_block = 128; size_t thread_per_block = 128;
size_t block_per_grid = ceilf(static_cast<float>(dihedral_numbers) / 128); size_t block_per_grid = ceilf(static_cast<float>(dihedral_numbers) / 128);
UNSIGNED_INT_VECTOR *uint_crd = UNSIGNED_INT_VECTOR *uint_crd =

2
mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/lj/lj_force_impl.cu
View File

@ -92,7 +92,7 @@ void LJForce(const int atom_numbers, const float cutoff_square, const int *uint_
const float *charge, const float *scaler_f, float *uint_crd_with_LJ, int *nl_atom_numbers, const float *charge, const float *scaler_f, float *uint_crd_with_LJ, int *nl_atom_numbers,
int *nl_atom_serial, int *nl, const float *d_LJ_A, const float *d_LJ_B, float *frc_f, int *nl_atom_serial, int *nl, const float *d_LJ_A, const float *d_LJ_B, float *frc_f,
cudaStream_t stream) { cudaStream_t stream) {
Reset_List<<<ceilf(static_cast<float>(3. * atom_numbers) / 128), 128>>>(3 * atom_numbers, frc_f, 0.); Reset_List<<<ceilf(static_cast<float>(3. * atom_numbers) / 128), 128, 0, stream>>>(3 * atom_numbers, frc_f, 0.);
VECTOR *frc = reinterpret_cast<VECTOR *>(frc_f); VECTOR *frc = reinterpret_cast<VECTOR *>(frc_f);
VECTOR *scaler = const_cast<VECTOR *>(reinterpret_cast<const VECTOR *>(scaler_f)); VECTOR *scaler = const_cast<VECTOR *>(reinterpret_cast<const VECTOR *>(scaler_f));
int max_neighbor_numbers = 800; int max_neighbor_numbers = 800;

2
mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/lj/lj_force_with_pme_direct_force_impl.cu
View File

@ -106,7 +106,7 @@ void LJForceWithPMEDirectForce(const int atom_numbers, const float cutoff, const
const int *LJtype, const float *charge, const float *scaler_f, float *uint_crd_with_LJ, const int *LJtype, const float *charge, const float *scaler_f, float *uint_crd_with_LJ,
int *nl_atom_numbers, int *nl_atom_serial, int *nl, const float *d_LJ_A, int *nl_atom_numbers, int *nl_atom_serial, int *nl, const float *d_LJ_A,
const float *d_LJ_B, float *frc_f, cudaStream_t stream) { const float *d_LJ_B, float *frc_f, cudaStream_t stream) {
Reset_List<<<ceilf(static_cast<float>(3. * atom_numbers) / 128), 128>>>(3 * atom_numbers, frc_f, 0.); Reset_List<<<ceilf(static_cast<float>(3. * atom_numbers) / 128), 128, 0, stream>>>(3 * atom_numbers, frc_f, 0.);
VECTOR *frc = reinterpret_cast<VECTOR *>(frc_f); VECTOR *frc = reinterpret_cast<VECTOR *>(frc_f);
VECTOR *scaler = const_cast<VECTOR *>(reinterpret_cast<const VECTOR *>(scaler_f)); VECTOR *scaler = const_cast<VECTOR *>(reinterpret_cast<const VECTOR *>(scaler_f));
int max_neighbor_numbers = 800; int max_neighbor_numbers = 800;

2
mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/nb14/dihedral_14_cf_atom_energy_impl.cu
View File

@ -66,7 +66,7 @@ void Dihedral14CFAtomEnergy(const int dihedral_14_numbers, const int atom_number
atom_numbers, uint_crd, uint_crd_with_LJ, LJtype, charge); atom_numbers, uint_crd, uint_crd_with_LJ, LJtype, charge);
VECTOR *boxlength = const_cast<VECTOR *>(reinterpret_cast<const VECTOR *>(boxlength_f)); VECTOR *boxlength = const_cast<VECTOR *>(reinterpret_cast<const VECTOR *>(boxlength_f));
Reset_List<<<ceilf(static_cast<float>(3. * atom_numbers) / 128), 128>>>(atom_numbers, ene, 0.); Reset_List<<<ceilf(static_cast<float>(3. * atom_numbers) / 128), 128, 0, stream>>>(atom_numbers, ene, 0.);
Dihedral14CFAtomEnergyKernel<<<block_per_grid, thread_per_block, 0, stream>>>( Dihedral14CFAtomEnergyKernel<<<block_per_grid, thread_per_block, 0, stream>>>(
dihedral_14_numbers, uint_crd_with_LJ, boxlength, a_14, b_14, cf_scale_factor, ene); dihedral_14_numbers, uint_crd_with_LJ, boxlength, a_14, b_14, cf_scale_factor, ene);

15
mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/nb14/dihedral_14_cf_energy_impl.cu
View File

@ -52,21 +52,20 @@ __global__ void Dihedral14CFEnergyKernel(const int dihedral_14_numbers, const UI
} }
void Dihedral14CFEnergy(const int dihedral_14_numbers, const int atom_numbers, const int *uint_crd_f, const int *LJtype, void Dihedral14CFEnergy(const int dihedral_14_numbers, const int atom_numbers, const int *uint_crd_f, const int *LJtype,
const float *charge, const float *boxlength_f, const int *a_14, const int *b_14, const float *charge, float *uint_crd_with_LJ_f, const float *boxlength_f, const int *a_14,
const float *cf_scale_factor, float *ene, cudaStream_t stream) { const int *b_14, const float *cf_scale_factor, float *ene, cudaStream_t stream) {
size_t thread_per_block = 128; size_t thread_per_block = 128;
size_t block_per_grid = ceilf(static_cast<float>(atom_numbers) / 128); size_t block_per_grid = ceilf(static_cast<float>(atom_numbers) / 128);
UINT_VECTOR_LJ_TYPE *uint_crd_with_LJ = NULL;
Cuda_Malloc_Safely(reinterpret_cast<void **>(&uint_crd_with_LJ), sizeof(UINT_VECTOR_LJ_TYPE) * atom_numbers);
UNSIGNED_INT_VECTOR *uint_crd = UNSIGNED_INT_VECTOR *uint_crd =
const_cast<UNSIGNED_INT_VECTOR *>(reinterpret_cast<const UNSIGNED_INT_VECTOR *>(uint_crd_f)); const_cast<UNSIGNED_INT_VECTOR *>(reinterpret_cast<const UNSIGNED_INT_VECTOR *>(uint_crd_f));
UINT_VECTOR_LJ_TYPE *uint_crd_with_LJ = reinterpret_cast<UINT_VECTOR_LJ_TYPE *>(uint_crd_with_LJ_f);
Copy_Crd_To_New_Crd_Start<<<ceilf(static_cast<float>(atom_numbers) / 32), 32, 0, stream>>>( Copy_Crd_To_New_Crd_Start<<<ceilf(static_cast<float>(atom_numbers) / 32), 32, 0, stream>>>(
atom_numbers, uint_crd, uint_crd_with_LJ, LJtype, charge); atom_numbers, uint_crd, uint_crd_with_LJ, LJtype, charge);
VECTOR *boxlength = const_cast<VECTOR *>(reinterpret_cast<const VECTOR *>(boxlength_f)); VECTOR *boxlength = const_cast<VECTOR *>(reinterpret_cast<const VECTOR *>(boxlength_f));
Reset_List<<<ceilf(static_cast<float>(3. * atom_numbers) / 128), 128>>>(atom_numbers, ene, 0.); Reset_List<<<ceilf(static_cast<float>(3. * atom_numbers) / 128), 128, 0, stream>>>(atom_numbers, ene, 0.);
Dihedral14CFEnergyKernel<<<block_per_grid, thread_per_block, 0, stream>>>( Dihedral14CFEnergyKernel<<<block_per_grid, thread_per_block, 0, stream>>>(
dihedral_14_numbers, uint_crd_with_LJ, boxlength, a_14, b_14, cf_scale_factor, ene); dihedral_14_numbers, uint_crd_with_LJ, boxlength, a_14, b_14, cf_scale_factor, ene);
@ -76,5 +75,5 @@ void Dihedral14CFEnergy(const int dihedral_14_numbers, const int atom_numbers, c
} }
void Dihedral14CFEnergy(const int dihedral_14_numbers, const int atom_numbers, const int *uint_crd_f, const int *LJtype, void Dihedral14CFEnergy(const int dihedral_14_numbers, const int atom_numbers, const int *uint_crd_f, const int *LJtype,
const float *charge, const float *boxlength_f, const int *a_14, const int *b_14, const float *charge, float *uint_crd_with_LJ_f, const float *boxlength_f, const int *a_14,
const float *cf_scale_factor, float *ene, cudaStream_t stream); const int *b_14, const float *cf_scale_factor, float *ene, cudaStream_t stream);

4
mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/nb14/dihedral_14_cf_energy_impl.cuh
View File

@ -20,6 +20,6 @@
#include "runtime/device/gpu/cuda_common.h" #include "runtime/device/gpu/cuda_common.h"
void Dihedral14CFEnergy(const int dihedral_14_numbers, const int atom_numbers, const int *uint_crd_f, const int *LJtype, void Dihedral14CFEnergy(const int dihedral_14_numbers, const int atom_numbers, const int *uint_crd_f, const int *LJtype,
const float *charge, const float *boxlength, const int *a_14, const int *b_14, const float *charge, float *uint_crd_with_LJ_f, const float *boxlength_f, const int *a_14,
const float *cf_scale_factor, float *ene, cudaStream_t stream); const int *b_14, const float *cf_scale_factor, float *ene, cudaStream_t stream);
#endif // MINDSPORE_CCSRC_KERNEL_GPU_CUDA_IMPL_SPONGE_NB14_DIHEDRAL_14_CF_ENERGY_IMPL_H #endif // MINDSPORE_CCSRC_KERNEL_GPU_CUDA_IMPL_SPONGE_NB14_DIHEDRAL_14_CF_ENERGY_IMPL_H

2
mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/nb14/dihedral_14_lj_atom_energy_impl.cu
View File

@ -87,7 +87,7 @@ void Dihedral14LJAtomEnergy(const int dihedral_14_numbers, const int atom_number
atom_numbers, uint_crd, uint_crd_with_LJ, LJtype, charge); atom_numbers, uint_crd, uint_crd_with_LJ, LJtype, charge);
VECTOR *boxlength = const_cast<VECTOR *>(reinterpret_cast<const VECTOR *>(boxlength_f)); VECTOR *boxlength = const_cast<VECTOR *>(reinterpret_cast<const VECTOR *>(boxlength_f));
Reset_List<<<ceilf(static_cast<float>(3. * atom_numbers) / 128), 128>>>(atom_numbers, ene, 0.); Reset_List<<<ceilf(static_cast<float>(3. * atom_numbers) / 128), 128, 0, stream>>>(atom_numbers, ene, 0.);
Dihedral14LJAtomEnergyKernel<<<block_per_grid, thread_per_block, 0, stream>>>( Dihedral14LJAtomEnergyKernel<<<block_per_grid, thread_per_block, 0, stream>>>(
dihedral_14_numbers, uint_crd_with_LJ, boxlength, a_14, b_14, lj_scale_factor, LJ_type_A, LJ_type_B, ene); dihedral_14_numbers, uint_crd_with_LJ, boxlength, a_14, b_14, lj_scale_factor, LJ_type_A, LJ_type_B, ene);

30
mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/nb14/dihedral_14_lj_cf_force_with_atom_energy_impl.cu
View File

@ -105,22 +105,23 @@ __global__ void Dihedral14LJCFForceWithAtomEnergyKernel(const int dihedral_14_nu
} }
void Dihedral14LJCFForceWithAtomEnergy(const int dihedral_14_numbers, const int atom_numbers, const int *uint_crd_f, void Dihedral14LJCFForceWithAtomEnergy(const int dihedral_14_numbers, const int atom_numbers, const int *uint_crd_f,
const int *LJtype, const float *charge, const float *boxlength_f, const int *LJtype, const float *charge, float *uint_crd_with_LJ_f,
const int *a_14, const int *b_14, const float *lj_scale_factor, const float *boxlength_f, const int *a_14, const int *b_14,
const float *cf_scale_factor, const float *LJ_type_A, const float *LJ_type_B, const float *lj_scale_factor, const float *cf_scale_factor,
float *frc_f, float *atom_energy, cudaStream_t stream) { const float *LJ_type_A, const float *LJ_type_B, float *frc_f, float *atom_energy,
cudaStream_t stream) {
size_t thread_per_block = 128; size_t thread_per_block = 128;
size_t block_per_grid = ceilf(static_cast<float>(atom_numbers) / 128); size_t block_per_grid = ceilf(static_cast<float>(atom_numbers) / 128);
UINT_VECTOR_LJ_TYPE *uint_crd_with_LJ = NULL;
Cuda_Malloc_Safely(reinterpret_cast<void **>(&uint_crd_with_LJ), sizeof(UINT_VECTOR_LJ_TYPE) * atom_numbers);
UNSIGNED_INT_VECTOR *uint_crd = UNSIGNED_INT_VECTOR *uint_crd =
const_cast<UNSIGNED_INT_VECTOR *>(reinterpret_cast<const UNSIGNED_INT_VECTOR *>(uint_crd_f)); const_cast<UNSIGNED_INT_VECTOR *>(reinterpret_cast<const UNSIGNED_INT_VECTOR *>(uint_crd_f));
UINT_VECTOR_LJ_TYPE *uint_crd_with_LJ = reinterpret_cast<UINT_VECTOR_LJ_TYPE *>(uint_crd_with_LJ_f);
Copy_Crd_To_New_Crd_Start<<<ceilf(static_cast<float>(atom_numbers) / 32), 32, 0, stream>>>( Copy_Crd_To_New_Crd_Start<<<ceilf(static_cast<float>(atom_numbers) / 32), 32, 0, stream>>>(
atom_numbers, uint_crd, uint_crd_with_LJ, LJtype, charge); atom_numbers, uint_crd, uint_crd_with_LJ, LJtype, charge);
Reset_List<<<ceilf(static_cast<float>(3. * atom_numbers) / 128), 128>>>(3 * atom_numbers, frc_f, 0.);
Reset_List<<<ceilf(static_cast<float>(3. * atom_numbers) / 128), 128>>>(atom_numbers, atom_energy, 0.); Reset_List<<<ceilf(static_cast<float>(3. * atom_numbers) / 128), 128, 0, stream>>>(3 * atom_numbers, frc_f, 0.);
Reset_List<<<ceilf(static_cast<float>(3. * atom_numbers) / 128), 128, 0, stream>>>(atom_numbers, atom_energy, 0.);
VECTOR *boxlength = const_cast<VECTOR *>(reinterpret_cast<const VECTOR *>(boxlength_f)); VECTOR *boxlength = const_cast<VECTOR *>(reinterpret_cast<const VECTOR *>(boxlength_f));
VECTOR *frc = const_cast<VECTOR *>(reinterpret_cast<const VECTOR *>(frc_f)); VECTOR *frc = const_cast<VECTOR *>(reinterpret_cast<const VECTOR *>(frc_f));
@ -133,8 +134,9 @@ void Dihedral14LJCFForceWithAtomEnergy(const int dihedral_14_numbers, const int
return; return;
} }
void Dihedral14LJForceWithDirectCF(const int dihedral_14_numbers, const int atom_numbers, const int *uint_crd_f, void Dihedral14LJCFForceWithAtomEnergy(const int dihedral_14_numbers, const int atom_numbers, const int *uint_crd_f,
const int *LJtype, const float *charge, const float *boxlength_f, const int *a_14, const int *LJtype, const float *charge, float *uint_crd_with_LJ_f,
const int *b_14, const float *lj_scale_factor, const float *cf_scale_factor, const float *boxlength_f, const int *a_14, const int *b_14,
const float *LJ_type_A, const float *LJ_type_B, float *frc, float *atom_energy, const float *lj_scale_factor, const float *cf_scale_factor,
cudaStream_t stream); const float *LJ_type_A, const float *LJ_type_B, float *frc_f, float *atom_energy,
cudaStream_t stream);

9
mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/nb14/dihedral_14_lj_cf_force_with_atom_energy_impl.cuh
View File

@ -20,8 +20,9 @@
#include "runtime/device/gpu/cuda_common.h" #include "runtime/device/gpu/cuda_common.h"
void Dihedral14LJCFForceWithAtomEnergy(const int dihedral_14_numbers, const int atom_numbers, const int *uint_crd_f, void Dihedral14LJCFForceWithAtomEnergy(const int dihedral_14_numbers, const int atom_numbers, const int *uint_crd_f,
const int *LJtype, const float *charge, const float *boxlength_f, const int *LJtype, const float *charge, float *uint_crd_with_LJ_f,
const int *a_14, const int *b_14, const float *lj_scale_factor, const float *boxlength_f, const int *a_14, const int *b_14,
const float *cf_scale_factor, const float *LJ_type_A, const float *LJ_type_B, const float *lj_scale_factor, const float *cf_scale_factor,
float *frc, float *atom_energy, cudaStream_t stream); const float *LJ_type_A, const float *LJ_type_B, float *frc_f, float *atom_energy,
cudaStream_t stream);
#endif // MINDSPORE_CCSRC_KERNEL_GPU_CUDA_IMPL_SPONGE_NB14_DIHEDRAL_14_LJ_CF_FORCE_WITH_ATOM_ENERGY_IMPL_H #endif // MINDSPORE_CCSRC_KERNEL_GPU_CUDA_IMPL_SPONGE_NB14_DIHEDRAL_14_LJ_CF_FORCE_WITH_ATOM_ENERGY_IMPL_H

19
mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/nb14/dihedral_14_lj_energy_impl.cu
View File

@ -72,20 +72,19 @@ __global__ void Dihedral14LJEnergyKernel(const int dihedral_14_numbers, const UI
} }
void Dihedral14LJEnergy(const int dihedral_14_numbers, const int atom_numbers, const int *uint_crd_f, const int *LJtype, void Dihedral14LJEnergy(const int dihedral_14_numbers, const int atom_numbers, const int *uint_crd_f, const int *LJtype,
const float *charge, const float *boxlength_f, const int *a_14, const int *b_14, const float *charge, float *uint_crd_with_LJ_f, const float *boxlength_f, const int *a_14,
const float *lj_scale_factor, const float *LJ_type_A, const float *LJ_type_B, float *ene, const int *b_14, const float *lj_scale_factor, const float *LJ_type_A, const float *LJ_type_B,
cudaStream_t stream) { float *ene, cudaStream_t stream) {
size_t thread_per_block = 128; size_t thread_per_block = 128;
size_t block_per_grid = ceilf(static_cast<float>(atom_numbers) / 128); size_t block_per_grid = ceilf(static_cast<float>(atom_numbers) / 128);
UINT_VECTOR_LJ_TYPE *uint_crd_with_LJ = NULL;
Cuda_Malloc_Safely(reinterpret_cast<void **>(&uint_crd_with_LJ), sizeof(UINT_VECTOR_LJ_TYPE) * atom_numbers);
UNSIGNED_INT_VECTOR *uint_crd = UNSIGNED_INT_VECTOR *uint_crd =
const_cast<UNSIGNED_INT_VECTOR *>(reinterpret_cast<const UNSIGNED_INT_VECTOR *>(uint_crd_f)); const_cast<UNSIGNED_INT_VECTOR *>(reinterpret_cast<const UNSIGNED_INT_VECTOR *>(uint_crd_f));
UINT_VECTOR_LJ_TYPE *uint_crd_with_LJ = reinterpret_cast<UINT_VECTOR_LJ_TYPE *>(uint_crd_with_LJ_f);
Copy_Crd_To_New_Crd_Start<<<ceilf(static_cast<float>(atom_numbers) / 32), 32, 0, stream>>>( Copy_Crd_To_New_Crd_Start<<<ceilf(static_cast<float>(atom_numbers) / 32), 32, 0, stream>>>(
atom_numbers, uint_crd, uint_crd_with_LJ, LJtype, charge); atom_numbers, uint_crd, uint_crd_with_LJ, LJtype, charge);
Reset_List<<<ceilf(static_cast<float>(3. * atom_numbers) / 128), 128>>>(dihedral_14_numbers, ene, 0.); Reset_List<<<ceilf(static_cast<float>(3. * atom_numbers) / 128), 128, 0, stream>>>(dihedral_14_numbers, ene, 0.);
VECTOR *boxlength = const_cast<VECTOR *>(reinterpret_cast<const VECTOR *>(boxlength_f)); VECTOR *boxlength = const_cast<VECTOR *>(reinterpret_cast<const VECTOR *>(boxlength_f));
Dihedral14LJEnergyKernel<<<block_per_grid, thread_per_block, 0, stream>>>( Dihedral14LJEnergyKernel<<<block_per_grid, thread_per_block, 0, stream>>>(
@ -97,6 +96,6 @@ void Dihedral14LJEnergy(const int dihedral_14_numbers, const int atom_numbers, c
} }
void Dihedral14LJEnergy(const int dihedral_14_numbers, const int atom_numbers, const int *uint_crd_f, const int *LJtype, void Dihedral14LJEnergy(const int dihedral_14_numbers, const int atom_numbers, const int *uint_crd_f, const int *LJtype,
const float *charge, const float *boxlength_f, const int *a_14, const int *b_14, const float *charge, float *uint_crd_with_LJ_f, const float *boxlength_f, const int *a_14,
const float *lj_scale_factor, const float *LJ_type_A, const float *LJ_type_B, float *ene, const int *b_14, const float *lj_scale_factor, const float *LJ_type_A, const float *LJ_type_B,
cudaStream_t stream); float *ene, cudaStream_t stream);

6
mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/nb14/dihedral_14_lj_energy_impl.cuh
View File

@ -20,8 +20,8 @@
#include "runtime/device/gpu/cuda_common.h" #include "runtime/device/gpu/cuda_common.h"
void Dihedral14LJEnergy(const int dihedral_14_numbers, const int atom_numbers, const int *uint_crd_f, const int *LJtype, void Dihedral14LJEnergy(const int dihedral_14_numbers, const int atom_numbers, const int *uint_crd_f, const int *LJtype,
const float *charge, const float *boxlength_f, const int *a_14, const int *b_14, const float *charge, float *uint_crd_with_LJ_f, const float *boxlength_f, const int *a_14,
const float *lj_scale_factor, const float *LJ_type_A, const float *LJ_type_B, float *ene, const int *b_14, const float *lj_scale_factor, const float *LJ_type_A, const float *LJ_type_B,
cudaStream_t stream); float *ene, cudaStream_t stream);
#endif // MINDSPORE_CCSRC_KERNEL_GPU_CUDA_IMPL_SPONGE_NB14_DIHEDRAL_14_LJ_ENERGY_IMPL_H #endif // MINDSPORE_CCSRC_KERNEL_GPU_CUDA_IMPL_SPONGE_NB14_DIHEDRAL_14_LJ_ENERGY_IMPL_H

2
mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/nb14/dihedral_14_lj_force_with_direct_cf_impl.cu
View File

@ -108,7 +108,7 @@ void Dihedral14LJForceWithDirectCF(const int dihedral_14_numbers, const int atom
atom_numbers, uint_crd, uint_crd_with_LJ, LJtype, charge); atom_numbers, uint_crd, uint_crd_with_LJ, LJtype, charge);
cudaStreamSynchronize(stream); cudaStreamSynchronize(stream);
VECTOR *boxlength = const_cast<VECTOR *>(reinterpret_cast<const VECTOR *>(boxlength_f)); VECTOR *boxlength = const_cast<VECTOR *>(reinterpret_cast<const VECTOR *>(boxlength_f));
Reset_List<<<ceilf(static_cast<float>(3. * atom_numbers) / 128), 128>>>(3 * atom_numbers, frc_f, 0.); Reset_List<<<ceilf(static_cast<float>(3. * atom_numbers) / 128), 128, 0, stream>>>(3 * atom_numbers, frc_f, 0.);
VECTOR *frc = const_cast<VECTOR *>(reinterpret_cast<const VECTOR *>(frc_f)); VECTOR *frc = const_cast<VECTOR *>(reinterpret_cast<const VECTOR *>(frc_f));
Dihedral14LJForceWithDirectCFKernel<<<block_per_grid, thread_per_block, 0, stream>>>( Dihedral14LJForceWithDirectCFKernel<<<block_per_grid, thread_per_block, 0, stream>>>(

3
mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/nvtit/md_iteration_leap_frog_impl.cu
View File

@ -97,9 +97,6 @@ void MDIterationLeapFrog(const int float4_numbers, const int atom_numbers, const
const float exp_gamma, const int is_max_velocity, const float max_velocity, const float exp_gamma, const int is_max_velocity, const float max_velocity,
const float *d_mass_inverse, const float *d_sqrt_mass, float *vel_f, float *crd_f, const float *d_mass_inverse, const float *d_sqrt_mass, float *vel_f, float *crd_f,
float *frc_f, float *acc_f, cudaStream_t stream) { float *frc_f, float *acc_f, cudaStream_t stream) {
Reset_List<<<ceilf(static_cast<float>(3. * atom_numbers) / 128), 128>>>(3 * atom_numbers, vel_f, 0.);
Reset_List<<<ceilf(static_cast<float>(3. * atom_numbers) / 128), 128>>>(3 * atom_numbers, crd_f, 0.);
Reset_List<<<ceilf(static_cast<float>(3. * atom_numbers) / 128), 128>>>(3 * atom_numbers, frc_f, 0.);
Reset_List<<<ceilf(static_cast<float>(3. * atom_numbers) / 128), 128>>>(3 * atom_numbers, acc_f, 0.); Reset_List<<<ceilf(static_cast<float>(3. * atom_numbers) / 128), 128>>>(3 * atom_numbers, acc_f, 0.);
VECTOR *frc = const_cast<VECTOR *>(reinterpret_cast<const VECTOR *>(frc_f)); VECTOR *frc = const_cast<VECTOR *>(reinterpret_cast<const VECTOR *>(frc_f));

2
mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/pme/pme_excluded_force_impl.cu
View File

@ -86,7 +86,7 @@ __global__ void PME_Excluded_Force_Correction(const int atom_numbers, const UNSI
void PMEExcludedForce(const int atom_numbers, const float pme_beta, const int *uint_crd_f, const float *sacler_f, void PMEExcludedForce(const int atom_numbers, const float pme_beta, const int *uint_crd_f, const float *sacler_f,
const float *charge, const int *excluded_list_start, const int *excluded_list, const float *charge, const int *excluded_list_start, const int *excluded_list,
const int *excluded_atom_numbers, float *frc_f, cudaStream_t stream) { const int *excluded_atom_numbers, float *frc_f, cudaStream_t stream) {
Reset_List<<<ceilf(static_cast<float>(3. * atom_numbers) / 128), 128>>>(3 * atom_numbers, frc_f, 0.); Reset_List<<<ceilf(static_cast<float>(3. * atom_numbers) / 128), 128, 0, stream>>>(3 * atom_numbers, frc_f, 0.);
UNSIGNED_INT_VECTOR *uint_crd = UNSIGNED_INT_VECTOR *uint_crd =
const_cast<UNSIGNED_INT_VECTOR *>(reinterpret_cast<const UNSIGNED_INT_VECTOR *>(uint_crd_f)); const_cast<UNSIGNED_INT_VECTOR *>(reinterpret_cast<const UNSIGNED_INT_VECTOR *>(uint_crd_f));
VECTOR *frc = reinterpret_cast<VECTOR *>(frc_f); VECTOR *frc = reinterpret_cast<VECTOR *>(frc_f);

2
mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/pme/pme_reciprocal_force_impl.cu
View File

@ -75,7 +75,7 @@ void PMEReciprocalForce(int fftx, int ffty, int fftz, int atom_numbers, float be
float *pme_frxyz, float *PME_Q, float *pme_fq, int *PME_atom_near, int *pme_kxyz, float *pme_frxyz, float *PME_Q, float *pme_fq, int *PME_atom_near, int *pme_kxyz,
const float *box_length_f, const int *uint_crd_f, const float *charge, float *force, const float *box_length_f, const int *uint_crd_f, const float *charge, float *force,
cudaStream_t stream) { cudaStream_t stream) {
Reset_List<<<ceilf(static_cast<float>(3. * atom_numbers) / 128), 128>>>(3 * atom_numbers, force, 0.); Reset_List<<<ceilf(static_cast<float>(3. * atom_numbers) / 128), 128, 0, stream>>>(3 * atom_numbers, force, 0.);
UNSIGNED_INT_VECTOR *uint_crd = UNSIGNED_INT_VECTOR *uint_crd =
const_cast<UNSIGNED_INT_VECTOR *>(reinterpret_cast<const UNSIGNED_INT_VECTOR *>(uint_crd_f)); const_cast<UNSIGNED_INT_VECTOR *>(reinterpret_cast<const UNSIGNED_INT_VECTOR *>(uint_crd_f));
UNSIGNED_INT_VECTOR *PME_uxyz = reinterpret_cast<UNSIGNED_INT_VECTOR *>(pme_uxyz); UNSIGNED_INT_VECTOR *PME_uxyz = reinterpret_cast<UNSIGNED_INT_VECTOR *>(pme_uxyz);

16
mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/nb14/dihedral_14_cf_energy_kernel.h
View File

@ -64,7 +64,7 @@ class Dihedral14CFEnergyGpuKernel : public GpuKernel {
const std::vector<size_t> &GetOutputSizeList() const override { return output_size_list_; } const std::vector<size_t> &GetOutputSizeList() const override { return output_size_list_; }
const std::vector<size_t> &GetWorkspaceSizeList() const override { return workspace_size_list_; } const std::vector<size_t> &GetWorkspaceSizeList() const override { return workspace_size_list_; }
bool Launch(const std::vector<AddressPtr> &inputs, const std::vector<AddressPtr> &, bool Launch(const std::vector<AddressPtr> &inputs, const std::vector<AddressPtr> &workspace,
const std::vector<AddressPtr> &outputs, void *stream_ptr) override { const std::vector<AddressPtr> &outputs, void *stream_ptr) override {
auto uint_crd_f = GetDeviceAddress<const T1>(inputs, 0); auto uint_crd_f = GetDeviceAddress<const T1>(inputs, 0);
auto LJtype = GetDeviceAddress<const T1>(inputs, 1); auto LJtype = GetDeviceAddress<const T1>(inputs, 1);
@ -74,9 +74,10 @@ class Dihedral14CFEnergyGpuKernel : public GpuKernel {
auto b_14 = GetDeviceAddress<const T1>(inputs, 5); auto b_14 = GetDeviceAddress<const T1>(inputs, 5);
auto cf_scale_factor = GetDeviceAddress<T>(inputs, 6); auto cf_scale_factor = GetDeviceAddress<T>(inputs, 6);
auto ene = GetDeviceAddress<T>(outputs, 0); auto ene = GetDeviceAddress<T>(outputs, 0);
auto uint_crd_with_LJ = GetDeviceAddress<T>(workspace, 0);
Dihedral14CFEnergy(dihedral_14_numbers, atom_numbers, uint_crd_f, LJtype, charge, boxlength_f, a_14, b_14, Dihedral14CFEnergy(dihedral_14_numbers, atom_numbers, uint_crd_f, LJtype, charge, uint_crd_with_LJ, boxlength_f,
cf_scale_factor, ene, reinterpret_cast<cudaStream_t>(stream_ptr)); a_14, b_14, cf_scale_factor, ene, reinterpret_cast<cudaStream_t>(stream_ptr));
return true; return true;
} }
@ -90,7 +91,7 @@ class Dihedral14CFEnergyGpuKernel : public GpuKernel {
input_size_list_.push_back(ele_a_14 * sizeof(T1)); input_size_list_.push_back(ele_a_14 * sizeof(T1));
input_size_list_.push_back(ele_b_14 * sizeof(T1)); input_size_list_.push_back(ele_b_14 * sizeof(T1));
input_size_list_.push_back(ele_cf_scale_factor * sizeof(T)); input_size_list_.push_back(ele_cf_scale_factor * sizeof(T));
workspace_size_list_.push_back(atom_numbers * sizeof(UINT_VECTOR_LJ_TYPE));
output_size_list_.push_back(atom_numbers * sizeof(T)); output_size_list_.push_back(atom_numbers * sizeof(T));
} }
@ -108,6 +109,13 @@ class Dihedral14CFEnergyGpuKernel : public GpuKernel {
std::vector<size_t> workspace_size_list_; std::vector<size_t> workspace_size_list_;
int dihedral_14_numbers; int dihedral_14_numbers;
int atom_numbers; int atom_numbers;
struct UINT_VECTOR_LJ_TYPE {
unsigned int uint_x;
unsigned int uint_y;
unsigned int uint_z;
int LJ_type;
float charge;
};
}; };
} // namespace kernel } // namespace kernel
} // namespace mindspore } // namespace mindspore

18
mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/nb14/dihedral_14_lj_cf_force_with_atom_energy_kernel.h
View File

@ -70,7 +70,7 @@ class Dihedral14LJCFForceWithAtomEnergyGpuKernel : public GpuKernel {
const std::vector<size_t> &GetOutputSizeList() const override { return output_size_list_; } const std::vector<size_t> &GetOutputSizeList() const override { return output_size_list_; }
const std::vector<size_t> &GetWorkspaceSizeList() const override { return workspace_size_list_; } const std::vector<size_t> &GetWorkspaceSizeList() const override { return workspace_size_list_; }
bool Launch(const std::vector<AddressPtr> &inputs, const std::vector<AddressPtr> &, bool Launch(const std::vector<AddressPtr> &inputs, const std::vector<AddressPtr> &workspace,
const std::vector<AddressPtr> &outputs, void *stream_ptr) override { const std::vector<AddressPtr> &outputs, void *stream_ptr) override {
auto uint_crd_f = GetDeviceAddress<const T1>(inputs, 0); auto uint_crd_f = GetDeviceAddress<const T1>(inputs, 0);
auto LJtype = GetDeviceAddress<const T1>(inputs, 1); auto LJtype = GetDeviceAddress<const T1>(inputs, 1);
@ -85,9 +85,11 @@ class Dihedral14LJCFForceWithAtomEnergyGpuKernel : public GpuKernel {
auto frc_f = GetDeviceAddress<T>(outputs, 0); auto frc_f = GetDeviceAddress<T>(outputs, 0);
auto atom_energy = GetDeviceAddress<T>(outputs, 1); auto atom_energy = GetDeviceAddress<T>(outputs, 1);
Dihedral14LJCFForceWithAtomEnergy(dihedral_14_numbers, atom_numbers, uint_crd_f, LJtype, charge, boxlength_f, a_14, auto uint_crd_with_LJ = GetDeviceAddress<T>(workspace, 0);
b_14, lj_scale_factor, cf_scale_factor, LJ_type_A, LJ_type_B, frc_f, atom_energy,
reinterpret_cast<cudaStream_t>(stream_ptr)); Dihedral14LJCFForceWithAtomEnergy(dihedral_14_numbers, atom_numbers, uint_crd_f, LJtype, charge, uint_crd_with_LJ,
boxlength_f, a_14, b_14, lj_scale_factor, cf_scale_factor, LJ_type_A, LJ_type_B,
frc_f, atom_energy, reinterpret_cast<cudaStream_t>(stream_ptr));
return true; return true;
} }
@ -104,6 +106,7 @@ class Dihedral14LJCFForceWithAtomEnergyGpuKernel : public GpuKernel {
input_size_list_.push_back(ele_cf_scale_factor * sizeof(T)); input_size_list_.push_back(ele_cf_scale_factor * sizeof(T));
input_size_list_.push_back(ele_LJ_type_A * sizeof(T)); input_size_list_.push_back(ele_LJ_type_A * sizeof(T));
input_size_list_.push_back(ele_LJ_type_B * sizeof(T)); input_size_list_.push_back(ele_LJ_type_B * sizeof(T));
workspace_size_list_.push_back(atom_numbers * sizeof(UINT_VECTOR_LJ_TYPE));
output_size_list_.push_back(3 * atom_numbers * sizeof(T)); output_size_list_.push_back(3 * atom_numbers * sizeof(T));
output_size_list_.push_back(atom_numbers * sizeof(T)); output_size_list_.push_back(atom_numbers * sizeof(T));
@ -126,6 +129,13 @@ class Dihedral14LJCFForceWithAtomEnergyGpuKernel : public GpuKernel {
std::vector<size_t> workspace_size_list_; std::vector<size_t> workspace_size_list_;
int dihedral_14_numbers; int dihedral_14_numbers;
int atom_numbers; int atom_numbers;
struct UINT_VECTOR_LJ_TYPE {
unsigned int uint_x;
unsigned int uint_y;
unsigned int uint_z;
int LJ_type;
float charge;
};
}; };
} // namespace kernel } // namespace kernel
} // namespace mindspore } // namespace mindspore

16
mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/nb14/dihedral_14_lj_energy_kernel.h
View File

@ -68,7 +68,7 @@ class Dihedral14LJEnergyGpuKernel : public GpuKernel {
const std::vector<size_t> &GetOutputSizeList() const override { return output_size_list_; } const std::vector<size_t> &GetOutputSizeList() const override { return output_size_list_; }
const std::vector<size_t> &GetWorkspaceSizeList() const override { return workspace_size_list_; } const std::vector<size_t> &GetWorkspaceSizeList() const override { return workspace_size_list_; }
bool Launch(const std::vector<AddressPtr> &inputs, const std::vector<AddressPtr> &, bool Launch(const std::vector<AddressPtr> &inputs, const std::vector<AddressPtr> &workspace,
const std::vector<AddressPtr> &outputs, void *stream_ptr) override { const std::vector<AddressPtr> &outputs, void *stream_ptr) override {
auto uint_crd_f = GetDeviceAddress<const T1>(inputs, 0); auto uint_crd_f = GetDeviceAddress<const T1>(inputs, 0);
auto LJtype = GetDeviceAddress<const T1>(inputs, 1); auto LJtype = GetDeviceAddress<const T1>(inputs, 1);
@ -80,9 +80,11 @@ class Dihedral14LJEnergyGpuKernel : public GpuKernel {
auto LJ_type_A = GetDeviceAddress<T>(inputs, 7); auto LJ_type_A = GetDeviceAddress<T>(inputs, 7);
auto LJ_type_B = GetDeviceAddress<T>(inputs, 8); auto LJ_type_B = GetDeviceAddress<T>(inputs, 8);
auto ene = GetDeviceAddress<T>(outputs, 0); auto ene = GetDeviceAddress<T>(outputs, 0);
auto uint_crd_with_LJ = GetDeviceAddress<T>(workspace, 0);
Dihedral14LJEnergy(dihedral_14_numbers, atom_numbers, uint_crd_f, LJtype, charge, boxlength_f, a_14, b_14, Dihedral14LJEnergy(dihedral_14_numbers, atom_numbers, uint_crd_f, LJtype, charge, uint_crd_with_LJ, boxlength_f,
lj_scale_factor, LJ_type_A, LJ_type_B, ene, reinterpret_cast<cudaStream_t>(stream_ptr)); a_14, b_14, lj_scale_factor, LJ_type_A, LJ_type_B, ene,
reinterpret_cast<cudaStream_t>(stream_ptr));
return true; return true;
} }
@ -98,6 +100,7 @@ class Dihedral14LJEnergyGpuKernel : public GpuKernel {
input_size_list_.push_back(ele_lj_scale_factor * sizeof(T)); input_size_list_.push_back(ele_lj_scale_factor * sizeof(T));
input_size_list_.push_back(ele_LJ_type_A * sizeof(T)); input_size_list_.push_back(ele_LJ_type_A * sizeof(T));
input_size_list_.push_back(ele_LJ_type_B * sizeof(T)); input_size_list_.push_back(ele_LJ_type_B * sizeof(T));
workspace_size_list_.push_back(atom_numbers * sizeof(UINT_VECTOR_LJ_TYPE));
output_size_list_.push_back(atom_numbers * sizeof(T)); output_size_list_.push_back(atom_numbers * sizeof(T));
} }
@ -118,6 +121,13 @@ class Dihedral14LJEnergyGpuKernel : public GpuKernel {
std::vector<size_t> workspace_size_list_; std::vector<size_t> workspace_size_list_;
int dihedral_14_numbers; int dihedral_14_numbers;
int atom_numbers; int atom_numbers;
struct UINT_VECTOR_LJ_TYPE {
unsigned int uint_x;
unsigned int uint_y;
unsigned int uint_z;
int LJ_type;
float charge;
};
}; };
} // namespace kernel } // namespace kernel
} // namespace mindspore } // namespace mindspore

3
mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/pme/pme_energy_kernel.h
View File

@ -93,7 +93,7 @@ class PMEEnergyGpuKernel : public GpuKernel {
input_size_list_.push_back(atom_numbers * sizeof(VECTOR)); input_size_list_.push_back(atom_numbers * sizeof(VECTOR));
input_size_list_.push_back(atom_numbers * sizeof(T1)); input_size_list_.push_back(atom_numbers * sizeof(T1));
input_size_list_.push_back(atom_numbers * sizeof(T1)); input_size_list_.push_back(excluded_numbers * sizeof(T1));
input_size_list_.push_back(atom_numbers * sizeof(T1)); input_size_list_.push_back(atom_numbers * sizeof(T1));
workspace_size_list_.push_back(atom_numbers * sizeof(UNSIGNED_INT_VECTOR)); workspace_size_list_.push_back(atom_numbers * sizeof(UNSIGNED_INT_VECTOR));
@ -118,6 +118,7 @@ class PMEEnergyGpuKernel : public GpuKernel {
std::vector<size_t> output_size_list_; std::vector<size_t> output_size_list_;
std::vector<size_t> workspace_size_list_; std::vector<size_t> workspace_size_list_;
int atom_numbers; int atom_numbers;
int excluded_numbers = 2719;
int max_nl_numbers = 800; int max_nl_numbers = 800;
int fftx; int fftx;
int ffty; int ffty;

3
mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/pme/pme_excluded_force_kernel.h
View File

@ -65,7 +65,7 @@ class PMEExcludedForceGpuKernel : public GpuKernel {
input_size_list_.push_back(atom_numbers * sizeof(VECTOR)); input_size_list_.push_back(atom_numbers * sizeof(VECTOR));
input_size_list_.push_back(atom_numbers * sizeof(T)); input_size_list_.push_back(atom_numbers * sizeof(T));
input_size_list_.push_back(atom_numbers * sizeof(T1)); input_size_list_.push_back(atom_numbers * sizeof(T1));
input_size_list_.push_back(atom_numbers * sizeof(T1)); input_size_list_.push_back(excluded_numbers * sizeof(T1));
input_size_list_.push_back(atom_numbers * sizeof(T1)); input_size_list_.push_back(atom_numbers * sizeof(T1));
output_size_list_.push_back(atom_numbers * 3 * sizeof(T)); output_size_list_.push_back(atom_numbers * 3 * sizeof(T));
@ -77,6 +77,7 @@ class PMEExcludedForceGpuKernel : public GpuKernel {
std::vector<size_t> output_size_list_; std::vector<size_t> output_size_list_;
std::vector<size_t> workspace_size_list_; std::vector<size_t> workspace_size_list_;
int atom_numbers; int atom_numbers;
int excluded_numbers = 2719;
float beta; float beta;
struct VECTOR { struct VECTOR {
float x; float x;

435
mindspore/ops/operations/sponge_ops.py
View File

@ -29,6 +29,9 @@ class BondForce(PrimitiveWithInfer):
.. math:: .. math::
dr = (x_1-x_2, y_1-y_2, z_1-z_2) dr = (x_1-x_2, y_1-y_2, z_1-z_2)
.. math::
F = (F_x, F_y, F_z) = 2*k*(1 - r_0/|dr|)*dr F = (F_x, F_y, F_z) = 2*k*(1 - r_0/|dr|)*dr
Args: Args:
@ -53,8 +56,8 @@ class BondForce(PrimitiveWithInfer):
@prim_attr_register @prim_attr_register
def __init__(self, bond_numbers, atom_numbers): def __init__(self, bond_numbers, atom_numbers):
assert isinstance(bond_numbers, int) validator.check_value_type('bond_numbers', bond_numbers, (int), self.name)
assert isinstance(atom_numbers, int) validator.check_value_type('atom_numbers', atom_numbers, (int), self.name)
self.bond_numbers = bond_numbers self.bond_numbers = bond_numbers
self.atom_numbers = atom_numbers self.atom_numbers = atom_numbers
self.add_prim_attr('bond_numbers', self.bond_numbers) self.add_prim_attr('bond_numbers', self.bond_numbers)
@ -66,8 +69,8 @@ class BondForce(PrimitiveWithInfer):
cls_name = self.name cls_name = self.name
N = self.atom_numbers N = self.atom_numbers
M = self.bond_numbers M = self.bond_numbers
validator.check_int(uint_crd_f_shape[0], N, Rel.EQ, "uint_crd_f", cls_name) validator.check_int(uint_crd_f_shape[0], N, Rel.EQ, "uint_crd_f[0]", cls_name)
validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f", cls_name) validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f[1]", cls_name)
validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f", cls_name) validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f", cls_name)
validator.check_int(atom_a_shape[0], M, Rel.EQ, "uint_crd_f", cls_name) validator.check_int(atom_a_shape[0], M, Rel.EQ, "uint_crd_f", cls_name)
validator.check_int(atom_b_shape[0], M, Rel.EQ, "atom_b", cls_name) validator.check_int(atom_b_shape[0], M, Rel.EQ, "atom_b", cls_name)
@ -93,6 +96,9 @@ class BondEnergy(PrimitiveWithInfer):
.. math:: .. math::
dr = (x_1-x_2, y_1-y_2, z_1-z_2) dr = (x_1-x_2, y_1-y_2, z_1-z_2)
.. math::
E = k*(|dr| - r_0)^2 E = k*(|dr| - r_0)^2
Args: Args:
@ -117,8 +123,8 @@ class BondEnergy(PrimitiveWithInfer):
@prim_attr_register @prim_attr_register
def __init__(self, bond_numbers, atom_numbers): def __init__(self, bond_numbers, atom_numbers):
assert isinstance(bond_numbers, int) validator.check_value_type('bond_numbers', bond_numbers, (int), self.name)
assert isinstance(atom_numbers, int) validator.check_value_type('atom_numbers', atom_numbers, (int), self.name)
self.bond_numbers = bond_numbers self.bond_numbers = bond_numbers
self.atom_numbers = atom_numbers self.atom_numbers = atom_numbers
self.add_prim_attr('bond_numbers', self.bond_numbers) self.add_prim_attr('bond_numbers', self.bond_numbers)
@ -130,8 +136,8 @@ class BondEnergy(PrimitiveWithInfer):
cls_name = self.name cls_name = self.name
N = self.atom_numbers N = self.atom_numbers
M = self.bond_numbers M = self.bond_numbers
validator.check_int(uint_crd_f_shape[0], N, Rel.EQ, "uint_crd_f", cls_name) validator.check_int(uint_crd_f_shape[0], N, Rel.EQ, "uint_crd_f[0]", cls_name)
validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f", cls_name) validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f[1]", cls_name)
validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f", cls_name) validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f", cls_name)
validator.check_int(atom_a_shape[0], M, Rel.EQ, "uint_crd_f", cls_name) validator.check_int(atom_a_shape[0], M, Rel.EQ, "uint_crd_f", cls_name)
validator.check_int(atom_b_shape[0], M, Rel.EQ, "atom_b", cls_name) validator.check_int(atom_b_shape[0], M, Rel.EQ, "atom_b", cls_name)
@ -179,8 +185,8 @@ class BondAtomEnergy(PrimitiveWithInfer):
@prim_attr_register @prim_attr_register
def __init__(self, bond_numbers, atom_numbers): def __init__(self, bond_numbers, atom_numbers):
assert isinstance(bond_numbers, int) validator.check_value_type('bond_numbers', bond_numbers, (int), self.name)
assert isinstance(atom_numbers, int) validator.check_value_type('atom_numbers', atom_numbers, (int), self.name)
self.bond_numbers = bond_numbers self.bond_numbers = bond_numbers
self.atom_numbers = atom_numbers self.atom_numbers = atom_numbers
self.add_prim_attr('bond_numbers', self.bond_numbers) self.add_prim_attr('bond_numbers', self.bond_numbers)
@ -192,8 +198,8 @@ class BondAtomEnergy(PrimitiveWithInfer):
cls_name = self.name cls_name = self.name
N = self.atom_numbers N = self.atom_numbers
M = self.bond_numbers M = self.bond_numbers
validator.check_int(uint_crd_f_shape[0], N, Rel.EQ, "uint_crd_f", cls_name) validator.check_int(uint_crd_f_shape[0], N, Rel.EQ, "uint_crd_f[0]", cls_name)
validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f", cls_name) validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f[1]", cls_name)
validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f", cls_name) validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f", cls_name)
validator.check_int(atom_a_shape[0], M, Rel.EQ, "uint_crd_f", cls_name) validator.check_int(atom_a_shape[0], M, Rel.EQ, "uint_crd_f", cls_name)
validator.check_int(atom_b_shape[0], M, Rel.EQ, "atom_b", cls_name) validator.check_int(atom_b_shape[0], M, Rel.EQ, "atom_b", cls_name)
@ -240,8 +246,8 @@ class BondForceWithAtomEnergy(PrimitiveWithInfer):
@prim_attr_register @prim_attr_register
def __init__(self, bond_numbers, atom_numbers): def __init__(self, bond_numbers, atom_numbers):
assert isinstance(bond_numbers, int) validator.check_value_type('bond_numbers', bond_numbers, (int), self.name)
assert isinstance(atom_numbers, int) validator.check_value_type('atom_numbers', atom_numbers, (int), self.name)
self.bond_numbers = bond_numbers self.bond_numbers = bond_numbers
self.atom_numbers = atom_numbers self.atom_numbers = atom_numbers
self.add_prim_attr('bond_numbers', self.bond_numbers) self.add_prim_attr('bond_numbers', self.bond_numbers)
@ -253,8 +259,8 @@ class BondForceWithAtomEnergy(PrimitiveWithInfer):
cls_name = self.name cls_name = self.name
N = self.atom_numbers N = self.atom_numbers
M = self.bond_numbers M = self.bond_numbers
validator.check_int(uint_crd_f_shape[0], N, Rel.EQ, "uint_crd_f", cls_name) validator.check_int(uint_crd_f_shape[0], N, Rel.EQ, "uint_crd_f[0]", cls_name)
validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f", cls_name) validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f[1]", cls_name)
validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f", cls_name) validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f", cls_name)
validator.check_int(atom_a_shape[0], M, Rel.EQ, "uint_crd_f", cls_name) validator.check_int(atom_a_shape[0], M, Rel.EQ, "uint_crd_f", cls_name)
validator.check_int(atom_b_shape[0], M, Rel.EQ, "atom_b", cls_name) validator.check_int(atom_b_shape[0], M, Rel.EQ, "atom_b", cls_name)
@ -286,6 +292,9 @@ class BondForceWithAtomVirial(PrimitiveWithInfer):
.. math:: .. math::
dr = (x_1-x_2, y_1-y_2, z_1-z_2) dr = (x_1-x_2, y_1-y_2, z_1-z_2)
.. math::
virial = |dr|*(|dr| - r_0)*k virial = |dr|*(|dr| - r_0)*k
Args: Args:
@ -311,8 +320,8 @@ class BondForceWithAtomVirial(PrimitiveWithInfer):
@prim_attr_register @prim_attr_register
def __init__(self, bond_numbers, atom_numbers): def __init__(self, bond_numbers, atom_numbers):
assert isinstance(bond_numbers, int) validator.check_value_type('bond_numbers', bond_numbers, (int), self.name)
assert isinstance(atom_numbers, int) validator.check_value_type('atom_numbers', atom_numbers, (int), self.name)
self.bond_numbers = bond_numbers self.bond_numbers = bond_numbers
self.atom_numbers = atom_numbers self.atom_numbers = atom_numbers
self.add_prim_attr('bond_numbers', self.bond_numbers) self.add_prim_attr('bond_numbers', self.bond_numbers)
@ -324,8 +333,8 @@ class BondForceWithAtomVirial(PrimitiveWithInfer):
cls_name = self.name cls_name = self.name
N = self.atom_numbers N = self.atom_numbers
M = self.bond_numbers M = self.bond_numbers
validator.check_int(uint_crd_f_shape[0], N, Rel.EQ, "uint_crd_f", cls_name) validator.check_int(uint_crd_f_shape[0], N, Rel.EQ, "uint_crd_f[0]", cls_name)
validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f", cls_name) validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f[1]", cls_name)
validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f", cls_name) validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f", cls_name)
validator.check_int(atom_a_shape[0], M, Rel.EQ, "uint_crd_f", cls_name) validator.check_int(atom_a_shape[0], M, Rel.EQ, "uint_crd_f", cls_name)
validator.check_int(atom_b_shape[0], M, Rel.EQ, "atom_b", cls_name) validator.check_int(atom_b_shape[0], M, Rel.EQ, "atom_b", cls_name)
@ -355,25 +364,36 @@ class DihedralForce(PrimitiveWithInfer):
.. math:: .. math::
dr_{ab} = (x_b-x_a, y_b-y_a, z_b-z_a) dr_{ab} = (x_b-x_a, y_b-y_a, z_b-z_a)
.. math::
dr_{cb} = (x_b-x_c, y_b-y_c, z_b-z_c) dr_{cb} = (x_b-x_c, y_b-y_c, z_b-z_c)
.. math::
dr_{cd} = (x_d-x_c, y_d-y_c, z_d-z_c) dr_{cd} = (x_d-x_c, y_d-y_c, z_d-z_c)
.. math::
r1 = dr_{ab}*dr_{cb} r1 = dr_{ab}*dr_{cb}
.. math::
r2 = dr_{cd}*dr_{cb} r2 = dr_{cd}*dr_{cb}
.. math::
phi = pi - sign(inner_product(r1*r2), dr_{cb}) phi = pi - sign(inner_product(r1*r2), dr_{cb})
* arccos(inner_product(r1, r2)/|r1|/|r2|) * arccos(inner_product(r1, r2)/|r1|/|r2|)
.. math::
dEdphi = n*phi*(k*cos(phi_0)*sin(n*phi) - k*sin(phi_0)*cos(n*phi))/sin(phi) dEdphi = n*phi*(k*cos(phi_0)*sin(n*phi) - k*sin(phi_0)*cos(n*phi))/sin(phi)
.. math::
dphidr1 = r2/|r1|/|r2| + cos(phi)/|r1|^2*r1 dphidr1 = r2/|r1|/|r2| + cos(phi)/|r1|^2*r1
.. math::
dphidr2 = r1/|r1|/|r2| + cos(phi)/|r2|^2*r2 dphidr2 = r1/|r1|/|r2| + cos(phi)/|r2|^2*r2
.. math::
dEdra = dEdphi * dr_{cb} * dphidr1 dEdra = dEdphi * dr_{cb} * dphidr1
.. math::
dEdrd = dEdphi * dphi_dr2 * dr_{cb} dEdrd = dEdphi * dphi_dr2 * dr_{cb}
.. math::
dEdrjpart = dEdphi * ((dr_{ab} * dphidr1) + (dr_{cd} * dphidr2)) dEdrjpart = dEdphi * ((dr_{ab} * dphidr1) + (dr_{cd} * dphidr2))
.. math::
F_a = dEdri F_a = dEdri
.. math::
F_b = dEdrjpart - dEdri F_b = dEdrjpart - dEdri
.. math::
F_c = - dEdrl - dEdrjpart F_c = - dEdrl - dEdrjpart
.. math::
F_d = dEdrl F_d = dEdrl
Args: Args:
@ -405,7 +425,7 @@ class DihedralForce(PrimitiveWithInfer):
@prim_attr_register @prim_attr_register
def __init__(self, dihedral_numbers): def __init__(self, dihedral_numbers):
assert isinstance(dihedral_numbers, int) validator.check_value_type('dihedral_numbers', dihedral_numbers, (int), self.name)
self.dihedral_numbers = dihedral_numbers self.dihedral_numbers = dihedral_numbers
self.init_prim_io_names(inputs=['uint_crd_f', 'scaler_f', 'atom_a', 'atom_b', 'atom_c', 'atom_d', 'ipn', 'pk', self.init_prim_io_names(inputs=['uint_crd_f', 'scaler_f', 'atom_a', 'atom_b', 'atom_c', 'atom_d', 'ipn', 'pk',
'gamc', 'gams', 'pn'], 'gamc', 'gams', 'pn'],
@ -415,8 +435,8 @@ class DihedralForce(PrimitiveWithInfer):
def infer_shape(self, uint_crd_f_shape, scaler_f_shape, atom_a_shape, atom_b_shape, atom_c_shape, atom_d_shape, def infer_shape(self, uint_crd_f_shape, scaler_f_shape, atom_a_shape, atom_b_shape, atom_c_shape, atom_d_shape,
ipn_shape, pk_shape, gamc_shape, gams_shape, pn_shape): ipn_shape, pk_shape, gamc_shape, gams_shape, pn_shape):
cls_name = self.name cls_name = self.name
M = atom_a_shape[0] M = self.dihedral_numbers
validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f", cls_name) validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f[1]", cls_name)
validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f", cls_name) validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f", cls_name)
validator.check_int(atom_a_shape[0], M, Rel.EQ, "atom_a", cls_name) validator.check_int(atom_a_shape[0], M, Rel.EQ, "atom_a", cls_name)
validator.check_int(atom_b_shape[0], M, Rel.EQ, "atom_b", cls_name) validator.check_int(atom_b_shape[0], M, Rel.EQ, "atom_b", cls_name)
@ -485,7 +505,7 @@ class DihedralEnergy(PrimitiveWithInfer):
@prim_attr_register @prim_attr_register
def __init__(self, dihedral_numbers): def __init__(self, dihedral_numbers):
assert isinstance(dihedral_numbers, int) validator.check_value_type('dihedral_numbers', dihedral_numbers, (int), self.name)
self.dihedral_numbers = dihedral_numbers self.dihedral_numbers = dihedral_numbers
self.init_prim_io_names(inputs=['uint_crd_f', 'scaler_f', 'atom_a', 'atom_b', 'atom_c', 'atom_d', 'ipn', 'pk', self.init_prim_io_names(inputs=['uint_crd_f', 'scaler_f', 'atom_a', 'atom_b', 'atom_c', 'atom_d', 'ipn', 'pk',
'gamc', 'gams', 'pn'], 'gamc', 'gams', 'pn'],
@ -495,8 +515,8 @@ class DihedralEnergy(PrimitiveWithInfer):
def infer_shape(self, uint_crd_f_shape, scaler_f_shape, atom_a_shape, atom_b_shape, atom_c_shape, atom_d_shape, def infer_shape(self, uint_crd_f_shape, scaler_f_shape, atom_a_shape, atom_b_shape, atom_c_shape, atom_d_shape,
ipn_shape, pk_shape, gamc_shape, gams_shape, pn_shape): ipn_shape, pk_shape, gamc_shape, gams_shape, pn_shape):
cls_name = self.name cls_name = self.name
M = atom_a_shape[0] M = self.dihedral_numbers
validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f", cls_name) validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f[1]", cls_name)
validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f", cls_name) validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f", cls_name)
validator.check_int(atom_a_shape[0], M, Rel.EQ, "atom_a", cls_name) validator.check_int(atom_a_shape[0], M, Rel.EQ, "atom_a", cls_name)
validator.check_int(atom_b_shape[0], M, Rel.EQ, "atom_b", cls_name) validator.check_int(atom_b_shape[0], M, Rel.EQ, "atom_b", cls_name)
@ -563,7 +583,7 @@ class DihedralAtomEnergy(PrimitiveWithInfer):
@prim_attr_register @prim_attr_register
def __init__(self, dihedral_numbers): def __init__(self, dihedral_numbers):
assert isinstance(dihedral_numbers, int) validator.check_value_type('dihedral_numbers', dihedral_numbers, (int), self.name)
self.dihedral_numbers = dihedral_numbers self.dihedral_numbers = dihedral_numbers
self.init_prim_io_names(inputs=['uint_crd_f', 'scaler_f', 'atom_a', 'atom_b', 'atom_c', 'atom_d', 'ipn', 'pk', self.init_prim_io_names(inputs=['uint_crd_f', 'scaler_f', 'atom_a', 'atom_b', 'atom_c', 'atom_d', 'ipn', 'pk',
'gamc', 'gams', 'pn'], 'gamc', 'gams', 'pn'],
@ -574,8 +594,8 @@ class DihedralAtomEnergy(PrimitiveWithInfer):
ipn_shape, pk_shape, gamc_shape, gams_shape, pn_shape): ipn_shape, pk_shape, gamc_shape, gams_shape, pn_shape):
cls_name = self.name cls_name = self.name
N = uint_crd_f_shape[0] N = uint_crd_f_shape[0]
M = atom_a_shape[0] M = self.dihedral_numbers
validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f", cls_name) validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f[1]", cls_name)
validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f", cls_name) validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f", cls_name)
validator.check_int(atom_a_shape[0], M, Rel.EQ, "atom_a", cls_name) validator.check_int(atom_a_shape[0], M, Rel.EQ, "atom_a", cls_name)
validator.check_int(atom_b_shape[0], M, Rel.EQ, "atom_b", cls_name) validator.check_int(atom_b_shape[0], M, Rel.EQ, "atom_b", cls_name)
@ -641,7 +661,7 @@ class DihedralForceWithAtomEnergy(PrimitiveWithInfer):
@prim_attr_register @prim_attr_register
def __init__(self, dihedral_numbers): def __init__(self, dihedral_numbers):
assert isinstance(dihedral_numbers, int) validator.check_value_type('dihedral_numbers', dihedral_numbers, (int), self.name)
self.dihedral_numbers = dihedral_numbers self.dihedral_numbers = dihedral_numbers
self.init_prim_io_names(inputs=['uint_crd_f', 'scaler_f', 'atom_a', 'atom_b', 'atom_c', 'atom_d', 'ipn', 'pk', self.init_prim_io_names(inputs=['uint_crd_f', 'scaler_f', 'atom_a', 'atom_b', 'atom_c', 'atom_d', 'ipn', 'pk',
'gamc', 'gams', 'pn'], 'gamc', 'gams', 'pn'],
@ -652,8 +672,8 @@ class DihedralForceWithAtomEnergy(PrimitiveWithInfer):
ipn_shape, pk_shape, gamc_shape, gams_shape, pn_shape): ipn_shape, pk_shape, gamc_shape, gams_shape, pn_shape):
cls_name = self.name cls_name = self.name
N = uint_crd_f_shape[0] N = uint_crd_f_shape[0]
M = atom_a_shape[0] M = self.dihedral_numbers
validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f", cls_name) validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f[1]", cls_name)
validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f", cls_name) validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f", cls_name)
validator.check_int(atom_a_shape[0], M, Rel.EQ, "atom_a", cls_name) validator.check_int(atom_a_shape[0], M, Rel.EQ, "atom_a", cls_name)
validator.check_int(atom_b_shape[0], M, Rel.EQ, "atom_b", cls_name) validator.check_int(atom_b_shape[0], M, Rel.EQ, "atom_b", cls_name)
@ -690,14 +710,18 @@ class AngleForce(PrimitiveWithInfer):
number of atoms is N. number of atoms is N.
.. math:: .. math::
dr_{ab} = (x_b-x_a, y_b-y_a, z_b-z_a) dr_{ab} = (x_b-x_a, y_b-y_a, z_b-z_a)
.. math::
dr_{cb} = (x_b-x_c, y_b-y_c, z_b-z_c) dr_{cb} = (x_b-x_c, y_b-y_c, z_b-z_c)
.. math::
theta = arccos(inner_product(dr_{ab}, dr_{cb})/|dr_{ab}|/|dr_{cb}|) theta = arccos(inner_product(dr_{ab}, dr_{cb})/|dr_{ab}|/|dr_{cb}|)
.. math::
F_a = -2*k*(theta-theta_0)/sin(theta)*[cos(theta)/|dr_{ab}|^2*dr_{ab} F_a = -2*k*(theta-theta_0)/sin(theta)*[cos(theta)/|dr_{ab}|^2*dr_{ab}
- 1/|dr_{ab}|/|dr_{cb}|*dr_{cb}] - 1/|dr_{ab}|/|dr_{cb}|*dr_{cb}]
.. math::
F_c = -2*k*(theta-theta_0)/sin(theta)*[cos(theta)/|dr_{cb}|^2*dr_{cb} F_c = -2*k*(theta-theta_0)/sin(theta)*[cos(theta)/|dr_{cb}|^2*dr_{cb}
- 1/|dr_{cb}|/|dr_{ab}|*dr_{ab}] - 1/|dr_{cb}|/|dr_{ab}|*dr_{ab}]
.. math::
F_b = -F_a - F_c F_b = -F_a - F_c
Args: Args:
@ -722,7 +746,7 @@ class AngleForce(PrimitiveWithInfer):
@prim_attr_register @prim_attr_register
def __init__(self, angle_numbers): def __init__(self, angle_numbers):
assert isinstance(angle_numbers, int) validator.check_value_type('angle_numbers', angle_numbers, (int), self.name)
self.angle_numbers = angle_numbers self.angle_numbers = angle_numbers
self.init_prim_io_names(inputs=['uint_crd_f', 'scaler_f', 'atom_a', 'atom_b', 'atom_c', 'angle_k', self.init_prim_io_names(inputs=['uint_crd_f', 'scaler_f', 'atom_a', 'atom_b', 'atom_c', 'angle_k',
'angle_theta0'], 'angle_theta0'],
@ -732,8 +756,8 @@ class AngleForce(PrimitiveWithInfer):
def infer_shape(self, uint_crd_f_shape, scaler_f_shape, atom_a_shape, atom_b_shape, atom_c_shape, angle_k_shape, def infer_shape(self, uint_crd_f_shape, scaler_f_shape, atom_a_shape, atom_b_shape, atom_c_shape, angle_k_shape,
angle_theta0_shape): angle_theta0_shape):
cls_name = self.name cls_name = self.name
M = atom_a_shape[0] M = self.angle_numbers
validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f", cls_name) validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f[1]", cls_name)
validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f", cls_name) validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f", cls_name)
validator.check_int(atom_a_shape[0], M, Rel.EQ, "atom_a", cls_name) validator.check_int(atom_a_shape[0], M, Rel.EQ, "atom_a", cls_name)
validator.check_int(atom_b_shape[0], M, Rel.EQ, "atom_b", cls_name) validator.check_int(atom_b_shape[0], M, Rel.EQ, "atom_b", cls_name)
@ -756,13 +780,16 @@ class AngleForce(PrimitiveWithInfer):
class AngleEnergy(PrimitiveWithInfer): class AngleEnergy(PrimitiveWithInfer):
""" """
Calculate the energy caused by 3-atoms angle term. Calculate the energy caused by 3-atoms angle term. Assume the number of angles is M and the
number of atoms is N.
.. math:: .. math::
dr_{ab} = (x_b-x_a, y_b-y_a, z_b-z_a) dr_{ab} = (x_b-x_a, y_b-y_a, z_b-z_a)
.. math::
dr_{cb} = (x_b-x_c, y_b-y_c, z_b-z_c) dr_{cb} = (x_b-x_c, y_b-y_c, z_b-z_c)
.. math::
theta = arccos(inner_product(dr_{ab}, dr_{cb})/|dr_{ab}|/|dr_{cb}|) theta = arccos(inner_product(dr_{ab}, dr_{cb})/|dr_{ab}|/|dr_{cb}|)
.. math::
E = k*(theta - theta_0)^2 E = k*(theta - theta_0)^2
Args: Args:
@ -787,7 +814,7 @@ class AngleEnergy(PrimitiveWithInfer):
@prim_attr_register @prim_attr_register
def __init__(self, angle_numbers): def __init__(self, angle_numbers):
assert isinstance(angle_numbers, int) validator.check_value_type('angle_numbers', angle_numbers, (int), self.name)
self.angle_numbers = angle_numbers self.angle_numbers = angle_numbers
self.init_prim_io_names(inputs=['uint_crd_f', 'scaler_f', 'atom_a', 'atom_b', 'atom_c', 'angle_k', self.init_prim_io_names(inputs=['uint_crd_f', 'scaler_f', 'atom_a', 'atom_b', 'atom_c', 'angle_k',
'angle_theta0'], 'angle_theta0'],
@ -797,8 +824,8 @@ class AngleEnergy(PrimitiveWithInfer):
def infer_shape(self, uint_crd_f_shape, scaler_f_shape, atom_a_shape, atom_b_shape, atom_c_shape, angle_k_shape, def infer_shape(self, uint_crd_f_shape, scaler_f_shape, atom_a_shape, atom_b_shape, atom_c_shape, angle_k_shape,
angle_theta0_shape): angle_theta0_shape):
cls_name = self.name cls_name = self.name
M = atom_a_shape[0] M = self.angle_numbers
validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f", cls_name) validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f[1]", cls_name)
validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f", cls_name) validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f", cls_name)
validator.check_int(atom_a_shape[0], M, Rel.EQ, "atom_a", cls_name) validator.check_int(atom_a_shape[0], M, Rel.EQ, "atom_a", cls_name)
validator.check_int(atom_b_shape[0], M, Rel.EQ, "atom_b", cls_name) validator.check_int(atom_b_shape[0], M, Rel.EQ, "atom_b", cls_name)
@ -822,7 +849,8 @@ class AngleEnergy(PrimitiveWithInfer):
class AngleAtomEnergy(PrimitiveWithInfer): class AngleAtomEnergy(PrimitiveWithInfer):
""" """
Add the potential energy caused by angle terms to the total potential Add the potential energy caused by angle terms to the total potential
energy of each atom. energy of each atom. Assume the number of angles is M and the
number of atoms is N.
The calculation formula is the same as operator AngleEnergy(). The calculation formula is the same as operator AngleEnergy().
@ -848,7 +876,7 @@ class AngleAtomEnergy(PrimitiveWithInfer):
@prim_attr_register @prim_attr_register
def __init__(self, angle_numbers): def __init__(self, angle_numbers):
assert isinstance(angle_numbers, int) validator.check_value_type('angle_numbers', angle_numbers, (int), self.name)
self.angle_numbers = angle_numbers self.angle_numbers = angle_numbers
self.init_prim_io_names(inputs=['uint_crd_f', 'scaler_f', 'atom_a', 'atom_b', 'atom_c', 'angle_k', self.init_prim_io_names(inputs=['uint_crd_f', 'scaler_f', 'atom_a', 'atom_b', 'atom_c', 'angle_k',
'angle_theta0'], 'angle_theta0'],
@ -859,8 +887,8 @@ class AngleAtomEnergy(PrimitiveWithInfer):
angle_theta0_shape): angle_theta0_shape):
cls_name = self.name cls_name = self.name
N = uint_crd_f_shape[0] N = uint_crd_f_shape[0]
M = atom_a_shape[0] M = self.angle_numbers
validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f", cls_name) validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f[1]", cls_name)
validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f", cls_name) validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f", cls_name)
validator.check_int(atom_a_shape[0], M, Rel.EQ, "atom_a", cls_name) validator.check_int(atom_a_shape[0], M, Rel.EQ, "atom_a", cls_name)
validator.check_int(atom_b_shape[0], M, Rel.EQ, "atom_b", cls_name) validator.check_int(atom_b_shape[0], M, Rel.EQ, "atom_b", cls_name)
@ -883,7 +911,8 @@ class AngleAtomEnergy(PrimitiveWithInfer):
class AngleForceWithAtomEnergy(PrimitiveWithInfer): class AngleForceWithAtomEnergy(PrimitiveWithInfer):
""" """
Calculate angle force and potential energy together. Calculate angle force and potential energy together. Assume the number of angles is M and the
number of atoms is N.
The calculation formula is the same as operator AngleForce() and AngleEnergy(). The calculation formula is the same as operator AngleForce() and AngleEnergy().
@ -910,7 +939,7 @@ class AngleForceWithAtomEnergy(PrimitiveWithInfer):
@prim_attr_register @prim_attr_register
def __init__(self, angle_numbers): def __init__(self, angle_numbers):
assert isinstance(angle_numbers, int) validator.check_value_type('angle_numbers', angle_numbers, (int), self.name)
self.angle_numbers = angle_numbers self.angle_numbers = angle_numbers
self.init_prim_io_names(inputs=['uint_crd_f', 'scaler_f', 'atom_a', 'atom_b', 'atom_c', 'angle_k', self.init_prim_io_names(inputs=['uint_crd_f', 'scaler_f', 'atom_a', 'atom_b', 'atom_c', 'angle_k',
'angle_theta0'], 'angle_theta0'],
@ -921,8 +950,8 @@ class AngleForceWithAtomEnergy(PrimitiveWithInfer):
angle_theta0_shape): angle_theta0_shape):
cls_name = self.name cls_name = self.name
N = uint_crd_f_shape[0] N = uint_crd_f_shape[0]
M = atom_a_shape[0] M = self.angle_numbers
validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f", cls_name) validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f[1]", cls_name)
validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f", cls_name) validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f", cls_name)
validator.check_int(atom_a_shape[0], M, Rel.EQ, "atom_a", cls_name) validator.check_int(atom_a_shape[0], M, Rel.EQ, "atom_a", cls_name)
validator.check_int(atom_b_shape[0], M, Rel.EQ, "atom_b", cls_name) validator.check_int(atom_b_shape[0], M, Rel.EQ, "atom_b", cls_name)
@ -954,8 +983,8 @@ class Dihedral14LJForce(PrimitiveWithInfer):
for all kinds of atom pairs. for all kinds of atom pairs.
.. math:: .. math::
dr = (x_a-x_b, y_a-y_b, z_a-z_b) dr = (x_a-x_b, y_a-y_b, z_a-z_b)
.. math::
F = k*(-12*A/|dr|^{14} + 6*B/|dr|^{8})*dr F = k*(-12*A/|dr|^{14} + 6*B/|dr|^{8})*dr
Args: Args:
@ -972,7 +1001,9 @@ class Dihedral14LJForce(PrimitiveWithInfer):
- **lj_scale_factor** (Tensor, float32) - [M,], the scale factor for the - **lj_scale_factor** (Tensor, float32) - [M,], the scale factor for the
Lennard-Jones part of force correction of each dihedral 1,4 term. Lennard-Jones part of force correction of each dihedral 1,4 term.
- **LJ_type_A** (Tensor, float32) - [Q,], the A parameter in Lennard-Jones scheme of each atom pair type. - **LJ_type_A** (Tensor, float32) - [Q,], the A parameter in Lennard-Jones scheme of each atom pair type.
Q is the number of atom pair.
- **LJ_type_B** (Tensor, float32) - [Q,], the B parameter in Lennard-Jones shceme of each atom pair type. - **LJ_type_B** (Tensor, float32) - [Q,], the B parameter in Lennard-Jones shceme of each atom pair type.
Q is the number of atom pair.
Outputs: Outputs:
- **frc_f** (Tensor, float32) - [N, 3], the force felt by each atom. - **frc_f** (Tensor, float32) - [N, 3], the force felt by each atom.
@ -983,8 +1014,8 @@ class Dihedral14LJForce(PrimitiveWithInfer):
@prim_attr_register @prim_attr_register
def __init__(self, nb14_numbers, atom_numbers): def __init__(self, nb14_numbers, atom_numbers):
assert isinstance(nb14_numbers, int) validator.check_value_type('nb14_numbers', nb14_numbers, (int), self.name)
assert isinstance(atom_numbers, int) validator.check_value_type('atom_numbers', atom_numbers, (int), self.name)
self.dihedral_14_numbers = nb14_numbers self.dihedral_14_numbers = nb14_numbers
self.atom_numbers = atom_numbers self.atom_numbers = atom_numbers
self.init_prim_io_names( self.init_prim_io_names(
@ -996,6 +1027,19 @@ class Dihedral14LJForce(PrimitiveWithInfer):
def infer_shape(self, uint_crd_f_shape, LJtype_shape, charge_shape, boxlength_f_shape, a_14_shape, b_14_shape, def infer_shape(self, uint_crd_f_shape, LJtype_shape, charge_shape, boxlength_f_shape, a_14_shape, b_14_shape,
lj_scale_factor_shape, LJ_type_A_shape, LJ_type_B_shape): lj_scale_factor_shape, LJ_type_A_shape, LJ_type_B_shape):
cls_name = self.name
N = self.atom_numbers
M = self.dihedral_14_numbers
Q = LJ_type_A_shape[0]
validator.check_int(uint_crd_f_shape[0], N, Rel.EQ, "uint_crd_f[0]", cls_name)
validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f[1]", cls_name)
validator.check_int(LJtype_shape[0], 3, Rel.EQ, "LJtype", cls_name)
validator.check_int(charge_shape[0], M, Rel.EQ, "charge", cls_name)
validator.check_int(boxlength_f_shape[0], 3, Rel.EQ, "boxlength_f", cls_name)
validator.check_int(a_14_shape[0], M, Rel.EQ, "a_14", cls_name)
validator.check_int(b_14_shape[0], M, Rel.EQ, "b_14", cls_name)
validator.check_int(lj_scale_factor_shape[0], M, Rel.EQ, "lj_scale_factor", cls_name)
validator.check_int(LJ_type_B_shape[0], Q, Rel.EQ, "LJ_type_B", cls_name)
return uint_crd_f_shape return uint_crd_f_shape
def infer_dtype(self, uint_crd_f_dtype, LJtype_dtype, charge_dtype, boxlength_f_type, a_14_type, b_14_type, def infer_dtype(self, uint_crd_f_dtype, LJtype_dtype, charge_dtype, boxlength_f_type, a_14_type, b_14_type,
@ -1019,9 +1063,9 @@ class Dihedral14LJEnergy(PrimitiveWithInfer):
Calculate the Lennard-Jones part of 1,4 dihedral energy correction for Calculate the Lennard-Jones part of 1,4 dihedral energy correction for
each necessary dihedral terms on the corresponding atoms. each necessary dihedral terms on the corresponding atoms.
.. math:, the : .. math::
dr = (x_a-x_b, y_a-y_b, z_a-z-b) dr = (x_a-x_b, y_a-y_b, z_a-z-b)
.. math::
E = k*(A/|dr|^{12} - B/|dr|^{6}) E = k*(A/|dr|^{12} - B/|dr|^{6})
Args: Args:
@ -1038,7 +1082,9 @@ class Dihedral14LJEnergy(PrimitiveWithInfer):
- **lj_scale_factor** (Tensor, float32) - [M,], the scale factor for the - **lj_scale_factor** (Tensor, float32) - [M,], the scale factor for the
Lennard-Jones part of force correction of each dihedral 1,4 term. Lennard-Jones part of force correction of each dihedral 1,4 term.
- **LJ_type_A** (Tensor, float32) - [Q,], the A parameter in Lennard-Jones scheme of each atom pair type. - **LJ_type_A** (Tensor, float32) - [Q,], the A parameter in Lennard-Jones scheme of each atom pair type.
Q is the number of atom pair.
- **LJ_type_B** (Tensor, float32) - [Q,], the B parameter in Lennard-Jones shceme of each atom pair type. - **LJ_type_B** (Tensor, float32) - [Q,], the B parameter in Lennard-Jones shceme of each atom pair type.
Q is the number of atom pair.
Outputs: Outputs:
- **ene** (Tensor, float32) - [M,], the Lennard-Jones potential - **ene** (Tensor, float32) - [M,], the Lennard-Jones potential
@ -1050,8 +1096,8 @@ class Dihedral14LJEnergy(PrimitiveWithInfer):
@prim_attr_register @prim_attr_register
def __init__(self, nb14_numbers, atom_numbers): def __init__(self, nb14_numbers, atom_numbers):
assert isinstance(nb14_numbers, int) validator.check_value_type('nb14_numbers', nb14_numbers, (int), self.name)
assert isinstance(atom_numbers, int) validator.check_value_type('atom_numbers', atom_numbers, (int), self.name)
self.dihedral_14_numbers = nb14_numbers self.dihedral_14_numbers = nb14_numbers
self.atom_numbers = atom_numbers self.atom_numbers = atom_numbers
@ -1064,6 +1110,19 @@ class Dihedral14LJEnergy(PrimitiveWithInfer):
def infer_shape(self, uint_crd_f_shape, LJtype_shape, charge_shape, boxlength_f_shape, a_14_shape, b_14_shape, def infer_shape(self, uint_crd_f_shape, LJtype_shape, charge_shape, boxlength_f_shape, a_14_shape, b_14_shape,
lj_scale_factor_shape, LJ_type_A_shape, LJ_type_B_shape): lj_scale_factor_shape, LJ_type_A_shape, LJ_type_B_shape):
cls_name = self.name
N = self.atom_numbers
M = self.dihedral_14_numbers
Q = LJ_type_A_shape[0]
validator.check_int(uint_crd_f_shape[0], N, Rel.EQ, "uint_crd_f[0]", cls_name)
validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f[1]", cls_name)
validator.check_int(LJtype_shape[0], 3, Rel.EQ, "LJtype", cls_name)
validator.check_int(charge_shape[0], M, Rel.EQ, "charge", cls_name)
validator.check_int(boxlength_f_shape[0], 3, Rel.EQ, "boxlength_f", cls_name)
validator.check_int(a_14_shape[0], M, Rel.EQ, "a_14", cls_name)
validator.check_int(b_14_shape[0], M, Rel.EQ, "b_14", cls_name)
validator.check_int(lj_scale_factor_shape[0], M, Rel.EQ, "lj_scale_factor", cls_name)
validator.check_int(LJ_type_B_shape[0], Q, Rel.EQ, "LJ_type_B", cls_name)
return [self.dihedral_14_numbers,] return [self.dihedral_14_numbers,]
def infer_dtype(self, uint_crd_f_dtype, LJtype_dtype, charge_dtype, boxlength_f_type, a_14_type, b_14_type, def infer_dtype(self, uint_crd_f_dtype, LJtype_dtype, charge_dtype, boxlength_f_type, a_14_type, b_14_type,
@ -1090,9 +1149,9 @@ class Dihedral14LJForceWithDirectCF(PrimitiveWithInfer):
Dihedral14LJForce(), and the Coulomb part is as follows: Dihedral14LJForce(), and the Coulomb part is as follows:
.. math:: .. math::
dr = (x_a-x_b, y_a-y_b, z_a-z_b)
dr = (x_a-x_b, y_a-y_b, z_a-z_b) .. math::
F = -k*q_a*q_b/|r|^3*dr F = -k*q_a*q_b/|r|^3*dr
Args: Args:
nb14_numbers (int32): the number of necessary dihedral 1,4 terms M. nb14_numbers (int32): the number of necessary dihedral 1,4 terms M.
@ -1110,7 +1169,9 @@ class Dihedral14LJForceWithDirectCF(PrimitiveWithInfer):
- **cf_scale_factor** (Tensor, float) - [M,], the scale factor for the - **cf_scale_factor** (Tensor, float) - [M,], the scale factor for the
Coulomb part of force correction for each dihedral 1,4 terms. Coulomb part of force correction for each dihedral 1,4 terms.
- **LJ_type_A** (Tensor, float32) - [Q,], the A parameter in Lennard-Jones scheme of each atom pair type. - **LJ_type_A** (Tensor, float32) - [Q,], the A parameter in Lennard-Jones scheme of each atom pair type.
Q is the number of atom pair.
- **LJ_type_B** (Tensor, float32) - [Q,], the B parameter in Lennard-Jones shceme of each atom pair type. - **LJ_type_B** (Tensor, float32) - [Q,], the B parameter in Lennard-Jones shceme of each atom pair type.
Q is the number of atom pair.
Outputs: Outputs:
- **frc_f** (Tensor, float) - [N, 3], the force felt by each atom. - **frc_f** (Tensor, float) - [N, 3], the force felt by each atom.
@ -1121,8 +1182,8 @@ class Dihedral14LJForceWithDirectCF(PrimitiveWithInfer):
@prim_attr_register @prim_attr_register
def __init__(self, nb14_numbers, atom_numbers): def __init__(self, nb14_numbers, atom_numbers):
assert isinstance(nb14_numbers, int) validator.check_value_type('nb14_numbers', nb14_numbers, (int), self.name)
assert isinstance(atom_numbers, int) validator.check_value_type('atom_numbers', atom_numbers, (int), self.name)
self.dihedral_14_numbers = nb14_numbers self.dihedral_14_numbers = nb14_numbers
self.atom_numbers = atom_numbers self.atom_numbers = atom_numbers
@ -1136,6 +1197,20 @@ class Dihedral14LJForceWithDirectCF(PrimitiveWithInfer):
def infer_shape(self, uint_crd_f_shape, LJtype_shape, charge_shape, boxlength_f_shape, a_14_shape, b_14_shape, def infer_shape(self, uint_crd_f_shape, LJtype_shape, charge_shape, boxlength_f_shape, a_14_shape, b_14_shape,
lj_scale_factor_shape, cf_scale_factor_shape, LJ_type_A_shape, LJ_type_B_shape): lj_scale_factor_shape, cf_scale_factor_shape, LJ_type_A_shape, LJ_type_B_shape):
cls_name = self.name
N = self.atom_numbers
M = self.dihedral_14_numbers
Q = LJ_type_A_shape[0]
validator.check_int(uint_crd_f_shape[0], N, Rel.EQ, "uint_crd_f[0]", cls_name)
validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f[1]", cls_name)
validator.check_int(LJtype_shape[0], 3, Rel.EQ, "LJtype", cls_name)
validator.check_int(charge_shape[0], M, Rel.EQ, "charge", cls_name)
validator.check_int(boxlength_f_shape[0], 3, Rel.EQ, "boxlength_f", cls_name)
validator.check_int(a_14_shape[0], M, Rel.EQ, "a_14", cls_name)
validator.check_int(b_14_shape[0], M, Rel.EQ, "b_14", cls_name)
validator.check_int(lj_scale_factor_shape[0], M, Rel.EQ, "lj_scale_factor", cls_name)
validator.check_int(cf_scale_factor_shape[0], M, Rel.EQ, "cf_scale_factor", cls_name)
validator.check_int(LJ_type_B_shape[0], Q, Rel.EQ, "LJ_type_B", cls_name)
return [self.atom_numbers, 3] return [self.atom_numbers, 3]
def infer_dtype(self, uint_crd_f_dtype, LJtype_dtype, charge_dtype, boxlength_f_type, a_14_type, b_14_type, def infer_dtype(self, uint_crd_f_dtype, LJtype_dtype, charge_dtype, boxlength_f_type, a_14_type, b_14_type,
@ -1180,7 +1255,9 @@ class Dihedral14LJCFForceWithAtomEnergy(PrimitiveWithInfer):
- **cf_scale_factor** (Tensor, float) - [M,], the scale factor for the - **cf_scale_factor** (Tensor, float) - [M,], the scale factor for the
Coulomb part of force correction for each dihedral 1,4 terms. Coulomb part of force correction for each dihedral 1,4 terms.
- **LJ_type_A** (Tensor, float32) - [Q,], the A parameter in Lennard-Jones scheme of each atom pair type. - **LJ_type_A** (Tensor, float32) - [Q,], the A parameter in Lennard-Jones scheme of each atom pair type.
Q is the number of atom pair.
- **LJ_type_B** (Tensor, float32) - [Q,], the B parameter in Lennard-Jones shceme of each atom pair type. - **LJ_type_B** (Tensor, float32) - [Q,], the B parameter in Lennard-Jones shceme of each atom pair type.
Q is the number of atom pair.
Outputs: Outputs:
- **frc_f** (Tensor, float32) - [N, 3], the force felt by each atom. - **frc_f** (Tensor, float32) - [N, 3], the force felt by each atom.
@ -1192,8 +1269,8 @@ class Dihedral14LJCFForceWithAtomEnergy(PrimitiveWithInfer):
@prim_attr_register @prim_attr_register
def __init__(self, nb14_numbers, atom_numbers): def __init__(self, nb14_numbers, atom_numbers):
assert isinstance(nb14_numbers, int) validator.check_value_type('nb14_numbers', nb14_numbers, (int), self.name)
assert isinstance(atom_numbers, int) validator.check_value_type('atom_numbers', atom_numbers, (int), self.name)
self.dihedral_14_numbers = nb14_numbers self.dihedral_14_numbers = nb14_numbers
self.atom_numbers = atom_numbers self.atom_numbers = atom_numbers
@ -1207,6 +1284,20 @@ class Dihedral14LJCFForceWithAtomEnergy(PrimitiveWithInfer):
def infer_shape(self, uint_crd_f_shape, LJtype_shape, charge_shape, boxlength_f_shape, a_14_shape, b_14_shape, def infer_shape(self, uint_crd_f_shape, LJtype_shape, charge_shape, boxlength_f_shape, a_14_shape, b_14_shape,
lj_scale_factor_shape, cf_scale_factor_shape, LJ_type_A_shape, LJ_type_B_shape): lj_scale_factor_shape, cf_scale_factor_shape, LJ_type_A_shape, LJ_type_B_shape):
cls_name = self.name
N = self.atom_numbers
M = self.dihedral_14_numbers
Q = LJ_type_A_shape[0]
validator.check_int(uint_crd_f_shape[0], N, Rel.EQ, "uint_crd_f[0]", cls_name)
validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f[1]", cls_name)
validator.check_int(LJtype_shape[0], 3, Rel.EQ, "LJtype", cls_name)
validator.check_int(charge_shape[0], M, Rel.EQ, "charge", cls_name)
validator.check_int(boxlength_f_shape[0], 3, Rel.EQ, "boxlength_f", cls_name)
validator.check_int(a_14_shape[0], M, Rel.EQ, "a_14", cls_name)
validator.check_int(b_14_shape[0], M, Rel.EQ, "b_14", cls_name)
validator.check_int(lj_scale_factor_shape[0], M, Rel.EQ, "lj_scale_factor", cls_name)
validator.check_int(cf_scale_factor_shape[0], M, Rel.EQ, "cf_scale_factor", cls_name)
validator.check_int(LJ_type_B_shape[0], Q, Rel.EQ, "LJ_type_B", cls_name)
return uint_crd_f_shape, charge_shape return uint_crd_f_shape, charge_shape
def infer_dtype(self, uint_crd_f_dtype, LJtype_dtype, charge_dtype, boxlength_f_type, a_14_type, b_14_type, def infer_dtype(self, uint_crd_f_dtype, LJtype_dtype, charge_dtype, boxlength_f_type, a_14_type, b_14_type,
@ -1246,7 +1337,9 @@ class Dihedral14LJAtomEnergy(PrimitiveWithInfer):
- **lj_scale_factor** (Tensor, float32) - [M,], the scale factor for the - **lj_scale_factor** (Tensor, float32) - [M,], the scale factor for the
Lennard-Jones part of force correction of each dihedral 1,4 term. Lennard-Jones part of force correction of each dihedral 1,4 term.
- **LJ_type_A** (Tensor, float32) - [Q,], the A parameter in Lennard-Jones scheme of each atom pair type. - **LJ_type_A** (Tensor, float32) - [Q,], the A parameter in Lennard-Jones scheme of each atom pair type.
Q is the number of atom pair.
- **LJ_type_B** (Tensor, float32) - [Q,], the B parameter in Lennard-Jones shceme of each atom pair type. - **LJ_type_B** (Tensor, float32) - [Q,], the B parameter in Lennard-Jones shceme of each atom pair type.
Q is the number of atom pair.
Outputs: Outputs:
- **ene** (Tensor, float32) - [N,], the accumulated potential energy of each atom. - **ene** (Tensor, float32) - [N,], the accumulated potential energy of each atom.
@ -1257,8 +1350,8 @@ class Dihedral14LJAtomEnergy(PrimitiveWithInfer):
@prim_attr_register @prim_attr_register
def __init__(self, nb14_numbers, atom_numbers): def __init__(self, nb14_numbers, atom_numbers):
assert isinstance(nb14_numbers, int) validator.check_value_type('nb14_numbers', nb14_numbers, (int), self.name)
assert isinstance(atom_numbers, int) validator.check_value_type('atom_numbers', atom_numbers, (int), self.name)
self.dihedral_14_numbers = nb14_numbers self.dihedral_14_numbers = nb14_numbers
self.atom_numbers = atom_numbers self.atom_numbers = atom_numbers
@ -1271,6 +1364,19 @@ class Dihedral14LJAtomEnergy(PrimitiveWithInfer):
def infer_shape(self, uint_crd_f_shape, LJtype_shape, charge_shape, boxlength_f_shape, a_14_shape, b_14_shape, def infer_shape(self, uint_crd_f_shape, LJtype_shape, charge_shape, boxlength_f_shape, a_14_shape, b_14_shape,
lj_scale_factor_shape, LJ_type_A_shape, LJ_type_B_shape): lj_scale_factor_shape, LJ_type_A_shape, LJ_type_B_shape):
cls_name = self.name
N = self.atom_numbers
M = self.dihedral_14_numbers
Q = LJ_type_A_shape[0]
validator.check_int(uint_crd_f_shape[0], N, Rel.EQ, "uint_crd_f[0]", cls_name)
validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f[1]", cls_name)
validator.check_int(LJtype_shape[0], 3, Rel.EQ, "LJtype", cls_name)
validator.check_int(charge_shape[0], M, Rel.EQ, "charge", cls_name)
validator.check_int(boxlength_f_shape[0], 3, Rel.EQ, "boxlength_f", cls_name)
validator.check_int(a_14_shape[0], M, Rel.EQ, "a_14", cls_name)
validator.check_int(b_14_shape[0], M, Rel.EQ, "b_14", cls_name)
validator.check_int(lj_scale_factor_shape[0], M, Rel.EQ, "lj_scale_factor", cls_name)
validator.check_int(LJ_type_B_shape[0], Q, Rel.EQ, "LJ_type_B", cls_name)
return LJtype_shape return LJtype_shape
def infer_dtype(self, uint_crd_f_dtype, LJtype_dtype, charge_dtype, boxlength_f_type, a_14_type, b_14_type, def infer_dtype(self, uint_crd_f_dtype, LJtype_dtype, charge_dtype, boxlength_f_type, a_14_type, b_14_type,
@ -1297,6 +1403,8 @@ class Dihedral14CFEnergy(PrimitiveWithInfer):
.. math:: .. math::
dr = (x_a-x_b, y_a-y_b, z_a-z_b) dr = (x_a-x_b, y_a-y_b, z_a-z_b)
.. math::
E = k*q_a*q_b/|dr| E = k*q_a*q_b/|dr|
Args: Args:
@ -1322,8 +1430,8 @@ class Dihedral14CFEnergy(PrimitiveWithInfer):
@prim_attr_register @prim_attr_register
def __init__(self, nb14_numbers, atom_numbers): def __init__(self, nb14_numbers, atom_numbers):
assert isinstance(nb14_numbers, int) validator.check_value_type('nb14_numbers', nb14_numbers, (int), self.name)
assert isinstance(atom_numbers, int) validator.check_value_type('atom_numbers', atom_numbers, (int), self.name)
self.dihedral_14_numbers = nb14_numbers self.dihedral_14_numbers = nb14_numbers
self.atom_numbers = atom_numbers self.atom_numbers = atom_numbers
@ -1335,6 +1443,17 @@ class Dihedral14CFEnergy(PrimitiveWithInfer):
def infer_shape(self, uint_crd_f_shape, LJtype_shape, charge_shape, boxlength_f_shape, a_14_shape, b_14_shape, def infer_shape(self, uint_crd_f_shape, LJtype_shape, charge_shape, boxlength_f_shape, a_14_shape, b_14_shape,
cf_scale_factor_shape): cf_scale_factor_shape):
cls_name = self.name
N = self.atom_numbers
M = self.dihedral_14_numbers
validator.check_int(uint_crd_f_shape[0], N, Rel.EQ, "uint_crd_f[0]", cls_name)
validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f[1]", cls_name)
validator.check_int(LJtype_shape[0], 3, Rel.EQ, "LJtype", cls_name)
validator.check_int(charge_shape[0], M, Rel.EQ, "charge", cls_name)
validator.check_int(boxlength_f_shape[0], 3, Rel.EQ, "boxlength_f", cls_name)
validator.check_int(a_14_shape[0], M, Rel.EQ, "a_14", cls_name)
validator.check_int(b_14_shape[0], M, Rel.EQ, "b_14", cls_name)
validator.check_int(cf_scale_factor_shape[0], M, Rel.EQ, "cf_scale_factor", cls_name)
return [self.dihedral_14_numbers,] return [self.dihedral_14_numbers,]
def infer_dtype(self, uint_crd_f_dtype, LJtype_dtype, charge_dtype, boxlength_f_type, a_14_type, b_14_type, def infer_dtype(self, uint_crd_f_dtype, LJtype_dtype, charge_dtype, boxlength_f_type, a_14_type, b_14_type,
@ -1382,8 +1501,8 @@ class Dihedral14CFAtomEnergy(PrimitiveWithInfer):
@prim_attr_register @prim_attr_register
def __init__(self, nb14_numbers, atom_numbers): def __init__(self, nb14_numbers, atom_numbers):
assert isinstance(nb14_numbers, int) validator.check_value_type('nb14_numbers', nb14_numbers, (int), self.name)
assert isinstance(atom_numbers, int) validator.check_value_type('atom_numbers', atom_numbers, (int), self.name)
self.dihedral_14_numbers = nb14_numbers self.dihedral_14_numbers = nb14_numbers
self.atom_numbers = atom_numbers self.atom_numbers = atom_numbers
@ -1395,6 +1514,17 @@ class Dihedral14CFAtomEnergy(PrimitiveWithInfer):
def infer_shape(self, uint_crd_f_shape, LJtype_shape, charge_shape, boxlength_f_shape, a_14_shape, b_14_shape, def infer_shape(self, uint_crd_f_shape, LJtype_shape, charge_shape, boxlength_f_shape, a_14_shape, b_14_shape,
cf_scale_factor_shape): cf_scale_factor_shape):
cls_name = self.name
N = self.atom_numbers
M = self.dihedral_14_numbers
validator.check_int(uint_crd_f_shape[0], N, Rel.EQ, "uint_crd_f[0]", cls_name)
validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f[1]", cls_name)
validator.check_int(LJtype_shape[0], 3, Rel.EQ, "LJtype", cls_name)
validator.check_int(charge_shape[0], M, Rel.EQ, "charge", cls_name)
validator.check_int(boxlength_f_shape[0], 3, Rel.EQ, "boxlength_f", cls_name)
validator.check_int(a_14_shape[0], M, Rel.EQ, "a_14", cls_name)
validator.check_int(b_14_shape[0], M, Rel.EQ, "b_14", cls_name)
validator.check_int(cf_scale_factor_shape[0], M, Rel.EQ, "cf_scale_factor", cls_name)
return LJtype_shape return LJtype_shape
def infer_dtype(self, uint_crd_f_dtype, LJtype_dtype, charge_dtype, boxlength_f_type, a_14_type, b_14_type, def infer_dtype(self, uint_crd_f_dtype, LJtype_dtype, charge_dtype, boxlength_f_type, a_14_type, b_14_type,
@ -1450,13 +1580,13 @@ class MDIterationLeapFrog(PrimitiveWithInfer):
@prim_attr_register @prim_attr_register
def __init__(self, float4_numbers, atom_numbers, half_dt, dt, exp_gamma, is_max_velocity, max_velocity): def __init__(self, float4_numbers, atom_numbers, half_dt, dt, exp_gamma, is_max_velocity, max_velocity):
assert isinstance(float4_numbers, int) validator.check_value_type('float4_numbers', float4_numbers, (int), self.name)
assert isinstance(atom_numbers, int) validator.check_value_type('atom_numbers', atom_numbers, (int), self.name)
assert isinstance(half_dt, float) validator.check_value_type('half_dt', half_dt, (float), self.name)
assert isinstance(dt, float) validator.check_value_type('dt', dt, (float), self.name)
assert isinstance(exp_gamma, float) validator.check_value_type('exp_gamma', exp_gamma, (float), self.name)
assert isinstance(is_max_velocity, int) validator.check_value_type('is_max_velocity', is_max_velocity, (int), self.name)
assert isinstance(max_velocity, float) validator.check_value_type('max_velocity', max_velocity, (float), self.name)
self.float4_numbers = float4_numbers self.float4_numbers = float4_numbers
self.atom_numbers = atom_numbers self.atom_numbers = atom_numbers
self.half_dt = half_dt self.half_dt = half_dt
@ -1477,6 +1607,10 @@ class MDIterationLeapFrog(PrimitiveWithInfer):
self.add_prim_attr('max_velocity', self.max_velocity) self.add_prim_attr('max_velocity', self.max_velocity)
def infer_shape(self, mass_inverse_shape, sqrt_mass_shape): def infer_shape(self, mass_inverse_shape, sqrt_mass_shape):
cls_name = self.name
N = self.atom_numbers
validator.check_int(mass_inverse_shape[0], N, Rel.EQ, "mass_inverse", cls_name)
validator.check_int(sqrt_mass_shape[0], N, Rel.EQ, "sqrt_mass", cls_name)
return [self.atom_numbers, 3], [self.atom_numbers, 3], [self.atom_numbers, 3], [self.atom_numbers, 3] return [self.atom_numbers, 3], [self.atom_numbers, 3], [self.atom_numbers, 3], [self.atom_numbers, 3]
def infer_dtype(self, mass_inverse_dtype, sqrt_mass_dtype): def infer_dtype(self, mass_inverse_dtype, sqrt_mass_dtype):
@ -1517,11 +1651,11 @@ class PMEReciprocalForce(PrimitiveWithInfer):
@prim_attr_register @prim_attr_register
def __init__(self, atom_numbers, beta, fftx, ffty, fftz): def __init__(self, atom_numbers, beta, fftx, ffty, fftz):
assert isinstance(atom_numbers, int) validator.check_value_type('atom_numbers', atom_numbers, (int), self.name)
assert isinstance(beta, float) validator.check_value_type('beta', beta, (float), self.name)
assert isinstance(fftx, int) validator.check_value_type('fftx', fftx, (int), self.name)
assert isinstance(ffty, int) validator.check_value_type('ffty', ffty, (int), self.name)
assert isinstance(fftz, int) validator.check_value_type('fftz', fftz, (int), self.name)
self.atom_numbers = atom_numbers self.atom_numbers = atom_numbers
self.beta = beta self.beta = beta
self.fftx = fftx self.fftx = fftx
@ -1536,6 +1670,12 @@ class PMEReciprocalForce(PrimitiveWithInfer):
self.add_prim_attr('fftz', self.fftz) self.add_prim_attr('fftz', self.fftz)
def infer_shape(self, boxlength_shape, uint_crd_shape, charge_shape): def infer_shape(self, boxlength_shape, uint_crd_shape, charge_shape):
cls_name = self.name
N = self.atom_numbers
validator.check_int(uint_crd_shape[0], N, Rel.EQ, "uint_crd[0]", cls_name)
validator.check_int(uint_crd_shape[1], 3, Rel.EQ, "uint_crd[1]", cls_name)
validator.check_int(boxlength_shape[0], 3, Rel.EQ, "boxlength", cls_name)
validator.check_int(charge_shape[0], N, Rel.EQ, "charge", cls_name)
return uint_crd_shape return uint_crd_shape
def infer_dtype(self, boxlength_type, uint_crd_type, charge_type): def infer_dtype(self, boxlength_type, uint_crd_type, charge_type):
@ -1564,7 +1704,7 @@ class PMEExcludedForce(PrimitiveWithInfer):
- **excluded_list_start** (Tensor, int32) - [N,], the start excluded index - **excluded_list_start** (Tensor, int32) - [N,], the start excluded index
in excluded list for each atom. in excluded list for each atom.
- **excluded_list** (Tensor, int32) - [E,], the contiguous join of excluded - **excluded_list** (Tensor, int32) - [E,], the contiguous join of excluded
list of each atom. list of each atom. E is the number of excluded atoms.
- **excluded_atom_numbers** (Tensor, int32) - [N,], the number of atom excluded - **excluded_atom_numbers** (Tensor, int32) - [N,], the number of atom excluded
in excluded list for each atom. in excluded list for each atom.
@ -1577,8 +1717,8 @@ class PMEExcludedForce(PrimitiveWithInfer):
@prim_attr_register @prim_attr_register
def __init__(self, atom_numbers, beta): def __init__(self, atom_numbers, beta):
assert isinstance(atom_numbers, int) validator.check_value_type('atom_numbers', atom_numbers, (int), self.name)
assert isinstance(beta, float) validator.check_value_type('beta', beta, (float), self.name)
self.atom_numbers = atom_numbers self.atom_numbers = atom_numbers
self.beta = beta self.beta = beta
self.init_prim_io_names( self.init_prim_io_names(
@ -1589,6 +1729,14 @@ class PMEExcludedForce(PrimitiveWithInfer):
def infer_shape(self, uint_crd_shape, sacler_shape, charge_shape, excluded_list_start_shape, excluded_list_shape, def infer_shape(self, uint_crd_shape, sacler_shape, charge_shape, excluded_list_start_shape, excluded_list_shape,
excluded_atom_numbers_shape): excluded_atom_numbers_shape):
cls_name = self.name
N = self.atom_numbers
validator.check_int(uint_crd_shape[0], N, Rel.EQ, "uint_crd[0]", cls_name)
validator.check_int(uint_crd_shape[1], 3, Rel.EQ, "uint_crd[1]", cls_name)
validator.check_int(sacler_shape[0], 3, Rel.EQ, "sacler", cls_name)
validator.check_int(charge_shape[0], N, Rel.EQ, "charge", cls_name)
validator.check_int(excluded_list_start_shape[0], N, Rel.EQ, "excluded_list_start", cls_name)
validator.check_int(excluded_atom_numbers_shape[0], N, Rel.EQ, "excluded_atom_numbers", cls_name)
return uint_crd_shape return uint_crd_shape
def infer_dtype(self, uint_crd_type, sacler_type, charge_type, excluded_list_start_type, excluded_list_type, def infer_dtype(self, uint_crd_type, sacler_type, charge_type, excluded_list_start_type, excluded_list_type,
@ -1632,7 +1780,7 @@ class PMEEnergy(PrimitiveWithInfer):
- **excluded_list_start** (Tensor, int32) - [N,], the start excluded index - **excluded_list_start** (Tensor, int32) - [N,], the start excluded index
in excluded list for each atom. in excluded list for each atom.
- **excluded_list** (Tensor, int32) - [E,], the contiguous join of excluded - **excluded_list** (Tensor, int32) - [E,], the contiguous join of excluded
list of each atom. list of each atom. E is the number of excluded atoms.
- **excluded_atom_numbers** (Tensor, int32) - [N,], the number of atom excluded - **excluded_atom_numbers** (Tensor, int32) - [N,], the number of atom excluded
in excluded list for each atom. in excluded list for each atom.
@ -1648,11 +1796,11 @@ class PMEEnergy(PrimitiveWithInfer):
@prim_attr_register @prim_attr_register
def __init__(self, atom_numbers, beta, fftx, ffty, fftz): def __init__(self, atom_numbers, beta, fftx, ffty, fftz):
assert isinstance(atom_numbers, int) validator.check_value_type('atom_numbers', atom_numbers, (int), self.name)
assert isinstance(beta, float) validator.check_value_type('beta', beta, (float), self.name)
assert isinstance(fftx, int) validator.check_value_type('fftx', fftx, (int), self.name)
assert isinstance(ffty, int) validator.check_value_type('ffty', ffty, (int), self.name)
assert isinstance(fftz, int) validator.check_value_type('fftz', fftz, (int), self.name)
self.atom_numbers = atom_numbers self.atom_numbers = atom_numbers
self.beta = beta self.beta = beta
self.fftx = fftx self.fftx = fftx
@ -1670,6 +1818,18 @@ class PMEEnergy(PrimitiveWithInfer):
def infer_shape(self, box_length, uint_crd, charge, nl_numbers, nl_serial, scaler, excluded_list_start, def infer_shape(self, box_length, uint_crd, charge, nl_numbers, nl_serial, scaler, excluded_list_start,
excluded_list, excluded_atom_numbers): excluded_list, excluded_atom_numbers):
cls_name = self.name
N = self.atom_numbers
validator.check_int(uint_crd[0], N, Rel.EQ, "uint_crd[0]", cls_name)
validator.check_int(uint_crd[1], 3, Rel.EQ, "uint_crd[1]", cls_name)
validator.check_int(box_length[0], 3, Rel.EQ, "box_length", cls_name)
validator.check_int(charge[0], N, Rel.EQ, "charge", cls_name)
validator.check_int(nl_numbers[0], N, Rel.EQ, "nl_numbers[0]", cls_name)
validator.check_int(nl_serial[0], N, Rel.LE, "nl_serial[0]", cls_name)
validator.check_int(nl_serial[1], 800, Rel.LE, "nl_serial[1]", cls_name)
validator.check_int(excluded_list_start[0], N, Rel.EQ, "excluded_list_start", cls_name)
validator.check_int(excluded_atom_numbers[0], N, Rel.EQ, "excluded_atom_numbers", cls_name)
validator.check_int(excluded_list[0], 0, Rel.GE, "excluded_list", cls_name)
return (1,), (1,), (1,), (1,) return (1,), (1,), (1,), (1,)
def infer_dtype(self, box_length, uint_crd, charge, nl_numbers, nl_serial, scaler, excluded_list_start, def infer_dtype(self, box_length, uint_crd, charge, nl_numbers, nl_serial, scaler, excluded_list_start,
@ -1701,6 +1861,8 @@ class LJEnergy(PrimitiveWithInfer):
.. math:: .. math::
dr = (x_a-x_b, y_a-y_b, z_a-z_b) dr = (x_a-x_b, y_a-y_b, z_a-z_b)
.. math::
E = A/|dr|^{12} - B/|dr|^{6} E = A/|dr|^{12} - B/|dr|^{6}
Agrs: Agrs:
@ -1716,7 +1878,9 @@ class LJEnergy(PrimitiveWithInfer):
- **nl_numbers** - (Tensor, int32) - [N,], the each atom. - **nl_numbers** - (Tensor, int32) - [N,], the each atom.
- **nl_serial** - (Tensor, int32) - [N, 800], the neighbor list of each atom, the max number is 800. - **nl_serial** - (Tensor, int32) - [N, 800], the neighbor list of each atom, the max number is 800.
- **d_LJ_A** (Tensor, float32) - [Q,], the Lennard-Jones A coefficient of each kind of atom pair. - **d_LJ_A** (Tensor, float32) - [Q,], the Lennard-Jones A coefficient of each kind of atom pair.
Q is the number of atom pair.
- **d_LJ_B** (Tensor, float32) - [Q,], the Lennard-Jones B coefficient of each kind of atom pair. - **d_LJ_B** (Tensor, float32) - [Q,], the Lennard-Jones B coefficient of each kind of atom pair.
Q is the number of atom pair.
Outputs: Outputs:
- **d_LJ_energy_atom** (Tensor, float32) - [N,], the Lennard-Jones potential energy of each atom. - **d_LJ_energy_atom** (Tensor, float32) - [N,], the Lennard-Jones potential energy of each atom.
@ -1728,8 +1892,8 @@ class LJEnergy(PrimitiveWithInfer):
@prim_attr_register @prim_attr_register
def __init__(self, atom_numbers, cutoff_square): def __init__(self, atom_numbers, cutoff_square):
assert isinstance(atom_numbers, int) validator.check_value_type('atom_numbers', atom_numbers, (int), self.name)
assert isinstance(cutoff_square, float) validator.check_value_type('cutoff_square', cutoff_square, (float), self.name)
self.atom_numbers = atom_numbers self.atom_numbers = atom_numbers
self.cutoff_square = cutoff_square self.cutoff_square = cutoff_square
self.init_prim_io_names( self.init_prim_io_names(
@ -1739,6 +1903,19 @@ class LJEnergy(PrimitiveWithInfer):
self.add_prim_attr('cutoff_square', self.cutoff_square) self.add_prim_attr('cutoff_square', self.cutoff_square)
def infer_shape(self, uint_crd, LJtype, charge, scaler, nl_numbers, nl_serial, d_LJ_A, d_LJ_B): def infer_shape(self, uint_crd, LJtype, charge, scaler, nl_numbers, nl_serial, d_LJ_A, d_LJ_B):
cls_name = self.name
N = self.atom_numbers
Q = d_LJ_A[0]
validator.check_int(uint_crd[0], N, Rel.EQ, "uint_crd[0]", cls_name)
validator.check_int(uint_crd[1], 3, Rel.EQ, "uint_crd[1]", cls_name)
validator.check_int(LJtype[0], N, Rel.EQ, "LJtype", cls_name)
validator.check_int(charge[0], 3, Rel.EQ, "charge", cls_name)
validator.check_int(scaler[0], 3, Rel.EQ, "scaler", cls_name)
validator.check_int(nl_numbers[0], N, Rel.EQ, "nl_numbers", cls_name)
validator.check_int(nl_serial[0], N, Rel.EQ, "nl_serial[0]", cls_name)
validator.check_int(nl_serial[1], 800, Rel.LE, "nl_serial[1]", cls_name)
validator.check_int(scaler[0], 3, Rel.EQ, "scaler", cls_name)
validator.check_int(d_LJ_B[0], Q, Rel.EQ, "d_LJ_B[0]", cls_name)
return charge return charge
def infer_dtype(self, uint_crd, LJtype, charge, scaler, nl_numbers, nl_serial, d_LJ_A, d_LJ_B): def infer_dtype(self, uint_crd, LJtype, charge, scaler, nl_numbers, nl_serial, d_LJ_A, d_LJ_B):
@ -1761,6 +1938,9 @@ class LJForce(PrimitiveWithInfer):
.. math:: .. math::
dr = (x_a-x_b, y_a-y_b, z_a-z_b) dr = (x_a-x_b, y_a-y_b, z_a-z_b)
.. math::
F = (-12*A/|dr|^{14} + 6*B/|dr|^{8}) * dr F = (-12*A/|dr|^{14} + 6*B/|dr|^{8}) * dr
Agrs: Agrs:
@ -1776,7 +1956,9 @@ class LJForce(PrimitiveWithInfer):
- **nl_numbers** - (Tensor, int32) - [N,], the each atom. - **nl_numbers** - (Tensor, int32) - [N,], the each atom.
- **nl_serial** - (Tensor, int32) - [N, 800], the neighbor list of each atom, the max number is 800. - **nl_serial** - (Tensor, int32) - [N, 800], the neighbor list of each atom, the max number is 800.
- **d_LJ_A** (Tensor, float32) - [Q,], the Lennard-Jones A coefficient of each kind of atom pair. - **d_LJ_A** (Tensor, float32) - [Q,], the Lennard-Jones A coefficient of each kind of atom pair.
Q is the number of atom pair.
- **d_LJ_B** (Tensor, float32) - [Q,], the Lennard-Jones B coefficient of each kind of atom pair. - **d_LJ_B** (Tensor, float32) - [Q,], the Lennard-Jones B coefficient of each kind of atom pair.
Q is the number of atom pair.
outputs: outputs:
- **frc** (Tensor, float32) - [N, 3], the force felt by each atom. - **frc** (Tensor, float32) - [N, 3], the force felt by each atom.
@ -1787,8 +1969,8 @@ class LJForce(PrimitiveWithInfer):
@prim_attr_register @prim_attr_register
def __init__(self, atom_numbers, cutoff_square): def __init__(self, atom_numbers, cutoff_square):
assert isinstance(atom_numbers, int) validator.check_value_type('atom_numbers', atom_numbers, (int), self.name)
assert isinstance(cutoff_square, float) validator.check_value_type('cutoff_square', cutoff_square, (float), self.name)
self.atom_numbers = atom_numbers self.atom_numbers = atom_numbers
self.cutoff_square = cutoff_square self.cutoff_square = cutoff_square
self.init_prim_io_names( self.init_prim_io_names(
@ -1798,6 +1980,19 @@ class LJForce(PrimitiveWithInfer):
self.add_prim_attr('cutoff_square', self.cutoff_square) self.add_prim_attr('cutoff_square', self.cutoff_square)
def infer_shape(self, uint_crd, LJtype, charge, scaler, nl_numbers, nl_serial, d_LJ_A, d_LJ_B): def infer_shape(self, uint_crd, LJtype, charge, scaler, nl_numbers, nl_serial, d_LJ_A, d_LJ_B):
cls_name = self.name
N = self.atom_numbers
Q = d_LJ_A[0]
validator.check_int(uint_crd[0], N, Rel.EQ, "uint_crd[0]", cls_name)
validator.check_int(uint_crd[1], 3, Rel.EQ, "uint_crd[1]", cls_name)
validator.check_int(LJtype[0], N, Rel.EQ, "LJtype", cls_name)
validator.check_int(charge[0], 3, Rel.EQ, "charge", cls_name)
validator.check_int(scaler[0], 3, Rel.EQ, "scaler", cls_name)
validator.check_int(nl_numbers[0], N, Rel.EQ, "nl_numbers", cls_name)
validator.check_int(nl_serial[0], N, Rel.EQ, "nl_serial[0]", cls_name)
validator.check_int(nl_serial[1], 800, Rel.LE, "nl_serial[1]", cls_name)
validator.check_int(scaler[0], 3, Rel.EQ, "scaler", cls_name)
validator.check_int(d_LJ_B[0], Q, Rel.EQ, "d_LJ_B[0]", cls_name)
return uint_crd return uint_crd
def infer_dtype(self, uint_crd, LJtype, charge, scaler, nl_numbers, nl_serial, d_LJ_A, d_LJ_B): def infer_dtype(self, uint_crd, LJtype, charge, scaler, nl_numbers, nl_serial, d_LJ_A, d_LJ_B):
@ -1833,7 +2028,9 @@ class LJForceWithPMEDirectForce(PrimitiveWithInfer):
- **nl_numbers** - (Tensor, int32) - [N,], the each atom. - **nl_numbers** - (Tensor, int32) - [N,], the each atom.
- **nl_serial** - (Tensor, int32) - [N, 800], the neighbor list of each atom, the max number is 800. - **nl_serial** - (Tensor, int32) - [N, 800], the neighbor list of each atom, the max number is 800.
- **d_LJ_A** (Tensor, float32) - [Q,], the Lennard-Jones A coefficient of each kind of atom pair. - **d_LJ_A** (Tensor, float32) - [Q,], the Lennard-Jones A coefficient of each kind of atom pair.
Q is the number of atom pair.
- **d_LJ_B** (Tensor, float32) - [Q,], the Lennard-Jones B coefficient of each kind of atom pair. - **d_LJ_B** (Tensor, float32) - [Q,], the Lennard-Jones B coefficient of each kind of atom pair.
Q is the number of atom pair.
Outputs: Outputs:
- **frc** (Tensor, float32), [N, 3], the force felt by each atom. - **frc** (Tensor, float32), [N, 3], the force felt by each atom.
@ -1844,9 +2041,9 @@ class LJForceWithPMEDirectForce(PrimitiveWithInfer):
@prim_attr_register @prim_attr_register
def __init__(self, atom_numbers, cutoff, pme_beta): def __init__(self, atom_numbers, cutoff, pme_beta):
assert isinstance(atom_numbers, int) validator.check_value_type('atom_numbers', atom_numbers, (int), self.name)
assert isinstance(cutoff, float) validator.check_value_type('cutoff', cutoff, (float), self.name)
assert isinstance(pme_beta, float) validator.check_value_type('pme_beta', pme_beta, (float), self.name)
self.atom_numbers = atom_numbers self.atom_numbers = atom_numbers
self.cutoff = cutoff self.cutoff = cutoff
self.pme_beta = pme_beta self.pme_beta = pme_beta
@ -1858,6 +2055,19 @@ class LJForceWithPMEDirectForce(PrimitiveWithInfer):
self.add_prim_attr('pme_beta', self.pme_beta) self.add_prim_attr('pme_beta', self.pme_beta)
def infer_shape(self, uint_crd, LJtype, charge, scaler, nl_numbers, nl_serial, d_LJ_A, d_LJ_B): def infer_shape(self, uint_crd, LJtype, charge, scaler, nl_numbers, nl_serial, d_LJ_A, d_LJ_B):
cls_name = self.name
N = self.atom_numbers
Q = d_LJ_A[0]
validator.check_int(uint_crd[0], N, Rel.EQ, "uint_crd[0]", cls_name)
validator.check_int(uint_crd[1], 3, Rel.EQ, "uint_crd[1]", cls_name)
validator.check_int(LJtype[0], N, Rel.EQ, "LJtype", cls_name)
validator.check_int(charge[0], 3, Rel.EQ, "charge", cls_name)
validator.check_int(scaler[0], 3, Rel.EQ, "scaler", cls_name)
validator.check_int(nl_numbers[0], N, Rel.EQ, "nl_numbers", cls_name)
validator.check_int(nl_serial[0], N, Rel.EQ, "nl_serial[0]", cls_name)
validator.check_int(nl_serial[1], 800, Rel.LE, "nl_serial[1]", cls_name)
validator.check_int(scaler[0], 3, Rel.EQ, "scaler", cls_name)
validator.check_int(d_LJ_B[0], Q, Rel.EQ, "d_LJ_B[0]", cls_name)
return uint_crd return uint_crd
def infer_dtype(self, uint_crd, LJtype, charge, scaler, nl_numbers, nl_serial, d_LJ_A, d_LJ_B): def infer_dtype(self, uint_crd, LJtype, charge, scaler, nl_numbers, nl_serial, d_LJ_A, d_LJ_B):
@ -1882,6 +2092,8 @@ class GetCenterOfGeometry(PrimitiveWithInfer):
@prim_attr_register @prim_attr_register
def __init__(self, center_numbers, center_numbers_inverse): def __init__(self, center_numbers, center_numbers_inverse):
validator.check_value_type('center_numbers', center_numbers, (int), self.name)
validator.check_value_type('center_numbers_inverse', center_numbers_inverse, (float), self.name)
self.center_numbers = center_numbers self.center_numbers = center_numbers
self.center_numbers_inverse = center_numbers_inverse self.center_numbers_inverse = center_numbers_inverse
self.add_prim_attr('center_numbers', self.center_numbers) self.add_prim_attr('center_numbers', self.center_numbers)
@ -1894,8 +2106,8 @@ class GetCenterOfGeometry(PrimitiveWithInfer):
cls_name = self.name cls_name = self.name
N = self.center_numbers N = self.center_numbers
validator.check_int(center_atoms_shape[0], N, Rel.EQ, "center_atoms", cls_name) validator.check_int(center_atoms_shape[0], N, Rel.EQ, "center_atoms", cls_name)
validator.check_int(crd_f_shape[0], N, Rel.EQ, "crd_f", cls_name) validator.check_int(crd_f_shape[0], N, Rel.EQ, "crd_f[0]", cls_name)
validator.check_int(crd_f_shape[1], 3, Rel.EQ, "crd_f", cls_name) validator.check_int(crd_f_shape[1], 3, Rel.EQ, "crd_f[1]", cls_name)
return [3,] return [3,]
def infer_dtype(self, center_atoms_dtype, crd_f_dtype): def infer_dtype(self, center_atoms_dtype, crd_f_dtype):
@ -1990,6 +2202,7 @@ class NeighborListUpdate(PrimitiveWithInfer):
between the unsigned int value and the real space coordinates. between the unsigned int value and the real space coordinates.
- **uint_crd** (Tensor, uint32) - [N, 3], the unsigned int coordinates value fo each atom. - **uint_crd** (Tensor, uint32) - [N, 3], the unsigned int coordinates value fo each atom.
- **gpointer** (Tensor, int32) - [G, 125], the 125 nearest neighbor grids (including self) of each grid. - **gpointer** (Tensor, int32) - [G, 125], the 125 nearest neighbor grids (including self) of each grid.
G is the number of nearest neighbor grids.
- **nl_atom_numbers** (Tensor, int32) - [N,], the number of atoms in neighbor list of each atom. - **nl_atom_numbers** (Tensor, int32) - [N,], the number of atoms in neighbor list of each atom.
- **nl_atom_serial** (Tensor, int32) - [N, L], the indices of atoms in neighbor list of each atom. - **nl_atom_serial** (Tensor, int32) - [N, L], the indices of atoms in neighbor list of each atom.
- **uint_dr_to_dr_cof** (Tensor, float32) - [3,], the scale factor between - **uint_dr_to_dr_cof** (Tensor, float32) - [3,], the scale factor between
@ -2182,13 +2395,13 @@ class MDIterationLeapFrogWithRF(PrimitiveWithInfer):
@prim_attr_register @prim_attr_register
def __init__(self, float4_numbers, atom_numbers, half_dt, dt, exp_gamma, is_max_velocity, max_velocity): def __init__(self, float4_numbers, atom_numbers, half_dt, dt, exp_gamma, is_max_velocity, max_velocity):
assert isinstance(float4_numbers, int) validator.check_value_type('float4_numbers', float4_numbers, (int), self.name)
assert isinstance(atom_numbers, int) validator.check_value_type('atom_numbers', atom_numbers, (int), self.name)
assert isinstance(half_dt, float) validator.check_value_type('half_dt', half_dt, (float), self.name)
assert isinstance(dt, float) validator.check_value_type('dt', dt, (float), self.name)
assert isinstance(exp_gamma, float) validator.check_value_type('exp_gamma', exp_gamma, (float), self.name)
assert isinstance(is_max_velocity, int) validator.check_value_type('is_max_velocity', is_max_velocity, (int), self.name)
assert isinstance(max_velocity, float) validator.check_value_type('max_velocity', max_velocity, (float), self.name)
self.float4_numbers = float4_numbers self.float4_numbers = float4_numbers
self.atom_numbers = atom_numbers self.atom_numbers = atom_numbers
self.half_dt = half_dt self.half_dt = half_dt

0
model_zoo/research/hpc/sponge/src/__init__.py Normal file
View File

2
model_zoo/research/hpc/sponge/src/nb14.py
View File

@ -82,7 +82,7 @@ class NON_BOND_14(nn.Cell):
self.cf_scale_type = [0] * self.dihedral_type_numbers self.cf_scale_type = [0] * self.dihedral_type_numbers
self.lj_scale_type = [0] * self.dihedral_type_numbers self.lj_scale_type = [0] * self.dihedral_type_numbers
self.process1(context)
self.h_atom_a = [0] * self.dihedral_numbers self.h_atom_a = [0] * self.dihedral_numbers
self.h_atom_b = [0] * self.dihedral_numbers self.h_atom_b = [0] * self.dihedral_numbers
self.h_lj_scale_factor = [0] * self.dihedral_numbers self.h_lj_scale_factor = [0] * self.dihedral_numbers

22
model_zoo/research/hpc/sponge/src/simulation_initial.py
View File

@ -18,15 +18,15 @@ import numpy as np
import mindspore.common.dtype as mstype import mindspore.common.dtype as mstype
from mindspore import Tensor, nn from mindspore import Tensor, nn
from Langevin_Liujian_md import Langevin_Liujian from .Langevin_Liujian_md import Langevin_Liujian
from angle import Angle from .angle import Angle
from bond import Bond from .bond import Bond
from dihedral import Dihedral from .dihedral import Dihedral
from lennard_jones import Lennard_Jones_Information from .lennard_jones import Lennard_Jones_Information
from md_information import md_information from .md_information import md_information
from nb14 import NON_BOND_14 from .nb14 import NON_BOND_14
from neighbor_list import nb_infomation from .neighbor_list import nb_infomation
from particle_mesh_ewald import Particle_Mesh_Ewald from .particle_mesh_ewald import Particle_Mesh_Ewald
class controller: class controller:
@ -102,7 +102,8 @@ class Simulation(nn.Cell):
nb_info = self.nb_info nb_info = self.nb_info
pme_method = self.pme_method pme_method = self.pme_method
bond_frc, _ = self.bond.Bond_Force_With_Atom_Energy(md_info.uint_crd, md_info.uint_dr_to_dr_cof) bond_frc, _ = self.bond.Bond_Force_With_Atom_Energy(md_info.uint_crd, md_info.uint_dr_to_dr_cof)
frc_t = bond_frc.asnumpy() frc_t = 0
frc_t += bond_frc.asnumpy()
angle_frc, _ = self.angle.Angle_Force_With_Atom_Energy(md_info.uint_crd, md_info.uint_dr_to_dr_cof) angle_frc, _ = self.angle.Angle_Force_With_Atom_Energy(md_info.uint_crd, md_info.uint_dr_to_dr_cof)
frc_t += angle_frc.asnumpy() frc_t += angle_frc.asnumpy()
@ -180,6 +181,7 @@ class Simulation(nn.Cell):
if md_info.mode > 0 and int(control.Command_Set["thermostat"]) == 1: if md_info.mode > 0 and int(control.Command_Set["thermostat"]) == 1:
md_info.vel, md_info.crd, md_info.frc, md_info.acc = liujian_info.MD_Iteration_Leap_Frog( md_info.vel, md_info.crd, md_info.frc, md_info.acc = liujian_info.MD_Iteration_Leap_Frog(
md_info.d_mass_inverse, md_info.vel, md_info.crd, md_info.frc) md_info.d_mass_inverse, md_info.vel, md_info.crd, md_info.frc)
self.Main_After_Iteration()
def Main_After_Iteration(self): def Main_After_Iteration(self):
"""main after iteration""" """main after iteration"""