From 28811fa958be4607231e2f42b084fd4702816a27 Mon Sep 17 00:00:00 2001 From: mamba_ni Date: Thu, 1 Apr 2021 15:05:39 +0800 Subject: [PATCH] master sponge performance --- .../sponge/common/crd_to_uint_crd_impl.cu | 51 + .../sponge/common/crd_to_uint_crd_impl.cuh | 26 + .../neighbor_list/neighbor_list_impl.cu | 21 +- .../neighbor_list/neighbor_list_impl.cuh | 2 +- ...md_iteration_leap_frog_liujian_gpu_impl.cu | 67 + ...d_iteration_leap_frog_liujian_gpu_impl.cuh | 28 + ...d_iteration_setup_random_state_gpu_impl.cu | 28 + ..._iteration_setup_random_state_gpu_impl.cuh | 23 + .../cuda_impl/sponge/pme/pme_energy_impl.cu | 105 +- .../cuda_impl/sponge/pme/pme_energy_impl.cuh | 16 +- .../sponge/pme/pme_reciprocal_force_impl.cu | 102 +- .../sponge/pme/pme_reciprocal_force_impl.cuh | 13 +- .../sponge/common/crd_to_uint_crd_kernel.cc | 27 + .../sponge/common/crd_to_uint_crd_kernel.h | 87 + .../neighbor_list_update_kernel.cc | 1 + .../neighbor_list_update_kernel.h | 15 +- .../nvtit/md_iteration_leap_frog_kernel.h | 8 - ..._iteration_leap_frog_liujian_gpu_kernel.cc | 35 + ...d_iteration_leap_frog_liujian_gpu_kernel.h | 100 + .../nvtit/md_iteration_setup_random_state.cc | 25 + .../nvtit/md_iteration_setup_random_state.h | 74 + .../gpu/sponge/pme/pme_energy_kernel.cc | 1 - .../gpu/sponge/pme/pme_energy_kernel.h | 159 +- .../sponge/pme/pme_reciprocal_force_kernel.cc | 11 +- .../sponge/pme/pme_reciprocal_force_kernel.h | 133 +- mindspore/ops/operations/__init__.py | 7 +- mindspore/ops/operations/sponge_ops.py | 5460 +++++++++-------- model_zoo/research/hpc/sponge/main.py | 86 +- model_zoo/research/hpc/sponge/src/__init__.py | 0 model_zoo/research/hpc/sponge/src/angle.py | 68 +- model_zoo/research/hpc/sponge/src/bond.py | 53 +- model_zoo/research/hpc/sponge/src/dihedral.py | 57 +- ...n_Liujian_md.py => langevin_liujian_md.py} | 44 +- .../research/hpc/sponge/src/lennard_jones.py | 44 +- .../research/hpc/sponge/src/md_information.py | 149 +- model_zoo/research/hpc/sponge/src/nb14.py | 85 +- .../research/hpc/sponge/src/neighbor_list.py | 125 +- .../hpc/sponge/src/particle_mesh_ewald.py | 50 +- .../research/hpc/sponge/src/simulation.py | 439 ++ .../hpc/sponge/src/simulation_initial.py | 245 - 40 files changed, 4381 insertions(+), 3689 deletions(-) create mode 100644 mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/common/crd_to_uint_crd_impl.cu create mode 100644 mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/common/crd_to_uint_crd_impl.cuh create mode 100644 mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/nvtit/md_iteration_leap_frog_liujian_gpu_impl.cu create mode 100644 mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/nvtit/md_iteration_leap_frog_liujian_gpu_impl.cuh create mode 100644 mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/nvtit/md_iteration_setup_random_state_gpu_impl.cu create mode 100644 mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/nvtit/md_iteration_setup_random_state_gpu_impl.cuh create mode 100644 mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/common/crd_to_uint_crd_kernel.cc create mode 100644 mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/common/crd_to_uint_crd_kernel.h create mode 100644 mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/nvtit/md_iteration_leap_frog_liujian_gpu_kernel.cc create mode 100644 mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/nvtit/md_iteration_leap_frog_liujian_gpu_kernel.h create mode 100644 mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/nvtit/md_iteration_setup_random_state.cc create mode 100644 mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/nvtit/md_iteration_setup_random_state.h delete mode 100644 model_zoo/research/hpc/sponge/src/__init__.py rename model_zoo/research/hpc/sponge/src/{Langevin_Liujian_md.py => langevin_liujian_md.py} (59%) create mode 100644 model_zoo/research/hpc/sponge/src/simulation.py delete mode 100644 model_zoo/research/hpc/sponge/src/simulation_initial.py diff --git a/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/common/crd_to_uint_crd_impl.cu b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/common/crd_to_uint_crd_impl.cu new file mode 100644 index 0000000000..59f92f702c --- /dev/null +++ b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/common/crd_to_uint_crd_impl.cu @@ -0,0 +1,51 @@ +/** + * Copyright 2021 Huawei Technologies Co., Ltd + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "backend/kernel_compiler/gpu/cuda_impl/sponge/common_sponge.cuh" +#include "backend/kernel_compiler/gpu/cuda_impl/sponge/common/crd_to_uint_crd_impl.cuh" + +__global__ void Crd_To_Uint_Crd(const int atom_numbers, const VECTOR *scale_factor, const VECTOR *crd, + UNSIGNED_INT_VECTOR *uint_crd) { + int atom_i = blockDim.x * blockIdx.x + threadIdx.x; + if (atom_i < atom_numbers) { + uint_crd[atom_i].uint_x = crd[atom_i].x * scale_factor[0].x; + uint_crd[atom_i].uint_y = crd[atom_i].y * scale_factor[0].y; + uint_crd[atom_i].uint_z = crd[atom_i].z * scale_factor[0].z; + /*uint_crd[atom_i].uint_x = 2 * uint_crd[atom_i].uint_x; + uint_crd[atom_i].uint_y = 2 * uint_crd[atom_i].uint_y; + uint_crd[atom_i].uint_z = 2 * uint_crd[atom_i].uint_z;*/ + uint_crd[atom_i].uint_x = uint_crd[atom_i].uint_x << 1; + uint_crd[atom_i].uint_y = uint_crd[atom_i].uint_y << 1; + uint_crd[atom_i].uint_z = uint_crd[atom_i].uint_z << 1; + } +} + +void CrdToUintCrd(const int atom_numbers, const float *crd_to_uint_crd_cof_f, const float *crd_f, + unsigned int *uint_crd_f, cudaStream_t stream) { + VECTOR *crd = const_cast(reinterpret_cast(crd_f)); + VECTOR *crd_to_uint_crd_cof = const_cast(reinterpret_cast(crd_to_uint_crd_cof_f)); + + UNSIGNED_INT_VECTOR *uint_crd = + const_cast(reinterpret_cast(uint_crd_f)); + + Crd_To_Uint_Crd<<(atom_numbers) / 128.0), 128, 0, stream>>>( + atom_numbers, crd_to_uint_crd_cof, crd, uint_crd); + + return; +} + +void CrdToUintCrd(const int atom_numbers, const float *crd_to_uint_crd_cof_f, const float *crd_f, + unsigned int *uint_crd_f, cudaStream_t stream); diff --git a/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/common/crd_to_uint_crd_impl.cuh b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/common/crd_to_uint_crd_impl.cuh new file mode 100644 index 0000000000..8cb5b4a12e --- /dev/null +++ b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/common/crd_to_uint_crd_impl.cuh @@ -0,0 +1,26 @@ +/** + * Copyright 2021 Huawei Technologies Co., Ltd + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#ifndef MINDSPORE_CCSRC_KERNEL_GPU_CUDA_IMPL_SPONGE_CRD_TO_UINT_CRD_IMPL_H_ +#define MINDSPORE_CCSRC_KERNEL_GPU_CUDA_IMPL_SPONGE_CRD_TO_UINT_CRD_IMPL_H_ + +#include +#include "runtime/device/gpu/cuda_common.h" + +void CrdToUintCrd(const int atom_numbers, const float *crd_to_uint_crd_cof_f, const float *crd_f, + unsigned int *uint_crd_f, cudaStream_t stream); + +#endif // MINDSPORE_CCSRC_KERNEL_GPU_CUDA_IMPL_SPONGE_CRD_TO_UINT_CRD_IMPL_H_ diff --git a/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/neighbor_list/neighbor_list_impl.cu b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/neighbor_list/neighbor_list_impl.cu index c8fbf5e70f..7a7744a97c 100644 --- a/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/neighbor_list/neighbor_list_impl.cu +++ b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/neighbor_list/neighbor_list_impl.cu @@ -14,7 +14,7 @@ * limitations under the License. */ #include "backend/kernel_compiler/gpu/cuda_impl/sponge/neighbor_list/neighbor_list_impl.cuh" - +#include __global__ void Copy_List(const int element_numbers, const int *origin_list, int *list) { int i = blockDim.x * blockIdx.x + threadIdx.x; if (i < element_numbers) { @@ -387,7 +387,7 @@ __global__ void Mul_half(float *src, float *dst) { } } -void Neighbor_List_Update(int grid_numbers, int atom_numbers, int refresh_count, int refresh_interval, +void Neighbor_List_Update(int grid_numbers, int atom_numbers, int *d_refresh_count, int refresh_interval, int not_first_time, float skin, int Nxy, float cutoff_square, float cutoff_with_skin_square, int *grid_N, float *box_length, int *atom_numbers_in_grid_bucket, float *grid_length_inverse, int *atom_in_grid_serial, GRID_BUCKET *bucket, float *crd, float *old_crd, @@ -397,15 +397,22 @@ void Neighbor_List_Update(int grid_numbers, int atom_numbers, int refresh_count, int *is_need_refresh_neighbor_list, cudaStream_t stream) { if (not_first_time) { if (refresh_interval > 0) { + std::vector refresh_count_list(1); + cudaMemcpyAsync(refresh_count_list.data(), d_refresh_count, sizeof(int), cudaMemcpyDeviceToHost, stream); + cudaStreamSynchronize(stream); + int refresh_count = refresh_count_list[0]; + if (refresh_count % refresh_interval == 0) { Mul_half<<<1, 3, 0, stream>>>(crd_to_uint_crd_cof, half_crd_to_uint_crd_cof); - Refresh_Neighbor_List_No_Check( - grid_numbers, atom_numbers, skin, Nxy, cutoff_square, grid_N, box_length, atom_numbers_in_grid_bucket, - grid_length_inverse, atom_in_grid_serial, bucket, reinterpret_cast(crd), - reinterpret_cast(old_crd), crd_to_uint_crd_cof, reinterpret_cast(uint_crd), - uint_dr_to_dr_cof, gpointer, d_nl, excluded_list_start, excluded_list, excluded_numbers, stream); + Refresh_Neighbor_List_No_Check(grid_numbers, atom_numbers, skin, Nxy, cutoff_square, grid_N, box_length, + atom_numbers_in_grid_bucket, grid_length_inverse, atom_in_grid_serial, bucket, + reinterpret_cast(crd), reinterpret_cast(old_crd), + half_crd_to_uint_crd_cof, reinterpret_cast(uint_crd), + uint_dr_to_dr_cof, gpointer, d_nl, excluded_list_start, excluded_list, + excluded_numbers, stream); } refresh_count += 1; + cudaMemcpyAsync(d_refresh_count, &refresh_count, sizeof(int), cudaMemcpyHostToDevice, stream); } else { Is_need_refresh_neighbor_list_cuda<<(atom_numbers) / 128), 128, 0, stream>>>( atom_numbers, reinterpret_cast(crd), reinterpret_cast(old_crd), half_skin_square, diff --git a/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/neighbor_list/neighbor_list_impl.cuh b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/neighbor_list/neighbor_list_impl.cuh index 5458800b3a..1464add3c9 100644 --- a/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/neighbor_list/neighbor_list_impl.cuh +++ b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/neighbor_list/neighbor_list_impl.cuh @@ -48,7 +48,7 @@ void Construct_Neighbor_List(int grid_numbers, int max_neighbor_numbers, int *nl void CopyNeighborListAtomNumber(int atom_numbers, NEIGHBOR_LIST *nl, int *nl_atom_numbers, cudaStream_t stream); -void Neighbor_List_Update(int grid_numbers, int atom_numbers, int refresh_count, int refresh_interval, +void Neighbor_List_Update(int grid_numbers, int atom_numbers, int* d_refresh_count, int refresh_interval, int not_first_time, float skin, int Nxy, float cutoff_square, float cutoff_with_skin_square, int *grid_N, float *box_length, int *atom_numbers_in_grid_bucket, float *grid_length_inverse, int *atom_in_grid_serial, GRID_BUCKET *bucket, float *crd, float *old_crd, diff --git a/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/nvtit/md_iteration_leap_frog_liujian_gpu_impl.cu b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/nvtit/md_iteration_leap_frog_liujian_gpu_impl.cu new file mode 100644 index 0000000000..4c9bda4b3f --- /dev/null +++ b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/nvtit/md_iteration_leap_frog_liujian_gpu_impl.cu @@ -0,0 +1,67 @@ +/** + * Copyright 2021 Huawei Technologies Co., Ltd + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "backend/kernel_compiler/gpu/cuda_impl/sponge/nvtit/md_iteration_leap_frog_liujian_gpu_impl.cuh" +#include "backend/kernel_compiler/gpu/cuda_impl/util.cuh" +#include "backend/kernel_compiler/gpu/cuda_impl/sponge/common_sponge.cuh" + +__global__ void MD_Iteration_Leap_Frog_With_LiuJian_kernel(const int atom_numbers, const float half_dt, const float dt, + const float exp_gamma, float *inverse_mass, + float *sqrt_mass_inverse, VECTOR *vel, VECTOR *crd, + VECTOR *frc, VECTOR *acc, VECTOR *random_frc, + VECTOR *output) { + int i = blockDim.x * blockIdx.x + threadIdx.x; + + if (i < atom_numbers) { + acc[i].x = inverse_mass[i] * frc[i].x; + acc[i].y = inverse_mass[i] * frc[i].y; + acc[i].z = inverse_mass[i] * frc[i].z; + + vel[i].x = vel[i].x + dt * acc[i].x; + vel[i].y = vel[i].y + dt * acc[i].y; + vel[i].z = vel[i].z + dt * acc[i].z; + + output[i].x = crd[i].x + half_dt * vel[i].x; + output[i].y = crd[i].y + half_dt * vel[i].y; + output[i].z = crd[i].z + half_dt * vel[i].z; + + vel[i].x = exp_gamma * vel[i].x + sqrt_mass_inverse[i] * random_frc[i].x; + vel[i].y = exp_gamma * vel[i].y + sqrt_mass_inverse[i] * random_frc[i].y; + vel[i].z = exp_gamma * vel[i].z + sqrt_mass_inverse[i] * random_frc[i].z; + + output[i].x = output[i].x + half_dt * vel[i].x; + output[i].y = output[i].y + half_dt * vel[i].y; + output[i].z = output[i].z + half_dt * vel[i].z; + } +} + +void MD_Iteration_Leap_Frog_With_LiuJian(const int atom_numbers, const float half_dt, const float dt, + const float exp_gamma, int float4_numbers, float *inverse_mass, + float *sqrt_mass_inverse, float *vel, float *crd, float *frc, float *acc, + curandStatePhilox4_32_10_t *rand_state, float *rand_frc, float *output, + cudaStream_t stream) { + Rand_Normal<<(float4_numbers) / 32.), 32, 0, stream>>>(float4_numbers, rand_state, + reinterpret_cast(rand_frc)); + VECTOR *d_vel = reinterpret_cast(vel); + VECTOR *d_crd = reinterpret_cast(crd); + VECTOR *d_frc = reinterpret_cast(frc); + VECTOR *d_acc = reinterpret_cast(acc); + VECTOR *d_rand_frc = reinterpret_cast(rand_frc); + VECTOR *d_out = reinterpret_cast(output); + MD_Iteration_Leap_Frog_With_LiuJian_kernel<<(atom_numbers) / 32), 32, 0, stream>>>( + atom_numbers, half_dt, dt, exp_gamma, inverse_mass, sqrt_mass_inverse, d_vel, d_crd, d_frc, d_acc, d_rand_frc, + d_out); +} diff --git a/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/nvtit/md_iteration_leap_frog_liujian_gpu_impl.cuh b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/nvtit/md_iteration_leap_frog_liujian_gpu_impl.cuh new file mode 100644 index 0000000000..c321d51310 --- /dev/null +++ b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/nvtit/md_iteration_leap_frog_liujian_gpu_impl.cuh @@ -0,0 +1,28 @@ +/** + * Copyright 2021 Huawei Technologies Co., Ltd + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#ifndef MINDSPORE_CCSRC_KERNEL_GPU_CUDA_IMPL_SPONGE_MD_ITERATION_LEAP_FROG_LIUJIAN_GPU_IMPL_H_ +#define MINDSPORE_CCSRC_KERNEL_GPU_CUDA_IMPL_SPONGE_MD_ITERATION_LEAP_FROG_LIUJIAN_GPU_IMPL_H_ + +#include +#include "runtime/device/gpu/cuda_common.h" +void MD_Iteration_Leap_Frog_With_LiuJian(const int atom_numbers, const float half_dt, const float dt, + const float exp_gamma, int float4_numbers, float *inverse_mass, + float *sqrt_mass_inverse, float *vel, float *crd, float *frc, float *acc, + curandStatePhilox4_32_10_t *rand_state, float *rand_frc, float *output, + cudaStream_t stream); + +#endif // MINDSPORE_CCSRC_KERNEL_GPU_CUDA_IMPL_SPONGE_MD_ITERATION_LEAP_FROG_LIUJIAN_GPU_IMPL_H_ diff --git a/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/nvtit/md_iteration_setup_random_state_gpu_impl.cu b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/nvtit/md_iteration_setup_random_state_gpu_impl.cu new file mode 100644 index 0000000000..1dbd34e39b --- /dev/null +++ b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/nvtit/md_iteration_setup_random_state_gpu_impl.cu @@ -0,0 +1,28 @@ +/** + * Copyright 2021 Huawei Technologies Co., Ltd + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "backend/kernel_compiler/gpu/cuda_impl/sponge/nvtit/md_iteration_setup_random_state_gpu_impl.cuh" +#include "backend/kernel_compiler/gpu/cuda_impl/util.cuh" +#include "backend/kernel_compiler/gpu/cuda_impl/sponge/common_sponge.cuh" + +void MD_Iteration_Setup_Random_State(int float4_numbers, curandStatePhilox4_32_10_t *rand_state, int seed, + cudaStream_t stream) { + Setup_Rand_Normal_Kernel<<(float4_numbers) / 32.), 32, 0, stream>>>(float4_numbers, + rand_state, seed); +} + +void MD_Iteration_Setup_Random_State(int float4_numbers, curandStatePhilox4_32_10_t *rand_state, int seed, + cudaStream_t stream); diff --git a/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/nvtit/md_iteration_setup_random_state_gpu_impl.cuh b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/nvtit/md_iteration_setup_random_state_gpu_impl.cuh new file mode 100644 index 0000000000..7dc45c061b --- /dev/null +++ b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/nvtit/md_iteration_setup_random_state_gpu_impl.cuh @@ -0,0 +1,23 @@ +/** + * Copyright 2021 Huawei Technologies Co., Ltd + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#ifndef MINDSPORE_CCSRC_KERNEL_GPU_CUDA_IMPL_SPONGE_MD_ITERATION_SETUP_RANDOM_STATE_GPU_IMPL_H_ +#define MINDSPORE_CCSRC_KERNEL_GPU_CUDA_IMPL_SPONGE_MD_ITERATION_SETUP_RANDOM_STATE_GPU_IMPL_H_ + +#include +#include "runtime/device/gpu/cuda_common.h" +void MD_Iteration_Setup_Random_State(int float4_numbers, curandStatePhilox4_32_10_t *rand_state, int seed, + cudaStream_t stream); +#endif diff --git a/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/pme/pme_energy_impl.cu b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/pme/pme_energy_impl.cu index 7ef1132eff..9cffb10bd5 100644 --- a/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/pme/pme_energy_impl.cu +++ b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/pme/pme_energy_impl.cu @@ -93,12 +93,13 @@ __global__ void PME_Excluded_Energy_Correction(const int atom_numbers, const UNS } } -void PMEEnergy(int fftx, int ffty, int fftz, int atom_numbers, float beta, float *box_length_f, float *PME_BC, - int *pme_uxyz, float *pme_frxyz, float *PME_Q, float *pme_fq, int *PME_atom_near, int *pme_kxyz, - const int *uint_crd_f, const float *charge, int *nl_atom_numbers, int *nl_atom_serial, int *nl, - const float *scaler_f, const int *excluded_list_start, const int *excluded_list, - const int *excluded_atom_numbers, float *d_reciprocal_ene, float *d_self_ene, float *d_direct_ene, - float *d_correction_ene, cudaStream_t stream) { +void PMEEnergy(int fftx, int ffty, int fftz, int atom_numbers, float beta, float *PME_BC, int *pme_uxyz, + float *pme_frxyz, float *PME_Q, float *pme_fq, int *PME_atom_near, int *pme_kxyz, const int *uint_crd_f, + const float *charge, int *nl_atom_numbers, int *nl_atom_serial, int *nl, const float *scaler_f, + const int *excluded_list_start, const int *excluded_list, const int *excluded_atom_numbers, + float *d_reciprocal_ene, float *d_self_ene, float *d_direct_ene, float *d_correction_ene, + dim3 thread_PME, int PME_Nin, int PME_Nfft, int PME_Nall, const cufftHandle &PME_plan_r2c, + const cufftHandle &PME_plan_c2r, cudaStream_t stream) { UNSIGNED_INT_VECTOR *uint_crd = const_cast(reinterpret_cast(uint_crd_f)); VECTOR *scaler = const_cast(reinterpret_cast(scaler_f)); @@ -106,97 +107,11 @@ void PMEEnergy(int fftx, int ffty, int fftz, int atom_numbers, float beta, float NEIGHBOR_LIST *nl_a = reinterpret_cast(nl); construct_neighbor_list_kernel<<(atom_numbers) / 128), 128, 0, stream>>>( atom_numbers, max_neighbor_numbers, nl_atom_numbers, nl_atom_serial, nl_a); - std::vector h_box_length(3); - cudaMemcpyAsync(h_box_length.data(), box_length_f, sizeof(float) * h_box_length.size(), cudaMemcpyDeviceToHost, - stream); - cudaStreamSynchronize(stream); - VECTOR *box_length = reinterpret_cast(h_box_length.data()); UNSIGNED_INT_VECTOR *PME_uxyz = reinterpret_cast(pme_uxyz); UNSIGNED_INT_VECTOR *PME_kxyz = reinterpret_cast(pme_kxyz); VECTOR *PME_frxyz = reinterpret_cast(pme_frxyz); cufftComplex *PME_FQ = reinterpret_cast(pme_fq); - cufftHandle PME_plan_r2c; - cufftHandle PME_plan_c2r; - cufftPlan3d(&PME_plan_r2c, fftx, ffty, fftz, CUFFT_R2C); - cufftPlan3d(&PME_plan_c2r, fftx, ffty, fftz, CUFFT_C2R); - cufftSetStream(PME_plan_r2c, stream); - cufftSetStream(PME_plan_c2r, stream); - thread_PME.x = 8; - thread_PME.y = 8; - int PME_Nin = ffty * fftz; - int PME_Nfft = fftx * ffty * (fftz / 2 + 1); - int PME_Nall = fftx * ffty * fftz; - float volume = box_length[0].x * box_length[0].y * box_length[0].z; - - UNSIGNED_INT_VECTOR *PME_kxyz_cpu; - Malloc_Safely(reinterpret_cast(&PME_kxyz_cpu), sizeof(UNSIGNED_INT_VECTOR) * 64); - - int kx, ky, kz, kxrp, kyrp, kzrp, index; - for (kx = 0; kx < 4; kx++) { - for (ky = 0; ky < 4; ky++) { - for (kz = 0; kz < 4; kz++) { - index = kx * 16 + ky * 4 + kz; - PME_kxyz_cpu[index].uint_x = kx; - PME_kxyz_cpu[index].uint_y = ky; - PME_kxyz_cpu[index].uint_z = kz; - } - } - } - cudaMemcpyAsync(PME_kxyz, PME_kxyz_cpu, sizeof(UNSIGNED_INT_VECTOR) * 64, cudaMemcpyHostToDevice, stream); - cudaStreamSynchronize(stream); - free(PME_kxyz_cpu); - - // initial start - float *B1, *B2, *B3, *PME_BC0; - B1 = reinterpret_cast(malloc(sizeof(float) * fftx)); - B2 = reinterpret_cast(malloc(sizeof(float) * ffty)); - B3 = reinterpret_cast(malloc(sizeof(float) * fftz)); - PME_BC0 = reinterpret_cast(malloc(sizeof(float) * PME_Nfft)); - - for (kx = 0; kx < fftx; kx++) { - B1[kx] = getb(kx, fftx, 4); - } - - for (ky = 0; ky < ffty; ky++) { - B2[ky] = getb(ky, ffty, 4); - } - - for (kz = 0; kz < fftz; kz++) { - B3[kz] = getb(kz, fftz, 4); - } - float mprefactor = PI * PI / -beta / beta; - - float msq; - for (kx = 0; kx < fftx; kx++) { - kxrp = kx; - if (kx > fftx / 2) kxrp = fftx - kx; - for (ky = 0; ky < ffty; ky++) { - kyrp = ky; - if (ky > ffty / 2) kyrp = ffty - ky; - for (kz = 0; kz <= fftz / 2; kz++) { - kzrp = kz; - - msq = kxrp * kxrp / box_length[0].x / box_length[0].x + kyrp * kyrp / box_length[0].y / box_length[0].y + - kzrp * kzrp / box_length[0].z / box_length[0].z; - index = kx * ffty * (fftz / 2 + 1) + ky * (fftz / 2 + 1) + kz; - if ((kx + ky + kz) == 0) { - PME_BC0[index] = 0; - } else { - PME_BC0[index] = 1.0 / PI / msq * exp(mprefactor * msq) / volume; - } - - PME_BC0[index] *= B1[kx] * B2[ky] * B3[kz]; - } - } - } - - cudaMemcpyAsync(PME_BC, PME_BC0, sizeof(float) * PME_Nfft, cudaMemcpyHostToDevice, stream); - cudaStreamSynchronize(stream); - free(B1); - free(B2); - free(B3); - free(PME_BC0); Reset_List<<<3 * atom_numbers / 32 + 1, 32, 0, stream>>>(3 * atom_numbers, reinterpret_cast(PME_uxyz), 1 << 30); @@ -226,9 +141,3 @@ void PMEEnergy(int fftx, int ffty, int fftz, int atom_numbers, float beta, float d_correction_ene); return; } -void PMEEnergy(int fftx, int ffty, int fftz, int atom_numbers, float beta, float *box_length_f, float *PME_BC, - int *pme_uxyz, float *pme_frxyz, float *PME_Q, float *pme_fq, int *PME_atom_near, int *pme_kxyz, - const int *uint_crd_f, const float *charge, int *nl_atom_numbers, int *nl_atom_serial, int *nl, - const float *scaler_f, const int *excluded_list_start, const int *excluded_list, - const int *excluded_atom_numbers, float *d_reciprocal_ene, float *d_self_ene, float *d_direct_ene, - float *d_correction_ene, cudaStream_t stream); diff --git a/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/pme/pme_energy_impl.cuh b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/pme/pme_energy_impl.cuh index 90457f1c90..333d4e125b 100644 --- a/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/pme/pme_energy_impl.cuh +++ b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/pme/pme_energy_impl.cuh @@ -16,15 +16,15 @@ #ifndef MINDSPORE_CCSRC_KERNEL_GPU_CUDA_IMPL_SPONGE_PME_PME_ENERGY_IMPL_H_ #define MINDSPORE_CCSRC_KERNEL_GPU_CUDA_IMPL_SPONGE_PME_PME_ENERGY_IMPL_H_ -#include -#include +#include #include "runtime/device/gpu/cuda_common.h" -void PMEEnergy(int fftx, int ffty, int fftz, int atom_numbers, float beta, float *box_length_f, float *PME_BC, - int *pme_uxyz, float *pme_frxyz, float *PME_Q, float *pme_fq, int *PME_atom_near, int *pme_kxyz, - const int *uint_crd_f, const float *charge, int *nl_atom_numbers, int *nl_atom_serial, int *nl, - const float *scaler_f, const int *excluded_list_start, const int *excluded_list, - const int *excluded_atom_numbers, float *d_reciprocal_ene, float *d_self_ene, float *d_direct_ene, - float *d_correction_ene, cudaStream_t stream); +void PMEEnergy(int fftx, int ffty, int fftz, int atom_numbers, float beta, float *PME_BC, int *pme_uxyz, + float *pme_frxyz, float *PME_Q, float *pme_fq, int *PME_atom_near, int *pme_kxyz, const int *uint_crd_f, + const float *charge, int *nl_atom_numbers, int *nl_atom_serial, int *nl, const float *scaler_f, + const int *excluded_list_start, const int *excluded_list, const int *excluded_atom_numbers, + float *d_reciprocal_ene, float *d_self_ene, float *d_direct_ene, float *d_correction_ene, + dim3 thread_PME, int PME_Nin, int PME_Nfft, int PME_Nall, const cufftHandle &PME_plan_r2c, + const cufftHandle &PME_plan_c2r, cudaStream_t stream); #endif diff --git a/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/pme/pme_reciprocal_force_impl.cu b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/pme/pme_reciprocal_force_impl.cu index c7e089c2f0..b9cb2f29d0 100644 --- a/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/pme/pme_reciprocal_force_impl.cu +++ b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/pme/pme_reciprocal_force_impl.cu @@ -28,7 +28,7 @@ __global__ void PME_BCFQ(cufftComplex *PME_FQ, float *PME_BC, int PME_Nfft) { __global__ void PME_Final(int *PME_atom_near, const float *charge, const float *PME_Q, VECTOR *force, const VECTOR *PME_frxyz, const UNSIGNED_INT_VECTOR *PME_kxyz, - const VECTOR PME_inverse_box_vector, const int atom_numbers) { + const _VECTOR PME_inverse_box_vector, const int atom_numbers) { int atom = blockDim.x * blockIdx.x + threadIdx.x; if (atom < atom_numbers) { int k, kx; @@ -73,8 +73,9 @@ __global__ void PME_Final(int *PME_atom_near, const float *charge, const float * void PMEReciprocalForce(int fftx, int ffty, int fftz, int atom_numbers, float beta, float *PME_BC, int *pme_uxyz, 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, - cudaStream_t stream) { + const int *uint_crd_f, const float *charge, float *force, int PME_Nin, int PME_Nall, + int PME_Nfft, const cufftHandle &PME_plan_r2c, const cufftHandle &PME_plan_c2r, + const _VECTOR &PME_inverse_box_vector, cudaStream_t stream) { Reset_List<<(3. * atom_numbers) / 128), 128, 0, stream>>>(3 * atom_numbers, force, 0.); UNSIGNED_INT_VECTOR *uint_crd = const_cast(reinterpret_cast(uint_crd_f)); @@ -86,98 +87,8 @@ void PMEReciprocalForce(int fftx, int ffty, int fftz, int atom_numbers, float be VECTOR *PME_frxyz = reinterpret_cast(pme_frxyz); VECTOR *frc = reinterpret_cast(force); - std::vector h_box_length(3); - cudaMemcpyAsync(h_box_length.data(), box_length_f, sizeof(float) * h_box_length.size(), cudaMemcpyDeviceToHost, - stream); - cudaStreamSynchronize(stream); - VECTOR *box_length = const_cast(reinterpret_cast(h_box_length.data())); cufftComplex *PME_FQ = reinterpret_cast(pme_fq); - VECTOR PME_inverse_box_vector; - PME_inverse_box_vector.x = static_cast(fftx) / box_length[0].x; - PME_inverse_box_vector.y = static_cast(ffty) / box_length[0].y; - PME_inverse_box_vector.z = static_cast(fftz) / box_length[0].z; - cufftHandle PME_plan_r2c; - cufftHandle PME_plan_c2r; - cufftPlan3d(&PME_plan_r2c, fftx, ffty, fftz, CUFFT_R2C); - cufftPlan3d(&PME_plan_c2r, fftx, ffty, fftz, CUFFT_C2R); - cufftSetStream(PME_plan_r2c, stream); - cufftSetStream(PME_plan_c2r, stream); - thread_PME.x = 8; - thread_PME.y = 8; - int PME_Nin = ffty * fftz; - int PME_Nfft = fftx * ffty * (fftz / 2 + 1); - int PME_Nall = fftx * ffty * fftz; - float volume = box_length[0].x * box_length[0].y * box_length[0].z; - - UNSIGNED_INT_VECTOR *PME_kxyz_cpu; - Malloc_Safely(reinterpret_cast(&PME_kxyz_cpu), sizeof(UNSIGNED_INT_VECTOR) * 64); - - int kx, ky, kz, kxrp, kyrp, kzrp, index; - for (kx = 0; kx < 4; kx++) { - for (ky = 0; ky < 4; ky++) { - for (kz = 0; kz < 4; kz++) { - index = kx * 16 + ky * 4 + kz; - PME_kxyz_cpu[index].uint_x = kx; - PME_kxyz_cpu[index].uint_y = ky; - PME_kxyz_cpu[index].uint_z = kz; - } - } - } - cudaMemcpyAsync(PME_kxyz, PME_kxyz_cpu, sizeof(UNSIGNED_INT_VECTOR) * 64, cudaMemcpyHostToDevice, stream); - cudaStreamSynchronize(stream); - free(PME_kxyz_cpu); - - // initial start - float *B1, *B2, *B3, *PME_BC0; - B1 = reinterpret_cast(malloc(sizeof(float) * fftx)); - B2 = reinterpret_cast(malloc(sizeof(float) * ffty)); - B3 = reinterpret_cast(malloc(sizeof(float) * fftz)); - PME_BC0 = reinterpret_cast(malloc(sizeof(float) * PME_Nfft)); - - for (kx = 0; kx < fftx; kx++) { - B1[kx] = getb(kx, fftx, 4); - } - - for (ky = 0; ky < ffty; ky++) { - B2[ky] = getb(ky, ffty, 4); - } - - for (kz = 0; kz < fftz; kz++) { - B3[kz] = getb(kz, fftz, 4); - } - float mprefactor = PI * PI / -beta / beta; - float msq; - for (kx = 0; kx < fftx; kx++) { - kxrp = kx; - if (kx > fftx / 2) kxrp = fftx - kx; - for (ky = 0; ky < ffty; ky++) { - kyrp = ky; - if (ky > ffty / 2) kyrp = ffty - ky; - for (kz = 0; kz <= fftz / 2; kz++) { - kzrp = kz; - - msq = kxrp * kxrp / box_length[0].x / box_length[0].x + kyrp * kyrp / box_length[0].y / box_length[0].y + - kzrp * kzrp / box_length[0].z / box_length[0].z; - index = kx * ffty * (fftz / 2 + 1) + ky * (fftz / 2 + 1) + kz; - if ((kx + ky + kz) == 0) { - PME_BC0[index] = 0; - } else { - PME_BC0[index] = 1.0 / PI / msq * exp(mprefactor * msq) / volume; - } - - PME_BC0[index] *= B1[kx] * B2[ky] * B3[kz]; - } - } - } - - cudaMemcpyAsync(PME_BC, PME_BC0, sizeof(float) * PME_Nfft, cudaMemcpyHostToDevice, stream); - cudaStreamSynchronize(stream); - free(B1); - free(B2); - free(B3); - free(PME_BC0); - // initial end Reset_List<<(3. * atom_numbers) / 128), 128, 0, stream>>>( 3 * atom_numbers, reinterpret_cast(frc), 0.); @@ -198,8 +109,3 @@ void PMEReciprocalForce(int fftx, int ffty, int fftz, int atom_numbers, float be PME_kxyz, PME_inverse_box_vector, atom_numbers); return; } - -void PMEReciprocalForce(int fftx, int ffty, int fftz, int atom_numbers, float beta, float *PME_BC, int *pme_uxyz, - 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, - cudaStream_t stream); diff --git a/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/pme/pme_reciprocal_force_impl.cuh b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/pme/pme_reciprocal_force_impl.cuh index 360ae6711f..8240f6d5b3 100644 --- a/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/pme/pme_reciprocal_force_impl.cuh +++ b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/pme/pme_reciprocal_force_impl.cuh @@ -16,13 +16,18 @@ #ifndef MINDSPORE_CCSRC_KERNEL_GPU_CUDA_IMPL_SPONGE_PME_PME_RECIPROCAL_FORCE_IMPL_H_ #define MINDSPORE_CCSRC_KERNEL_GPU_CUDA_IMPL_SPONGE_PME_PME_RECIPROCAL_FORCE_IMPL_H_ -#include -#include +#include #include "runtime/device/gpu/cuda_common.h" +struct _VECTOR { + float x; + float y; + float z; +}; void PMEReciprocalForce(int fftx, int ffty, int fftz, int atom_numbers, float beta, float *PME_BC, int *pme_uxyz, 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, - cudaStream_t stream); + const int *uint_crd_f, const float *charge, float *force, int PME_Nin, int PME_Nall, + int PME_Nfft, const cufftHandle &PME_plan_r2c, const cufftHandle &PME_plan_c2r, + const _VECTOR &PME_inverse_box_vector, cudaStream_t stream); #endif diff --git a/mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/common/crd_to_uint_crd_kernel.cc b/mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/common/crd_to_uint_crd_kernel.cc new file mode 100644 index 0000000000..f1aa894eeb --- /dev/null +++ b/mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/common/crd_to_uint_crd_kernel.cc @@ -0,0 +1,27 @@ +/** + * Copyright 2021 Huawei Technologies Co., Ltd + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "backend/kernel_compiler/gpu/sponge/common/crd_to_uint_crd_kernel.h" + +namespace mindspore { +namespace kernel { +MS_REG_GPU_KERNEL_TWO( + CrdToUintCrd, + KernelAttr().AddInputAttr(kNumberTypeFloat32).AddInputAttr(kNumberTypeFloat32).AddOutputAttr(kNumberTypeUInt32), + CrdToUintCrdGpuKernel, float, unsigned int) + +} // namespace kernel +} // namespace mindspore diff --git a/mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/common/crd_to_uint_crd_kernel.h b/mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/common/crd_to_uint_crd_kernel.h new file mode 100644 index 0000000000..01bb66c058 --- /dev/null +++ b/mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/common/crd_to_uint_crd_kernel.h @@ -0,0 +1,87 @@ +/** + * Copyright 2021 Huawei Technologies Co., Ltd + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#ifndef MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_GPU_SPONG_COMMON_CRD_TO_UINT_CRD_KERNEL_H_ +#define MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_GPU_SPONG_COMMON_CRD_TO_UINT_CRD_KERNEL_H_ + +#include +#include +#include +#include +#include "backend/kernel_compiler/gpu/gpu_kernel.h" +#include "backend/kernel_compiler/gpu/gpu_kernel_factory.h" +#include "runtime/device/gpu/cuda_common.h" +#include "backend/kernel_compiler/gpu/cuda_impl/sponge/common/crd_to_uint_crd_impl.cuh" + +namespace mindspore { +namespace kernel { + +template +class CrdToUintCrdGpuKernel : public GpuKernel { + public: + CrdToUintCrdGpuKernel() : ele_crd(1) {} + ~CrdToUintCrdGpuKernel() override = default; + + bool Init(const CNodePtr &kernel_node) override { + kernel_node_ = kernel_node; + atom_numbers = static_cast(GetAttr(kernel_node, "atom_numbers")); + + auto shape_crd_to_uint_crd_cof = AnfAlgo::GetPrevNodeOutputInferShape(kernel_node, 0); + auto shape_crd = AnfAlgo::GetPrevNodeOutputInferShape(kernel_node, 1); + + for (size_t i = 0; i < shape_crd_to_uint_crd_cof.size(); i++) + ele_crd_to_uint_crd_cof *= shape_crd_to_uint_crd_cof[i]; + for (size_t i = 0; i < shape_crd.size(); i++) ele_crd *= shape_crd[i]; + + InitSizeLists(); + return true; + } + + const std::vector &GetInputSizeList() const override { return input_size_list_; } + const std::vector &GetOutputSizeList() const override { return output_size_list_; } + const std::vector &GetWorkspaceSizeList() const override { return workspace_size_list_; } + + bool Launch(const std::vector &inputs, const std::vector &, + const std::vector &outputs, void *stream_ptr) override { + auto crd_to_uint_crd_cof = GetDeviceAddress(inputs, 0); + auto crd = GetDeviceAddress(inputs, 1); + + auto uint_crd = GetDeviceAddress(outputs, 0); + + CrdToUintCrd(atom_numbers, crd_to_uint_crd_cof, crd, uint_crd, reinterpret_cast(stream_ptr)); + + return true; + } + + protected: + void InitSizeLists() override { + input_size_list_.push_back(ele_crd_to_uint_crd_cof * sizeof(T)); + input_size_list_.push_back(ele_crd * sizeof(T)); + + output_size_list_.push_back(3 * atom_numbers * sizeof(T)); + } + + private: + size_t ele_crd_to_uint_crd_cof = 1; + size_t ele_crd = 1; + std::vector input_size_list_; + std::vector output_size_list_; + std::vector workspace_size_list_; + int atom_numbers; +}; +} // namespace kernel +} // namespace mindspore +#endif // MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_GPU_SPONG_COMMON_CRD_TO_UINT_CRD_KERNEL_H_ diff --git a/mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/neighbor_list/neighbor_list_update_kernel.cc b/mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/neighbor_list/neighbor_list_update_kernel.cc index f0e097c7a4..455eead002 100644 --- a/mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/neighbor_list/neighbor_list_update_kernel.cc +++ b/mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/neighbor_list/neighbor_list_update_kernel.cc @@ -38,6 +38,7 @@ MS_REG_GPU_KERNEL_TWO(NeighborListUpdate, .AddInputAttr(kNumberTypeInt32) .AddInputAttr(kNumberTypeInt32) .AddInputAttr(kNumberTypeInt32) + .AddInputAttr(kNumberTypeInt32) .AddOutputAttr(kNumberTypeFloat32), NeighborListUpdateGpuKernel, int, float) diff --git a/mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/neighbor_list/neighbor_list_update_kernel.h b/mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/neighbor_list/neighbor_list_update_kernel.h index 5efbc50f43..6e13cf0d6e 100644 --- a/mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/neighbor_list/neighbor_list_update_kernel.h +++ b/mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/neighbor_list/neighbor_list_update_kernel.h @@ -36,7 +36,6 @@ class NeighborListUpdateGpuKernel : public GpuKernel { bool Init(const CNodePtr &kernel_node) override { grid_numbers = static_cast(GetAttr(kernel_node, "grid_numbers")); atom_numbers = static_cast(GetAttr(kernel_node, "atom_numbers")); - refresh_count = static_cast(GetAttr(kernel_node, "refresh_count")); refresh_interval = static_cast(GetAttr(kernel_node, "refresh_interval")); not_first_time = static_cast(GetAttr(kernel_node, "not_first_time")); Nxy = static_cast(GetAttr(kernel_node, "Nxy")); @@ -47,7 +46,8 @@ class NeighborListUpdateGpuKernel : public GpuKernel { cutoff_with_skin = static_cast(GetAttr(kernel_node, "cutoff_with_skin")); half_cutoff_with_skin = static_cast(GetAttr(kernel_node, "half_cutoff_with_skin")); cutoff_with_skin_square = static_cast(GetAttr(kernel_node, "cutoff_with_skin_square")); - + h_bucket.resize(grid_numbers); + h_gpointer.resize(grid_numbers); InitSizeLists(); return true; } @@ -76,17 +76,18 @@ class NeighborListUpdateGpuKernel : public GpuKernel { auto excluded_list = GetDeviceAddress(inputs, 15); auto excluded_numbers = GetDeviceAddress(inputs, 16); auto need_refresh_flag = GetDeviceAddress(inputs, 17); + auto d_refresh_count = GetDeviceAddress(inputs, 18); GRID_BUCKET *d_bucket = reinterpret_cast(GetDeviceAddress(workspaces, 0)); GRID_POINTER *d_gpointer = reinterpret_cast(GetDeviceAddress(workspaces, 1)); NEIGHBOR_LIST *nl = GetDeviceAddress(workspaces, 2); float *half_crd_to_uint_crd_cof = GetDeviceAddress(workspaces, 3); - std::vector h_bucket(grid_numbers); + // std::vector h_bucket(grid_numbers); for (size_t i = 0; i < h_bucket.size(); i += 1) { h_bucket[i].atom_serial = bucket + i * max_atom_in_grid_numbers; } - std::vector h_gpointer(grid_numbers); + // std::vector h_gpointer(grid_numbers); for (size_t i = 0; i < h_gpointer.size(); i += 1) { h_gpointer[i].grid_serial = gpointer + i * 125; } @@ -98,7 +99,7 @@ class NeighborListUpdateGpuKernel : public GpuKernel { Construct_Neighbor_List(atom_numbers, max_neighbor_numbers, nl_atom_numbers, nl_atom_serial, nl, reinterpret_cast(stream_ptr)); - Neighbor_List_Update(grid_numbers, atom_numbers, refresh_count, refresh_interval, not_first_time, skin, Nxy, + Neighbor_List_Update(grid_numbers, atom_numbers, d_refresh_count, refresh_interval, not_first_time, skin, Nxy, cutoff_square, cutoff_with_skin_square, grid_N, box_length, atom_numbers_in_grid_bucket, grid_length_inverse, atom_in_grid_serial, d_bucket, crd, old_crd, crd_to_uint_crd_cof, half_crd_to_uint_crd_cof, uint_crd, uint_dr_to_dr_cof, d_gpointer, nl, excluded_list_start, @@ -132,6 +133,7 @@ class NeighborListUpdateGpuKernel : public GpuKernel { input_size_list_.push_back(sizeof(int) * excluded_atom_numbers); input_size_list_.push_back(sizeof(int) * atom_numbers); + input_size_list_.push_back(sizeof(int)); input_size_list_.push_back(sizeof(int)); workspace_size_list_.push_back(sizeof(GRID_BUCKET) * grid_numbers); @@ -148,7 +150,6 @@ class NeighborListUpdateGpuKernel : public GpuKernel { int not_first_time; int atom_numbers; int grid_numbers; - int refresh_count; int refresh_interval; int Nxy; int max_atom_in_grid_numbers; @@ -163,6 +164,8 @@ class NeighborListUpdateGpuKernel : public GpuKernel { std::vector input_size_list_; std::vector output_size_list_; std::vector workspace_size_list_; + std::vector h_bucket; + std::vector h_gpointer; }; } // namespace kernel } // namespace mindspore diff --git a/mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/nvtit/md_iteration_leap_frog_kernel.h b/mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/nvtit/md_iteration_leap_frog_kernel.h index c7a121ce1d..cb092c5190 100644 --- a/mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/nvtit/md_iteration_leap_frog_kernel.h +++ b/mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/nvtit/md_iteration_leap_frog_kernel.h @@ -45,14 +45,6 @@ class MDIterationLeapFrogGpuKernel : public GpuKernel { is_max_velocity = static_cast(GetAttr(kernel_node, "is_max_velocity")); max_velocity = static_cast(GetAttr(kernel_node, "max_velocity")); - // printf("float4_numbers: %d", float4_numbers); - // printf("atom_numbers: %d", atom_numbers); - // printf("half_dt: %f", half_dt); - // printf("dt: %f", dt); - // printf("exp_gamma: %f", exp_gamma); - // printf("is_max_velocity: %d", is_max_velocity); - // printf("max_velocity: %f", max_velocity); - auto shape_mass_inverse = AnfAlgo::GetPrevNodeOutputInferShape(kernel_node, 0); auto shape_qrt_mass = AnfAlgo::GetPrevNodeOutputInferShape(kernel_node, 1); diff --git a/mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/nvtit/md_iteration_leap_frog_liujian_gpu_kernel.cc b/mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/nvtit/md_iteration_leap_frog_liujian_gpu_kernel.cc new file mode 100644 index 0000000000..425eaa6214 --- /dev/null +++ b/mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/nvtit/md_iteration_leap_frog_liujian_gpu_kernel.cc @@ -0,0 +1,35 @@ +/** + * Copyright 2021 Huawei Technologies Co., Ltd + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "backend/kernel_compiler/gpu/sponge/nvtit/md_iteration_leap_frog_liujian_gpu_kernel.h" + +namespace mindspore { +namespace kernel { +MS_REG_GPU_KERNEL_TWO(MDIterationLeapFrogLiujian, + KernelAttr() + .AddInputAttr(kNumberTypeFloat32) + .AddInputAttr(kNumberTypeFloat32) + .AddInputAttr(kNumberTypeFloat32) + .AddInputAttr(kNumberTypeFloat32) + .AddInputAttr(kNumberTypeFloat32) + .AddInputAttr(kNumberTypeFloat32) + .AddInputAttr(kNumberTypeFloat32) + .AddInputAttr(kNumberTypeFloat32) + .AddOutputAttr(kNumberTypeFloat32), + MDIterationLeapFrogLiujianCudaGpuKernel, float, int) + +} // namespace kernel +} // namespace mindspore diff --git a/mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/nvtit/md_iteration_leap_frog_liujian_gpu_kernel.h b/mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/nvtit/md_iteration_leap_frog_liujian_gpu_kernel.h new file mode 100644 index 0000000000..e37f5cf3f4 --- /dev/null +++ b/mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/nvtit/md_iteration_leap_frog_liujian_gpu_kernel.h @@ -0,0 +1,100 @@ +/** + * Copyright 2021 Huawei Technologies Co., Ltd + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#ifndef MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_GPU_MD_ITERATION_LEAP_FROG_LIUJIAN_GPU_KERNEL_H_ +#define MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_GPU_MD_ITERATION_LEAP_FROG_LIUJIAN_GPU_KERNEL_H_ + +#include "backend/kernel_compiler/gpu/cuda_impl/sponge/nvtit/md_iteration_leap_frog_liujian_gpu_impl.cuh" +#include +#include +#include +#include + +#include "backend/kernel_compiler/gpu/gpu_kernel.h" +#include "backend/kernel_compiler/gpu/gpu_kernel_factory.h" +#include "runtime/device/gpu/cuda_common.h" + +namespace mindspore { +namespace kernel { +template +class MDIterationLeapFrogLiujianCudaGpuKernel : public GpuKernel { + public: + MDIterationLeapFrogLiujianCudaGpuKernel() {} + ~MDIterationLeapFrogLiujianCudaGpuKernel() override = default; + + bool Init(const CNodePtr &kernel_node) override { + // get bond_numbers + kernel_node_ = kernel_node; + atom_numbers = static_cast(GetAttr(kernel_node, "atom_numbers")); + half_dt = static_cast(GetAttr(kernel_node, "half_dt")); + dt = static_cast(GetAttr(kernel_node, "dt")); + exp_gamma = static_cast(GetAttr(kernel_node, "exp_gamma")); + float4_numbers = ceil(3. * static_cast(atom_numbers) / 4.); + InitSizeLists(); + return true; + } + + const std::vector &GetInputSizeList() const override { return input_size_list_; } + const std::vector &GetOutputSizeList() const override { return output_size_list_; } + const std::vector &GetWorkspaceSizeList() const override { return workspace_size_list_; } + + bool Launch(const std::vector &inputs, const std::vector &workspace, + const std::vector &outputs, void *stream_ptr) override { + auto inverse_mass = GetDeviceAddress(inputs, 0); + auto sqrt_mass_inverse = GetDeviceAddress(inputs, 1); + auto vel = GetDeviceAddress(inputs, 2); + auto crd = GetDeviceAddress(inputs, 3); + auto frc = GetDeviceAddress(inputs, 4); + auto acc = GetDeviceAddress(inputs, 5); + auto rand_state = GetDeviceAddress(inputs, 6); + auto rand_frc = GetDeviceAddress(inputs, 7); + + auto output = GetDeviceAddress(outputs, 0); + + MD_Iteration_Leap_Frog_With_LiuJian(atom_numbers, half_dt, dt, exp_gamma, float4_numbers, inverse_mass, + sqrt_mass_inverse, vel, crd, frc, acc, + reinterpret_cast(rand_state), rand_frc, output, + reinterpret_cast(stream_ptr)); + return true; + } + + protected: + void InitSizeLists() override { + input_size_list_.push_back(atom_numbers * sizeof(float)); + input_size_list_.push_back(atom_numbers * sizeof(float)); + input_size_list_.push_back(atom_numbers * 3 * sizeof(float)); + input_size_list_.push_back(atom_numbers * 3 * sizeof(float)); + input_size_list_.push_back(atom_numbers * 3 * sizeof(float)); + input_size_list_.push_back(atom_numbers * 3 * sizeof(float)); + input_size_list_.push_back(float4_numbers * sizeof(curandStatePhilox4_32_10_t)); + input_size_list_.push_back(atom_numbers * 3 * sizeof(float)); + + output_size_list_.push_back(atom_numbers * 3 * sizeof(T)); + } + + private: + std::vector input_size_list_; + std::vector output_size_list_; + std::vector workspace_size_list_; + int atom_numbers; + float half_dt; + float dt; + float exp_gamma; + int float4_numbers; +}; +} // namespace kernel +} // namespace mindspore +#endif diff --git a/mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/nvtit/md_iteration_setup_random_state.cc b/mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/nvtit/md_iteration_setup_random_state.cc new file mode 100644 index 0000000000..37c5289c3f --- /dev/null +++ b/mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/nvtit/md_iteration_setup_random_state.cc @@ -0,0 +1,25 @@ +/** + * Copyright 2021 Huawei Technologies Co., Ltd + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "backend/kernel_compiler/gpu/sponge/nvtit/md_iteration_setup_random_state.h" + +namespace mindspore { +namespace kernel { +MS_REG_GPU_KERNEL_TWO(MDIterationSetupRandState, KernelAttr().AddOutputAttr(kNumberTypeFloat32), + MDIterationSetupRandStateGpuKernel, float, int) + +} // namespace kernel +} // namespace mindspore diff --git a/mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/nvtit/md_iteration_setup_random_state.h b/mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/nvtit/md_iteration_setup_random_state.h new file mode 100644 index 0000000000..bbc3dc2513 --- /dev/null +++ b/mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/nvtit/md_iteration_setup_random_state.h @@ -0,0 +1,74 @@ +/** + * Copyright 2021 Huawei Technologies Co., Ltd + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#ifndef MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_GPU_MD_ITERATION_SETUP_RANDOM_STATE_GPU_KERNEL_H_ +#define MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_GPU_MD_ITERATION_SETUP_RANDOM_STATE_GPU_KERNEL_H_ + +#include "backend/kernel_compiler/gpu/cuda_impl/sponge/nvtit/md_iteration_setup_random_state_gpu_impl.cuh" +#include +#include +#include +#include + +#include "backend/kernel_compiler/gpu/gpu_kernel.h" +#include "backend/kernel_compiler/gpu/gpu_kernel_factory.h" +#include "runtime/device/gpu/cuda_common.h" + +namespace mindspore { +namespace kernel { +template +class MDIterationSetupRandStateGpuKernel : public GpuKernel { + public: + MDIterationSetupRandStateGpuKernel() {} + ~MDIterationSetupRandStateGpuKernel() override = default; + + bool Init(const CNodePtr &kernel_node) override { + // get bond_numbers + kernel_node_ = kernel_node; + atom_numbers = static_cast(GetAttr(kernel_node, "atom_numbers")); + seed = static_cast(GetAttr(kernel_node, "seed")); + float4_numbers = ceil(3. * static_cast(atom_numbers) / 4.); + InitSizeLists(); + return true; + } + + const std::vector &GetInputSizeList() const override { return input_size_list_; } + const std::vector &GetOutputSizeList() const override { return output_size_list_; } + const std::vector &GetWorkspaceSizeList() const override { return workspace_size_list_; } + + bool Launch(const std::vector &inputs, const std::vector &, + const std::vector &outputs, void *stream_ptr) override { + auto output = GetDeviceAddress(outputs, 0); + curandStatePhilox4_32_10_t *rand_state = reinterpret_cast(output); + MD_Iteration_Setup_Random_State(float4_numbers, rand_state, seed, reinterpret_cast(stream_ptr)); + + return true; + } + + protected: + void InitSizeLists() override { output_size_list_.push_back(sizeof(curandStatePhilox4_32_10_t) * float4_numbers); } + + private: + std::vector input_size_list_; + std::vector output_size_list_; + std::vector workspace_size_list_; + int atom_numbers; + int seed; + int float4_numbers; +}; +} // namespace kernel +} // namespace mindspore +#endif diff --git a/mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/pme/pme_energy_kernel.cc b/mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/pme/pme_energy_kernel.cc index 94e01ff08d..8c96955d27 100644 --- a/mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/pme/pme_energy_kernel.cc +++ b/mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/pme/pme_energy_kernel.cc @@ -19,7 +19,6 @@ namespace mindspore { namespace kernel { MS_REG_GPU_KERNEL_TWO(PMEEnergy, KernelAttr() - .AddInputAttr(kNumberTypeFloat32) .AddInputAttr(kNumberTypeUInt32) .AddInputAttr(kNumberTypeFloat32) .AddInputAttr(kNumberTypeInt32) diff --git a/mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/pme/pme_energy_kernel.h b/mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/pme/pme_energy_kernel.h index 4f151040ba..d51895e0f0 100644 --- a/mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/pme/pme_energy_kernel.h +++ b/mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/pme/pme_energy_kernel.h @@ -18,8 +18,6 @@ #include #include #include -#include -#include #include "backend/kernel_compiler/gpu/gpu_kernel.h" #include "backend/kernel_compiler/gpu/gpu_kernel_factory.h" #include "runtime/device/gpu/cuda_common.h" @@ -40,8 +38,76 @@ class PMEEnergyGpuKernel : public GpuKernel { fftx = static_cast(GetAttr(kernel_node, "fftx")); ffty = static_cast(GetAttr(kernel_node, "ffty")); fftz = static_cast(GetAttr(kernel_node, "fftz")); - PME_Nall = fftx * ffty * fftz; + + float box_length_0 = static_cast(GetAttr(kernel_node, "box_length_0")); + float box_length_1 = static_cast(GetAttr(kernel_node, "box_length_1")); + float box_length_2 = static_cast(GetAttr(kernel_node, "box_length_2")); + std::vector h_box_length(3); + h_box_length[0] = box_length_0; + h_box_length[1] = box_length_1; + h_box_length[2] = box_length_2; + VECTOR *box_length = reinterpret_cast(h_box_length.data()); + cufftPlan3d(&PME_plan_r2c, fftx, ffty, fftz, CUFFT_R2C); + cufftPlan3d(&PME_plan_c2r, fftx, ffty, fftz, CUFFT_C2R); + _thread_PME.x = 8; + _thread_PME.y = 8; + PME_Nin = ffty * fftz; PME_Nfft = fftx * ffty * (fftz / 2 + 1); + PME_Nall = fftx * ffty * fftz; + PME_kxyz_cpu.resize(64); + volume = box_length[0].x * box_length[0].y * box_length[0].z; + int kx, ky, kz, kxrp, kyrp, kzrp, index; + for (kx = 0; kx < 4; kx++) { + for (ky = 0; ky < 4; ky++) { + for (kz = 0; kz < 4; kz++) { + index = kx * 16 + ky * 4 + kz; + PME_kxyz_cpu[index].uint_x = kx; + PME_kxyz_cpu[index].uint_y = ky; + PME_kxyz_cpu[index].uint_z = kz; + } + } + } + + B1.resize(fftx); + B2.resize(ffty); + B3.resize(fftz); + PME_BC0.resize(PME_Nfft); + for (kx = 0; kx < fftx; kx++) { + B1[kx] = getb(kx, fftx, 4); + } + + for (ky = 0; ky < ffty; ky++) { + B2[ky] = getb(ky, ffty, 4); + } + + for (kz = 0; kz < fftz; kz++) { + B3[kz] = getb(kz, fftz, 4); + } + float mprefactor = PI * PI / -beta / beta; + + float msq; + for (kx = 0; kx < fftx; kx++) { + kxrp = kx; + if (kx > fftx / 2) kxrp = fftx - kx; + for (ky = 0; ky < ffty; ky++) { + kyrp = ky; + if (ky > ffty / 2) kyrp = ffty - ky; + for (kz = 0; kz <= fftz / 2; kz++) { + kzrp = kz; + + msq = kxrp * kxrp / box_length[0].x / box_length[0].x + kyrp * kyrp / box_length[0].y / box_length[0].y + + kzrp * kzrp / box_length[0].z / box_length[0].z; + index = kx * ffty * (fftz / 2 + 1) + ky * (fftz / 2 + 1) + kz; + if ((kx + ky + kz) == 0) { + PME_BC0[index] = 0; + } else { + PME_BC0[index] = 1.0 / PI / msq * exp(mprefactor * msq) / volume; + } + + PME_BC0[index] *= B1[kx] * B2[ky] * B3[kz]; + } + } + } InitSizeLists(); return true; @@ -53,15 +119,14 @@ class PMEEnergyGpuKernel : public GpuKernel { bool Launch(const std::vector &inputs, const std::vector &workspace, const std::vector &outputs, void *stream_ptr) override { - auto boxlength = GetDeviceAddress(inputs, 0); - auto uint_crd = GetDeviceAddress(inputs, 1); - auto charge = GetDeviceAddress(inputs, 2); - auto nl_numbers = GetDeviceAddress(inputs, 3); - auto nl_serial = GetDeviceAddress(inputs, 4); - auto scaler = GetDeviceAddress(inputs, 5); - auto excluded_list_start = GetDeviceAddress(inputs, 6); - auto excluded_list = GetDeviceAddress(inputs, 7); - auto excluded_atom_numbers = GetDeviceAddress(inputs, 8); + auto uint_crd = GetDeviceAddress(inputs, 0); + auto charge = GetDeviceAddress(inputs, 1); + auto nl_numbers = GetDeviceAddress(inputs, 2); + auto nl_serial = GetDeviceAddress(inputs, 3); + auto scaler = GetDeviceAddress(inputs, 4); + auto excluded_list_start = GetDeviceAddress(inputs, 5); + auto excluded_list = GetDeviceAddress(inputs, 6); + auto excluded_atom_numbers = GetDeviceAddress(inputs, 7); auto pme_uxyz = GetDeviceAddress(workspace, 0); // workspace auto pme_frxyz = GetDeviceAddress(workspace, 1); // workspace @@ -77,16 +142,22 @@ class PMEEnergyGpuKernel : public GpuKernel { auto direct_ene = GetDeviceAddress(outputs, 2); auto correction_ene = GetDeviceAddress(outputs, 3); - PMEEnergy(fftx, ffty, fftz, atom_numbers, beta, boxlength, pme_bc, pme_uxyz, pme_frxyz, pme_q, pme_fq, - pme_atom_near, pme_kxyz, uint_crd, charge, nl_numbers, nl_serial, nl, scaler, excluded_list_start, - excluded_list, excluded_atom_numbers, reciprocal_ene, self_ene, direct_ene, correction_ene, - reinterpret_cast(stream_ptr)); + cufftSetStream(PME_plan_r2c, reinterpret_cast(stream_ptr)); + cufftSetStream(PME_plan_c2r, reinterpret_cast(stream_ptr)); + cudaMemcpyAsync(pme_kxyz, PME_kxyz_cpu.data(), sizeof(UNSIGNED_INT_VECTOR) * 64, cudaMemcpyHostToDevice, + reinterpret_cast(stream_ptr)); + cudaMemcpyAsync(pme_bc, PME_BC0.data(), sizeof(float) * PME_Nfft, cudaMemcpyHostToDevice, + reinterpret_cast(stream_ptr)); + + PMEEnergy(fftx, ffty, fftz, atom_numbers, beta, pme_bc, pme_uxyz, pme_frxyz, pme_q, pme_fq, pme_atom_near, pme_kxyz, + uint_crd, charge, nl_numbers, nl_serial, nl, scaler, excluded_list_start, excluded_list, + excluded_atom_numbers, reciprocal_ene, self_ene, direct_ene, correction_ene, _thread_PME, PME_Nin, + PME_Nfft, PME_Nall, PME_plan_r2c, PME_plan_c2r, reinterpret_cast(stream_ptr)); return true; } protected: void InitSizeLists() override { - input_size_list_.push_back(sizeof(VECTOR)); input_size_list_.push_back(atom_numbers * sizeof(UNSIGNED_INT_VECTOR)); input_size_list_.push_back(atom_numbers * sizeof(VECTOR)); input_size_list_.push_back(atom_numbers * sizeof(T1)); @@ -112,12 +183,56 @@ class PMEEnergyGpuKernel : public GpuKernel { output_size_list_.push_back(sizeof(T)); } + cufftComplex expc(cufftComplex z) { + cufftComplex res; + float t = expf(z.x); + sincosf(z.y, &res.y, &res.x); + res.x *= t; + res.y *= t; + return res; + } + float M_(float u, int n) { + if (n == 2) { + if (u > 2 || u < 0) return 0; + return 1 - abs(u - 1); + } else { + return u / (n - 1) * M_(u, n - 1) + (n - u) / (n - 1) * M_(u - 1, n - 1); + } + } + float getb(int k, int NFFT, int B_order) { + cufftComplex tempc, tempc2, res; + float tempf; + tempc2.x = 0; + tempc2.y = 0; + + tempc.x = 0; + tempc.y = 2 * (B_order - 1) * PI * k / NFFT; + res = expc(tempc); + + for (int kk = 0; kk < (B_order - 1); kk++) { + tempc.x = 0; + tempc.y = 2 * PI * k / NFFT * kk; + tempc = expc(tempc); + tempf = M_(kk + 1, B_order); + tempc2.x += tempf * tempc.x; + tempc2.y += tempf * tempc.y; + } + res = cuCdivf(res, tempc2); + return res.x * res.x + res.y * res.y; + } + private: size_t ele_uint_crd = 1; std::vector input_size_list_; std::vector output_size_list_; std::vector workspace_size_list_; + + std::vector B1; + std::vector B2; + std::vector B3; + std::vector PME_BC0; + int atom_numbers; int excluded_numbers; int max_nl_numbers = 800; @@ -125,8 +240,16 @@ class PMEEnergyGpuKernel : public GpuKernel { int ffty; int fftz; float beta; + int PME_Nin; int PME_Nall; int PME_Nfft; + float volume; + float PI = 3.1415926; + cufftHandle PME_plan_r2c; + cufftHandle PME_plan_c2r; + + dim3 _thread_PME; + struct VECTOR { float x; float y; @@ -138,7 +261,7 @@ class PMEEnergyGpuKernel : public GpuKernel { unsigned int uint_y; unsigned int uint_z; }; - + std::vector PME_kxyz_cpu; struct NEIGHBOR_LIST { int atom_numbers; int *atom_serial; diff --git a/mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/pme/pme_reciprocal_force_kernel.cc b/mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/pme/pme_reciprocal_force_kernel.cc index 94c2e7130f..8ca56eb00b 100644 --- a/mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/pme/pme_reciprocal_force_kernel.cc +++ b/mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/pme/pme_reciprocal_force_kernel.cc @@ -17,13 +17,10 @@ namespace mindspore { namespace kernel { -MS_REG_GPU_KERNEL_TWO(PMEReciprocalForce, - KernelAttr() - .AddInputAttr(kNumberTypeFloat32) - .AddInputAttr(kNumberTypeUInt32) - .AddInputAttr(kNumberTypeFloat32) - .AddOutputAttr(kNumberTypeFloat32), - PMEReciprocalForceGpuKernel, float, int) +MS_REG_GPU_KERNEL_TWO( + PMEReciprocalForce, + KernelAttr().AddInputAttr(kNumberTypeUInt32).AddInputAttr(kNumberTypeFloat32).AddOutputAttr(kNumberTypeFloat32), + PMEReciprocalForceGpuKernel, float, int) } // namespace kernel } // namespace mindspore diff --git a/mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/pme/pme_reciprocal_force_kernel.h b/mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/pme/pme_reciprocal_force_kernel.h index 161fb6ccc5..7bf21d136a 100644 --- a/mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/pme/pme_reciprocal_force_kernel.h +++ b/mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/pme/pme_reciprocal_force_kernel.h @@ -41,6 +41,75 @@ class PMEReciprocalForceGpuKernel : public GpuKernel { fftz = static_cast(GetAttr(kernel_node, "fftz")); PME_Nall = fftx * ffty * fftz; PME_Nfft = fftx * ffty * (fftz / 2 + 1); + PME_Nin = ffty * fftz; + + float box_length_0 = static_cast(GetAttr(kernel_node, "box_length_0")); + float box_length_1 = static_cast(GetAttr(kernel_node, "box_length_1")); + float box_length_2 = static_cast(GetAttr(kernel_node, "box_length_2")); + std::vector h_box_length(3); + h_box_length[0] = box_length_0; + h_box_length[1] = box_length_1; + h_box_length[2] = box_length_2; + VECTOR *box_length = reinterpret_cast(h_box_length.data()); + PME_inverse_box_vector.x = static_cast(fftx) / box_length[0].x; + PME_inverse_box_vector.y = static_cast(ffty) / box_length[0].y; + PME_inverse_box_vector.z = static_cast(fftz) / box_length[0].z; + cufftPlan3d(&PME_plan_r2c, fftx, ffty, fftz, CUFFT_R2C); + cufftPlan3d(&PME_plan_c2r, fftx, ffty, fftz, CUFFT_C2R); + + float volume = box_length[0].x * box_length[0].y * box_length[0].z; + PME_kxyz_cpu.resize(64); + int kx, ky, kz, kxrp, kyrp, kzrp, index; + for (kx = 0; kx < 4; kx++) { + for (ky = 0; ky < 4; ky++) { + for (kz = 0; kz < 4; kz++) { + index = kx * 16 + ky * 4 + kz; + PME_kxyz_cpu[index].uint_x = kx; + PME_kxyz_cpu[index].uint_y = ky; + PME_kxyz_cpu[index].uint_z = kz; + } + } + } + B1.resize(fftx); + B2.resize(ffty); + B3.resize(fftz); + PME_BC0.resize(PME_Nfft); + + for (kx = 0; kx < fftx; kx++) { + B1[kx] = getb(kx, fftx, 4); + } + + for (ky = 0; ky < ffty; ky++) { + B2[ky] = getb(ky, ffty, 4); + } + + for (kz = 0; kz < fftz; kz++) { + B3[kz] = getb(kz, fftz, 4); + } + float mprefactor = PI * PI / -beta / beta; + float msq; + for (kx = 0; kx < fftx; kx++) { + kxrp = kx; + if (kx > fftx / 2) kxrp = fftx - kx; + for (ky = 0; ky < ffty; ky++) { + kyrp = ky; + if (ky > ffty / 2) kyrp = ffty - ky; + for (kz = 0; kz <= fftz / 2; kz++) { + kzrp = kz; + + msq = kxrp * kxrp / box_length[0].x / box_length[0].x + kyrp * kyrp / box_length[0].y / box_length[0].y + + kzrp * kzrp / box_length[0].z / box_length[0].z; + index = kx * ffty * (fftz / 2 + 1) + ky * (fftz / 2 + 1) + kz; + if ((kx + ky + kz) == 0) { + PME_BC0[index] = 0; + } else { + PME_BC0[index] = 1.0 / PI / msq * exp(mprefactor * msq) / volume; + } + + PME_BC0[index] *= B1[kx] * B2[ky] * B3[kz]; + } + } + } InitSizeLists(); return true; @@ -52,9 +121,8 @@ class PMEReciprocalForceGpuKernel : public GpuKernel { bool Launch(const std::vector &inputs, const std::vector &workspace, const std::vector &outputs, void *stream_ptr) override { - auto boxlength = GetDeviceAddress(inputs, 0); - auto uint_crd = GetDeviceAddress(inputs, 1); - auto charge = GetDeviceAddress(inputs, 2); + auto uint_crd = GetDeviceAddress(inputs, 0); + auto charge = GetDeviceAddress(inputs, 1); auto pme_uxyz = GetDeviceAddress(workspace, 0); // workspace auto pme_frxyz = GetDeviceAddress(workspace, 1); // workspace @@ -65,9 +133,15 @@ class PMEReciprocalForceGpuKernel : public GpuKernel { auto pme_kxyz = GetDeviceAddress(workspace, 6); // workspace auto force = GetDeviceAddress(outputs, 0); - + cufftSetStream(PME_plan_r2c, reinterpret_cast(stream_ptr)); + cufftSetStream(PME_plan_c2r, reinterpret_cast(stream_ptr)); + cudaMemcpyAsync(pme_kxyz, PME_kxyz_cpu.data(), sizeof(UNSIGNED_INT_VECTOR) * 64, cudaMemcpyHostToDevice, + reinterpret_cast(stream_ptr)); + cudaMemcpyAsync(pme_bc, PME_BC0.data(), sizeof(float) * PME_Nfft, cudaMemcpyHostToDevice, + reinterpret_cast(stream_ptr)); PMEReciprocalForce(fftx, ffty, fftz, atom_numbers, beta, pme_bc, pme_uxyz, pme_frxyz, pme_q, pme_fq, pme_atom_near, - pme_kxyz, boxlength, uint_crd, charge, force, reinterpret_cast(stream_ptr)); + pme_kxyz, uint_crd, charge, force, PME_Nin, PME_Nall, PME_Nfft, PME_plan_r2c, PME_plan_c2r, + PME_inverse_box_vector, reinterpret_cast(stream_ptr)); return true; } @@ -88,6 +162,44 @@ class PMEReciprocalForceGpuKernel : public GpuKernel { output_size_list_.push_back(atom_numbers * sizeof(VECTOR)); } + cufftComplex expc(cufftComplex z) { + cufftComplex res; + float t = expf(z.x); + sincosf(z.y, &res.y, &res.x); + res.x *= t; + res.y *= t; + return res; + } + float M_(float u, int n) { + if (n == 2) { + if (u > 2 || u < 0) return 0; + return 1 - abs(u - 1); + } else { + return u / (n - 1) * M_(u, n - 1) + (n - u) / (n - 1) * M_(u - 1, n - 1); + } + } + float getb(int k, int NFFT, int B_order) { + cufftComplex tempc, tempc2, res; + float tempf; + tempc2.x = 0; + tempc2.y = 0; + + tempc.x = 0; + tempc.y = 2 * (B_order - 1) * PI * k / NFFT; + res = expc(tempc); + + for (int kk = 0; kk < (B_order - 1); kk++) { + tempc.x = 0; + tempc.y = 2 * PI * k / NFFT * kk; + tempc = expc(tempc); + tempf = M_(kk + 1, B_order); + tempc2.x += tempf * tempc.x; + tempc2.y += tempf * tempc.y; + } + res = cuCdivf(res, tempc2); + return res.x * res.x + res.y * res.y; + } + private: size_t ele_uint_crd = 1; @@ -101,18 +213,27 @@ class PMEReciprocalForceGpuKernel : public GpuKernel { float beta; int PME_Nall; int PME_Nfft; + int PME_Nin; + float PI = 3.1415926; + std::vector B1; + std::vector B2; + std::vector B3; + std::vector PME_BC0; + cufftHandle PME_plan_r2c; + cufftHandle PME_plan_c2r; struct VECTOR { float x; float y; float z; }; - + _VECTOR PME_inverse_box_vector; struct UNSIGNED_INT_VECTOR { unsigned int uint_x; unsigned int uint_y; unsigned int uint_z; }; + std::vector PME_kxyz_cpu; }; } // namespace kernel } // namespace mindspore diff --git a/mindspore/ops/operations/__init__.py b/mindspore/ops/operations/__init__.py index 335821db1f..d67b10bc52 100644 --- a/mindspore/ops/operations/__init__.py +++ b/mindspore/ops/operations/__init__.py @@ -105,7 +105,8 @@ from .sponge_ops import (BondForce, BondEnergy, BondAtomEnergy, BondForceWithAto LJForce, LJEnergy, LJForceWithPMEDirectForce, PMEExcludedForce, PMEEnergy, Dihedral14LJForce, Dihedral14LJForceWithDirectCF, Dihedral14LJEnergy, Dihedral14LJCFForceWithAtomEnergy, Dihedral14LJAtomEnergy, Dihedral14CFEnergy, Dihedral14CFAtomEnergy, MDIterationLeapFrog, - GetCenterOfGeometry, MDTemperature, NeighborListUpdate) + GetCenterOfGeometry, MDTemperature, NeighborListUpdate, MDIterationLeapFrogLiujian, + CrdToUintCrd, MDIterationSetupRandState) __all__ = [ @@ -465,7 +466,9 @@ __all__ = [ "GetCenterOfGeometry", "MDTemperature", "NeighborListUpdate", - + "MDIterationLeapFrogLiujian", + "CrdToUintCrd", + "MDIterationSetupRandState", ] __all__.sort() diff --git a/mindspore/ops/operations/sponge_ops.py b/mindspore/ops/operations/sponge_ops.py index 7d8160d7f7..380e5f6c7a 100644 --- a/mindspore/ops/operations/sponge_ops.py +++ b/mindspore/ops/operations/sponge_ops.py @@ -1,2650 +1,2810 @@ -# Copyright 2021 Huawei Technologies Co., Ltd -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ============================================================================ - -"""Operators for sponge.""" - -from ..primitive import PrimitiveWithInfer, prim_attr_register -from ..._checkparam import Rel -from ..._checkparam import Validator as validator -from ...common import dtype as mstype - - -class BondForce(PrimitiveWithInfer): - """ - Calculate the force exerted by the simple harmonic bond on the corresponding atoms. - Assume the number of harmonic bonds is M and the number of atoms is N. - - .. math:: - - 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 - - Args: - atom_numbers(int32): the number of atoms N. - bond_numbers(int32): the number of harmonic bonds M. - - Inputs: - - **uint_crd_f** (Tensor, uint32 ) - [N, 3], the unsigned int coordinate value of each atom. - - **scaler_f** (Tensor, float32) - [3,], the 3-D scale factor (x, y, z), - between the real space float coordinates and the unsigned int coordinates. - - **atom_a** (Tensor, int32) - [M,], the first atom index of each bond. - - **atom_b** (Tensor, int32) - [M,], the second atom index of each bond. - - **bond_k** (Tensor, float32) - [M,], the force constant of each bond. - - **bond_r0** (Tensor, float32) - [M,], the equlibrium length of each bond. - - Outputs: - - **frc_f** (float32 Tensor) - [N, 3], the force felt by each atom. - - Supported Platforms: - ``GPU`` - """ - - @prim_attr_register - def __init__(self, bond_numbers, atom_numbers): - validator.check_value_type('bond_numbers', bond_numbers, (int), self.name) - validator.check_value_type('atom_numbers', atom_numbers, (int), self.name) - self.bond_numbers = bond_numbers - self.atom_numbers = atom_numbers - self.add_prim_attr('bond_numbers', self.bond_numbers) - self.add_prim_attr('atom_numbers', self.atom_numbers) - self.init_prim_io_names(inputs=['uint_crd_f', 'scaler_f', 'atom_a', 'atom_b', 'bond_k', 'bond_r0'], - outputs=['frc_f']) - - def infer_shape(self, uint_crd_f_shape, scaler_f_shape, atom_a_shape, atom_b_shape, bond_k_shape, bond_r0_shape): - cls_name = self.name - N = self.atom_numbers - M = self.bond_numbers - validator.check_int(len(uint_crd_f_shape), 2, Rel.EQ, "uint_crd_f_dim", cls_name) - validator.check_int(len(scaler_f_shape), 1, Rel.EQ, "scaler_f_dim", cls_name) - validator.check_int(len(atom_a_shape), 1, Rel.EQ, "atom_a_dim", cls_name) - validator.check_int(len(atom_b_shape), 1, Rel.EQ, "atom_b_dim", cls_name) - validator.check_int(len(bond_k_shape), 1, Rel.EQ, "bond_k_dim", cls_name) - validator.check_int(len(bond_r0_shape), 1, Rel.EQ, "bond_r0_dim", cls_name) - validator.check_int(uint_crd_f_shape[0], N, Rel.EQ, "uint_crd_f_shape[0]", cls_name) - validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f_shape[1]", cls_name) - validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f_shape", cls_name) - validator.check_int(atom_a_shape[0], M, Rel.EQ, "uint_crd_f_shape", cls_name) - validator.check_int(atom_b_shape[0], M, Rel.EQ, "atom_b_shape", cls_name) - validator.check_int(bond_k_shape[0], M, Rel.EQ, "bond_k_shape", cls_name) - validator.check_int(bond_r0_shape[0], M, Rel.EQ, "bond_r0_shape", cls_name) - return uint_crd_f_shape - - def infer_dtype(self, uint_crd_f_dtype, scaler_f_type, atom_a_type, atom_b_type, bond_k_type, bond_r0_type): - validator.check_tensor_dtype_valid('uint_crd_f', uint_crd_f_dtype, [mstype.uint32], self.name) - validator.check_tensor_dtype_valid('scaler_f', scaler_f_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('atom_a', atom_a_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('atom_b', atom_b_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('bond_k', bond_k_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('bond_r0', bond_r0_type, [mstype.float32], self.name) - return bond_r0_type - - -class BondEnergy(PrimitiveWithInfer): - """ - Calculate the harmonic potential energy between each bonded atom pair. - Assume our system has N atoms and M harmonic bonds. - - .. math:: - - dr = (x_1-x_2, y_1-y_2, z_1-z_2) - - .. math:: - - E = k*(|dr| - r_0)^2 - - Args: - atom_numbers(int32): the number of atoms N. - bond_numbers(int32): the number of harmonic bonds M. - - Inputs: - - **uint_crd_f** (Tensor, uint32 ) - [N, 3], the unsigned int coordinate value of each atom. - - **scaler_f** (Tensor, float32) - [3,], the 3-D scale factor (x, y, z), - between the real space float coordinates and the unsigned int coordinates. - - **atom_a** (Tensor, int32) - [M,], the first atom index of each bond. - - **atom_b** (Tensor, int32) - [M,], the second atom index of each bond. - - **bond_k** (Tensor, float32) - [M,], the force constant of each bond. - - **bond_r0** (Tensor, float32) - [M,], the equlibrium length of each bond. - - Outputs: - - **bond_ene** (Tensor, float32) - [M,], the harmonic potential energy for each bond. - - Supported Platforms: - ``GPU`` - """ - - @prim_attr_register - def __init__(self, bond_numbers, atom_numbers): - validator.check_value_type('bond_numbers', bond_numbers, (int), self.name) - validator.check_value_type('atom_numbers', atom_numbers, (int), self.name) - self.bond_numbers = bond_numbers - self.atom_numbers = atom_numbers - self.add_prim_attr('bond_numbers', self.bond_numbers) - self.add_prim_attr('atom_numbers', self.atom_numbers) - self.init_prim_io_names(inputs=['uint_crd_f', 'scaler_f', 'atom_a', 'atom_b', 'bond_k', 'bond_r0'], - outputs=['bond_ene']) - - def infer_shape(self, uint_crd_f_shape, scaler_f_shape, atom_a_shape, atom_b_shape, bond_k_shape, bond_r0_shape): - cls_name = self.name - N = self.atom_numbers - M = self.bond_numbers - validator.check_int(len(uint_crd_f_shape), 2, Rel.EQ, "uint_crd_f_dim", cls_name) - validator.check_int(len(scaler_f_shape), 1, Rel.EQ, "scaler_f_dim", cls_name) - validator.check_int(len(atom_a_shape), 1, Rel.EQ, "atom_a_dim", cls_name) - validator.check_int(len(atom_b_shape), 1, Rel.EQ, "atom_b_dim", cls_name) - validator.check_int(len(bond_k_shape), 1, Rel.EQ, "bond_k_dim", cls_name) - validator.check_int(len(bond_r0_shape), 1, Rel.EQ, "bond_r0_dim", cls_name) - validator.check_int(uint_crd_f_shape[0], N, Rel.EQ, "uint_crd_f_shape[0]", cls_name) - validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f_shape[1]", cls_name) - validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f_shape", cls_name) - validator.check_int(atom_a_shape[0], M, Rel.EQ, "uint_crd_f_shape", cls_name) - validator.check_int(atom_b_shape[0], M, Rel.EQ, "atom_b_shape", cls_name) - validator.check_int(bond_k_shape[0], M, Rel.EQ, "bond_k_shape", cls_name) - validator.check_int(bond_r0_shape[0], M, Rel.EQ, "bond_r0_shape", cls_name) - - return bond_k_shape - - def infer_dtype(self, uint_crd_f_dtype, scaler_f_type, atom_a_type, atom_b_type, bond_k_type, bond_r0_type): - validator.check_tensor_dtype_valid('uint_crd_f', uint_crd_f_dtype, [mstype.uint32], self.name) - validator.check_tensor_dtype_valid('scaler_f', scaler_f_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('atom_a', atom_a_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('atom_b', atom_b_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('bond_k', bond_k_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('bond_r0', bond_r0_type, [mstype.float32], self.name) - return bond_r0_type - - -class BondAtomEnergy(PrimitiveWithInfer): - """ - Add the potential energy caused by simple harmonic bonds to the total - potential energy of each atom. - - The calculation formula is the same as operator BondEnergy(). - - Args: - atom_numbers(int32): the number of atoms N. - bond_numbers(int32): the number of harmonic bonds M. - - Inputs: - - **uint_crd_f** (Tensor, uint32 ) - [N, 3], the unsigned int coordinate value of each atom. - - **scaler_f** (Tensor, float32) - [3,], the 3-D scale factor (x, y, z), - between the real space float coordinates and the unsigned int coordinates. - - **atom_a** (Tensor, int32) - [M,], the first atom index of each bond. - - **atom_b** (Tensor, int32) - [M,], the second atom index of each bond. - - **bond_k** (Tensor, float32) - [M,], the force constant of each bond. - - **bond_r0** (Tensor, float32) - [M,], the equlibrium length of each bond. - - Outputs: - - **atom_ene** (Tensor, float32) - [N,], the accumulated potential energy for each atom. - - Supported Platforms: - ``GPU`` - """ - - @prim_attr_register - def __init__(self, bond_numbers, atom_numbers): - validator.check_value_type('bond_numbers', bond_numbers, (int), self.name) - validator.check_value_type('atom_numbers', atom_numbers, (int), self.name) - self.bond_numbers = bond_numbers - self.atom_numbers = atom_numbers - self.add_prim_attr('bond_numbers', self.bond_numbers) - self.add_prim_attr('atom_numbers', self.atom_numbers) - self.init_prim_io_names(inputs=['uint_crd_f', 'scaler_f', 'atom_a', 'atom_b', 'bond_k', 'bond_r0'], - outputs=['atom_ene']) - - def infer_shape(self, uint_crd_f_shape, scaler_f_shape, atom_a_shape, atom_b_shape, bond_k_shape, bond_r0_shape): - cls_name = self.name - N = self.atom_numbers - M = self.bond_numbers - validator.check_int(len(uint_crd_f_shape), 2, Rel.EQ, "uint_crd_f_dim", cls_name) - validator.check_int(len(scaler_f_shape), 1, Rel.EQ, "scaler_f_dim", cls_name) - validator.check_int(len(atom_a_shape), 1, Rel.EQ, "atom_a_dim", cls_name) - validator.check_int(len(atom_b_shape), 1, Rel.EQ, "atom_b_dim", cls_name) - validator.check_int(len(bond_k_shape), 1, Rel.EQ, "bond_k_dim", cls_name) - validator.check_int(len(bond_r0_shape), 1, Rel.EQ, "bond_r0_dim", cls_name) - validator.check_int(uint_crd_f_shape[0], N, Rel.EQ, "uint_crd_f_shape[0]", cls_name) - validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f_shape[1]", cls_name) - validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f_shape", cls_name) - validator.check_int(atom_a_shape[0], M, Rel.EQ, "uint_crd_f_shape", cls_name) - validator.check_int(atom_b_shape[0], M, Rel.EQ, "atom_b_shape", cls_name) - validator.check_int(bond_k_shape[0], M, Rel.EQ, "bond_k_shape", cls_name) - validator.check_int(bond_r0_shape[0], M, Rel.EQ, "bond_r0_shape", cls_name) - return [N,] - - def infer_dtype(self, uint_crd_f_dtype, scaler_f_type, atom_a_type, atom_b_type, bond_k_type, bond_r0_type): - validator.check_tensor_dtype_valid('uint_crd_f', uint_crd_f_dtype, [mstype.uint32], self.name) - validator.check_tensor_dtype_valid('scaler_f', scaler_f_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('atom_a', atom_a_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('atom_b', atom_b_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('bond_k', bond_k_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('bond_r0', bond_r0_type, [mstype.float32], self.name) - return bond_r0_type - - -class BondForceWithAtomEnergy(PrimitiveWithInfer): - """ - Calculate bond force and harmonic potential energy together. - - The calculation formula is the same as operator BondForce() and BondEnergy(). - - Args: - atom_numbers(int32): the number of atoms N. - bond_numbers(int32): the number of harmonic bonds M. - - Inputs: - - **uint_crd_f** (Tensor, uint32 ) - [N, 3], the unsigned int coordinate value of each atom. - - **scaler_f** (Tensor, float32) - [3,], the 3-D scale factor (x, y, z), - between the real space float coordinates and the unsigned int coordinates. - - **atom_a** (Tensor, int32) - [M,], the first atom index of each bond. - - **atom_b** (Tensor, int32) - [M,], the second atom index of each bond. - - **bond_k** (Tensor, float32) - [M,], the force constant of each bond. - - **bond_r0** (Tensor, float32) - [M,], the equlibrium length of each bond. - - Outputs: - - **frc_f** (Tensor, float32) - [N, 3], same as operator BondForce(). - - **atom_e** (Tensor, float32) - [N,], same as atom_ene in operator BondAtomEnergy(). - - Supported Platforms: - ``GPU`` - """ - - @prim_attr_register - def __init__(self, bond_numbers, atom_numbers): - validator.check_value_type('bond_numbers', bond_numbers, (int), self.name) - validator.check_value_type('atom_numbers', atom_numbers, (int), self.name) - self.bond_numbers = bond_numbers - self.atom_numbers = atom_numbers - self.add_prim_attr('bond_numbers', self.bond_numbers) - self.add_prim_attr('atom_numbers', self.atom_numbers) - self.init_prim_io_names(inputs=['uint_crd_f', 'scaler_f', 'atom_a', 'atom_b', 'bond_k', 'bond_r0'], - outputs=['frc_f', 'atom_e']) - - def infer_shape(self, uint_crd_f_shape, scaler_f_shape, atom_a_shape, atom_b_shape, bond_k_shape, bond_r0_shape): - cls_name = self.name - N = self.atom_numbers - M = self.bond_numbers - validator.check_int(len(uint_crd_f_shape), 2, Rel.EQ, "uint_crd_f_dim", cls_name) - validator.check_int(len(scaler_f_shape), 1, Rel.EQ, "scaler_f_dim", cls_name) - validator.check_int(len(atom_a_shape), 1, Rel.EQ, "atom_a_dim", cls_name) - validator.check_int(len(atom_b_shape), 1, Rel.EQ, "atom_b_dim", cls_name) - validator.check_int(len(bond_k_shape), 1, Rel.EQ, "bond_k_dim", cls_name) - validator.check_int(len(bond_r0_shape), 1, Rel.EQ, "bond_r0_dim", cls_name) - validator.check_int(uint_crd_f_shape[0], N, Rel.EQ, "uint_crd_f_shape[0]", cls_name) - validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f_shape[1]", cls_name) - validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f_shape", cls_name) - validator.check_int(atom_a_shape[0], M, Rel.EQ, "uint_crd_f_shape", cls_name) - validator.check_int(atom_b_shape[0], M, Rel.EQ, "atom_b_shape", cls_name) - validator.check_int(bond_k_shape[0], M, Rel.EQ, "bond_k_shape", cls_name) - validator.check_int(bond_r0_shape[0], M, Rel.EQ, "bond_r0_shape", cls_name) - - return uint_crd_f_shape, [N,] - - def infer_dtype(self, uint_crd_f_dtype, scaler_f_type, atom_a_type, atom_b_type, bond_k_type, bond_r0_type): - validator.check_tensor_dtype_valid('uint_crd_f', uint_crd_f_dtype, [mstype.uint32], self.name) - validator.check_tensor_dtype_valid('scaler_f', scaler_f_type, [mstype.float32], self.name) - - validator.check_tensor_dtype_valid('atom_a', atom_a_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('atom_b', atom_b_type, [mstype.int32], self.name) - - validator.check_tensor_dtype_valid('bond_k', bond_k_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('bond_r0', bond_r0_type, [mstype.float32], self.name) - return bond_r0_type, bond_r0_type - - -class BondForceWithAtomVirial(PrimitiveWithInfer): - """ - Calculate bond force and the virial coefficient caused by simple harmonic - bond for each atom together. - - The calculation formula of the force part is the same as operator BondForce(). - The Virial part is as follows: - - .. math:: - - dr = (x_1-x_2, y_1-y_2, z_1-z_2) - - .. math:: - - virial = |dr|*(|dr| - r_0)*k - - Args: - atom_numbers(int32): the number of atoms N. - bond_numbers(int32): the number of harmonic bonds M. - - Inputs: - - **uint_crd_f** (Tensor, uint32 ) - [N, 3], the unsigned int coordinate value of each atom. - - **scaler_f** (Tensor, float32) - [3,], the 3-D scale factor (x, y, z), - between the real space float coordinates and the unsigned int coordinates. - - **atom_a** (Tensor, int32) - [M,], the first atom index of each bond. - - **atom_b** (Tensor, int32) - [M,], the second atom index of each bond. - - **bond_k** (Tensor, float32) - [M,], the force constant of each bond. - - **bond_r0** (Tensor, float32) - [M,], the equlibrium length of each bond. - - Outputs: - - **frc_f** (Tensor, float32) - [N, 3], same as operator BondForce(). - - **atom_v** (Tensor, float32) - [N,], the accumulated virial coefficient for each atom. - - Supported Platforms: - ``GPU`` - """ - - @prim_attr_register - def __init__(self, bond_numbers, atom_numbers): - validator.check_value_type('bond_numbers', bond_numbers, (int), self.name) - validator.check_value_type('atom_numbers', atom_numbers, (int), self.name) - self.bond_numbers = bond_numbers - self.atom_numbers = atom_numbers - self.add_prim_attr('bond_numbers', self.bond_numbers) - self.add_prim_attr('atom_numbers', self.atom_numbers) - self.init_prim_io_names(inputs=['uint_crd_f', 'scaler_f', 'atom_a', 'atom_b', 'bond_k', 'bond_r0'], - outputs=['frc_f', 'atom_v']) - - def infer_shape(self, uint_crd_f_shape, scaler_f_shape, atom_a_shape, atom_b_shape, bond_k_shape, bond_r0_shape): - cls_name = self.name - N = self.atom_numbers - M = self.bond_numbers - validator.check_int(len(uint_crd_f_shape), 2, Rel.EQ, "uint_crd_f_dim", cls_name) - validator.check_int(len(scaler_f_shape), 1, Rel.EQ, "scaler_f_dim", cls_name) - validator.check_int(len(atom_a_shape), 1, Rel.EQ, "atom_a_dim", cls_name) - validator.check_int(len(atom_b_shape), 1, Rel.EQ, "atom_b_dim", cls_name) - validator.check_int(len(bond_k_shape), 1, Rel.EQ, "bond_k_dim", cls_name) - validator.check_int(len(bond_r0_shape), 1, Rel.EQ, "bond_r0_dim", cls_name) - validator.check_int(uint_crd_f_shape[0], N, Rel.EQ, "uint_crd_f_shape[0]", cls_name) - validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f_shape[1]", cls_name) - validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f_shape", cls_name) - validator.check_int(atom_a_shape[0], M, Rel.EQ, "uint_crd_f_shape", cls_name) - validator.check_int(atom_b_shape[0], M, Rel.EQ, "atom_b_shape", cls_name) - validator.check_int(bond_k_shape[0], M, Rel.EQ, "bond_k_shape", cls_name) - validator.check_int(bond_r0_shape[0], M, Rel.EQ, "bond_r0_shape", cls_name) - - return uint_crd_f_shape, [N,] - - def infer_dtype(self, uint_crd_f_dtype, scaler_f_type, atom_a_type, atom_b_type, bond_k_type, bond_r0_type): - validator.check_tensor_dtype_valid('uint_crd_f', uint_crd_f_dtype, [mstype.uint32], self.name) - validator.check_tensor_dtype_valid('scaler_f', scaler_f_type, [mstype.float32], self.name) - - validator.check_tensor_dtype_valid('atom_a', atom_a_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('atom_b', atom_b_type, [mstype.int32], self.name) - - validator.check_tensor_dtype_valid('bond_k', bond_k_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('bond_r0', bond_r0_type, [mstype.float32], self.name) - return bond_r0_type, bond_r0_type - - -class DihedralForce(PrimitiveWithInfer): - """ - Calculate the force exerted by the dihedral term which made of 4-atoms - on the corresponding atoms. Assume the number of dihedral terms is M and - the number of atoms is N. - - .. math:: - - 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) - .. math:: - dr_{cd} = (x_d-x_c, y_d-y_c, z_d-z_c) - .. math:: - r1 = dr_{ab}*dr_{cb} - .. math:: - r2 = dr_{cd}*dr_{cb} - .. math:: - phi = pi - sign(inner_product(r1*r2), dr_{cb}) - * 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) - .. math:: - dphidr1 = r2/|r1|/|r2| + cos(phi)/|r1|^2*r1 - .. math:: - dphidr2 = r1/|r1|/|r2| + cos(phi)/|r2|^2*r2 - .. math:: - dEdra = dEdphi * dr_{cb} * dphidr1 - .. math:: - dEdrd = dEdphi * dphi_dr2 * dr_{cb} - .. math:: - dEdrjpart = dEdphi * ((dr_{ab} * dphidr1) + (dr_{cd} * dphidr2)) - .. math:: - F_a = dEdri - .. math:: - F_b = dEdrjpart - dEdri - .. math:: - F_c = - dEdrl - dEdrjpart - .. math:: - F_d = dEdrl - - Args: - dihedral_numbers(int32): the number of dihedral terms M. - - Inputs: - - **dihedral_numbers** (int32) - the number of dihedral terms M. - - **uint_crd_f** (Tensor, uint32) - [N, 3], the unsigned int coordinates - value of each atom. - - **scaler_f** (Tensor, float32) - [3,], the 3-D scale factor between - the real space float coordinates and the unsigned int coordinates. - - **atom_a** (Tensor, int32) - [M,], the 1st atom index of each dihedral. - - **atom_b** (Tensor, int32) - [M,], the 2nd atom index of each dihedral. - - **atom_c** (Tensor, int32) - [M,], the 3rd atom index of each dihedral. - - **atom_d** (Tensor, int32) - [M,], the 4th atom index of each dihedral. - 4 atoms are connected in the form a-b-c-d. - - **ipn** (Tensor, int32) - [M,], the period of dihedral angle of each dihedral. - - **pk** (Tensor, float32) - [M,], the force constant of each dihedral. - - **gamc** (Tensor, float32) - [M,], k*cos(phi_0) of each dihedral. - - **gams** (Tensor, float32) - [M,], k*sin(phi_0) of each dihedral. - - **pn** (Tensor, float32) - [M,], the floating point form of ipn. - - Outputs: - - **frc_f** (Tensor, float32) - [N, 3], the force felt by each atom. - - Supported Platforms: - ``GPU`` - """ - - @prim_attr_register - def __init__(self, dihedral_numbers): - validator.check_value_type('dihedral_numbers', dihedral_numbers, (int), self.name) - 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', - 'gamc', 'gams', 'pn'], - outputs=['frc_f']) - self.add_prim_attr('dihedral_numbers', self.dihedral_numbers) - - 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): - cls_name = self.name - M = self.dihedral_numbers - validator.check_int(len(uint_crd_f_shape), 2, Rel.EQ, "uint_crd_f_dim", cls_name) - validator.check_int(len(scaler_f_shape), 1, Rel.EQ, "scaler_f_dim", cls_name) - validator.check_int(len(atom_a_shape), 1, Rel.EQ, "atom_a_dim", cls_name) - validator.check_int(len(atom_b_shape), 1, Rel.EQ, "atom_b_dim", cls_name) - validator.check_int(len(atom_c_shape), 1, Rel.EQ, "atom_c_dim", cls_name) - validator.check_int(len(atom_d_shape), 1, Rel.EQ, "atom_d_dim", cls_name) - validator.check_int(len(ipn_shape), 1, Rel.EQ, "ipn_dim", cls_name) - validator.check_int(len(pk_shape), 1, Rel.EQ, "pk_dim", cls_name) - validator.check_int(len(gamc_shape), 1, Rel.EQ, "gamc_dim", cls_name) - validator.check_int(len(gams_shape), 1, Rel.EQ, "gams_dim", cls_name) - validator.check_int(len(pn_shape), 1, Rel.EQ, "pn_dim", cls_name) - - validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f_shape[1]", cls_name) - validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f_shape", cls_name) - validator.check_int(atom_a_shape[0], M, Rel.EQ, "atom_a_shape", cls_name) - validator.check_int(atom_b_shape[0], M, Rel.EQ, "atom_b_shape", cls_name) - validator.check_int(atom_c_shape[0], M, Rel.EQ, "atom_c_shape", cls_name) - validator.check_int(atom_d_shape[0], M, Rel.EQ, "atom_d_shape", cls_name) - validator.check_int(ipn_shape[0], M, Rel.EQ, "ipn_shape", cls_name) - validator.check_int(pk_shape[0], M, Rel.EQ, "pk_shape", cls_name) - validator.check_int(gamc_shape[0], M, Rel.EQ, "gamc_shape", cls_name) - validator.check_int(gams_shape[0], M, Rel.EQ, "gams_shape", cls_name) - validator.check_int(pn_shape[0], M, Rel.EQ, "pn_shape", cls_name) - return uint_crd_f_shape - - def infer_dtype(self, uint_crd_f_dtype, scaler_f_type, atom_a_type, atom_b_type, atom_c_type, atom_d_type, - ipn_type, pk_type, gamc_type, gams_type, pn_type): - validator.check_tensor_dtype_valid('uint_crd_f', uint_crd_f_dtype, [mstype.uint32], self.name) - validator.check_tensor_dtype_valid('scaler_f', scaler_f_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('atom_a', atom_a_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('atom_b', atom_b_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('atom_c', atom_c_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('atom_d', atom_d_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('ipn', ipn_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('pk', pk_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('gamc', gamc_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('gams', gams_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('pn', pn_type, [mstype.float32], self.name) - - return pn_type - - -class DihedralEnergy(PrimitiveWithInfer): - """ - Calculate the potential energy caused by dihedral terms for each 4-atom pair. - Assume our system has N atoms and M dihedral terms. - - .. math:: - - E = k(1 + cos(n*phi - phi_0)) - - Args: - dihedral_numbers(int32): the number of dihedral terms M. - - Inputs: - - **dihedral_numbers** (int32) - the number of dihedral terms M. - - **uint_crd_f** (Tensor, uint32) - [N, 3], the unsigned int coordinates - value of each atom. - - **scaler_f** (Tensor, float32) - [3,], the 3-D scale factor between - the real space float coordinates and the unsigned int coordinates. - - **atom_a** (Tensor, int32) - [M,], the 1st atom index of each dihedral. - - **atom_b** (Tensor, int32) - [M,], the 2nd atom index of each dihedral. - - **atom_c** (Tensor, int32) - [M,], the 3rd atom index of each dihedral. - - **atom_d** (Tensor, int32) - [M,], the 4th atom index of each dihedral. - 4 atoms are connected in the form a-b-c-d. - - **ipn** (Tensor, int32) - [M,], the period of dihedral angle of each dihedral. - - **pk** (Tensor, float32) - [M,], the force constant of each dihedral. - - **gamc** (Tensor, float32) - [M,], k*cos(phi_0) of each dihedral. - - **gams** (Tensor, float32) - [M,], k*sin(phi_0) of each dihedral. - - **pn** (Tensor, float32) - [M,], the floating point form of ipn. - - Outputs: - - **ene** (Tensor, float32) - [M,], the potential energy for each - dihedral term. - - Supported Platforms: - ``GPU`` - """ - - @prim_attr_register - def __init__(self, dihedral_numbers): - validator.check_value_type('dihedral_numbers', dihedral_numbers, (int), self.name) - 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', - 'gamc', 'gams', 'pn'], - outputs=['ene']) - self.add_prim_attr('dihedral_numbers', self.dihedral_numbers) - - 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): - cls_name = self.name - M = self.dihedral_numbers - validator.check_int(len(uint_crd_f_shape), 2, Rel.EQ, "uint_crd_f_dim", cls_name) - validator.check_int(len(scaler_f_shape), 1, Rel.EQ, "scaler_f_dim", cls_name) - validator.check_int(len(atom_a_shape), 1, Rel.EQ, "atom_a_dim", cls_name) - validator.check_int(len(atom_b_shape), 1, Rel.EQ, "atom_b_dim", cls_name) - validator.check_int(len(atom_c_shape), 1, Rel.EQ, "atom_c_dim", cls_name) - validator.check_int(len(atom_d_shape), 1, Rel.EQ, "atom_d_dim", cls_name) - validator.check_int(len(ipn_shape), 1, Rel.EQ, "ipn_dim", cls_name) - validator.check_int(len(pk_shape), 1, Rel.EQ, "pk_dim", cls_name) - validator.check_int(len(gamc_shape), 1, Rel.EQ, "gamc_dim", cls_name) - validator.check_int(len(gams_shape), 1, Rel.EQ, "gams_dim", cls_name) - validator.check_int(len(pn_shape), 1, Rel.EQ, "pn_dim", cls_name) - - validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f_shape[1]", cls_name) - validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f_shape", cls_name) - validator.check_int(atom_a_shape[0], M, Rel.EQ, "atom_a_shape", cls_name) - validator.check_int(atom_b_shape[0], M, Rel.EQ, "atom_b_shape", cls_name) - validator.check_int(atom_c_shape[0], M, Rel.EQ, "atom_c_shape", cls_name) - validator.check_int(atom_d_shape[0], M, Rel.EQ, "atom_d_shape", cls_name) - validator.check_int(ipn_shape[0], M, Rel.EQ, "ipn_shape", cls_name) - validator.check_int(pk_shape[0], M, Rel.EQ, "pk_shape", cls_name) - validator.check_int(gamc_shape[0], M, Rel.EQ, "gamc_shape", cls_name) - validator.check_int(gams_shape[0], M, Rel.EQ, "gams_shape", cls_name) - validator.check_int(pn_shape[0], M, Rel.EQ, "pn_shape", cls_name) - return [M,] - - def infer_dtype(self, uint_crd_f_dtype, scaler_f_type, atom_a_type, atom_b_type, atom_c_type, atom_d_type, - ipn_type, pk_type, gamc_type, gams_type, pn_type): - validator.check_tensor_dtype_valid('uint_crd_f', uint_crd_f_dtype, [mstype.uint32], self.name) - validator.check_tensor_dtype_valid('scaler_f', scaler_f_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('atom_a', atom_a_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('atom_b', atom_b_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('atom_c', atom_c_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('atom_d', atom_d_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('ipn', ipn_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('pk', pk_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('gamc', gamc_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('gams', gams_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('pn', pn_type, [mstype.float32], self.name) - - return pn_type - - -class DihedralAtomEnergy(PrimitiveWithInfer): - """ - Add the potential energy caused by dihedral terms to the total potential - energy of each atom. - - The calculation formula is the same as operator DihedralEnergy(). - - Args: - dihedral_numbers(int32): the number of dihedral terms M. - - Inputs: - - **dihedral_numbers** (int32) - the number of dihedral terms M. - - **uint_crd_f** (Tensor, uint32) - [N, 3], the unsigned int coordinates - value of each atom. - - **scaler_f** (Tensor, float32) - [3,], the 3-D scale factor between - the real space float coordinates and the unsigned int coordinates. - - **atom_a** (Tensor, int32) - [M,], the 1st atom index of each dihedral. - - **atom_b** (Tensor, int32) - [M,], the 2nd atom index of each dihedral. - - **atom_c** (Tensor, int32) - [M,], the 3rd atom index of each dihedral. - - **atom_d** (Tensor, int32) - [M,], the 4th atom index of each dihedral. - 4 atoms are connected in the form a-b-c-d. - - **ipn** (Tensor, int32) - [M,], the period of dihedral angle of each dihedral. - - **pk** (Tensor, float32) - [M,], the force constant of each dihedral. - - **gamc** (Tensor, float32) - [M,], k*cos(phi_0) of each dihedral. - - **gams** (Tensor, float32) - [M,], k*sin(phi_0) of each dihedral. - - **pn** (Tensor, float32) - [M,], the floating point form of ipn. - - Outputs: - - **ene** (Tensor, float32) - [N,], the accumulated potential - energy for each atom. - - Supported Platforms: - ``GPU`` - """ - - @prim_attr_register - def __init__(self, dihedral_numbers): - validator.check_value_type('dihedral_numbers', dihedral_numbers, (int), self.name) - 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', - 'gamc', 'gams', 'pn'], - outputs=['ene']) - self.add_prim_attr('dihedral_numbers', self.dihedral_numbers) - - 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): - cls_name = self.name - N = uint_crd_f_shape[0] - M = self.dihedral_numbers - validator.check_int(len(uint_crd_f_shape), 2, Rel.EQ, "uint_crd_f_dim", cls_name) - validator.check_int(len(scaler_f_shape), 1, Rel.EQ, "scaler_f_dim", cls_name) - validator.check_int(len(atom_a_shape), 1, Rel.EQ, "atom_a_dim", cls_name) - validator.check_int(len(atom_b_shape), 1, Rel.EQ, "atom_b_dim", cls_name) - validator.check_int(len(atom_c_shape), 1, Rel.EQ, "atom_c_dim", cls_name) - validator.check_int(len(atom_d_shape), 1, Rel.EQ, "atom_d_dim", cls_name) - validator.check_int(len(ipn_shape), 1, Rel.EQ, "ipn_dim", cls_name) - validator.check_int(len(pk_shape), 1, Rel.EQ, "pk_dim", cls_name) - validator.check_int(len(gamc_shape), 1, Rel.EQ, "gamc_dim", cls_name) - validator.check_int(len(gams_shape), 1, Rel.EQ, "gams_dim", cls_name) - validator.check_int(len(pn_shape), 1, Rel.EQ, "pn_dim", cls_name) - - validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f_shape[1]", cls_name) - validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f_shape", cls_name) - validator.check_int(atom_a_shape[0], M, Rel.EQ, "atom_a_shape", cls_name) - validator.check_int(atom_b_shape[0], M, Rel.EQ, "atom_b_shape", cls_name) - validator.check_int(atom_c_shape[0], M, Rel.EQ, "atom_c_shape", cls_name) - validator.check_int(atom_d_shape[0], M, Rel.EQ, "atom_d_shape", cls_name) - validator.check_int(ipn_shape[0], M, Rel.EQ, "ipn_shape", cls_name) - validator.check_int(pk_shape[0], M, Rel.EQ, "pk_shape", cls_name) - validator.check_int(gamc_shape[0], M, Rel.EQ, "gamc_shape", cls_name) - validator.check_int(gams_shape[0], M, Rel.EQ, "gams_shape", cls_name) - validator.check_int(pn_shape[0], M, Rel.EQ, "pn_shape", cls_name) - return [N,] - - def infer_dtype(self, uint_crd_f_dtype, scaler_f_type, atom_a_type, atom_b_type, atom_c_type, atom_d_type, - ipn_type, pk_type, gamc_type, gams_type, pn_type): - validator.check_tensor_dtype_valid('uint_crd_f', uint_crd_f_dtype, [mstype.uint32], self.name) - validator.check_tensor_dtype_valid('scaler_f', scaler_f_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('atom_a', atom_a_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('atom_b', atom_b_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('atom_c', atom_c_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('atom_d', atom_d_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('ipn', ipn_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('pk', pk_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('gamc', gamc_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('gams', gams_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('pn', pn_type, [mstype.float32], self.name) - - return pn_type - - -class DihedralForceWithAtomEnergy(PrimitiveWithInfer): - """ - Calculate dihedral force and potential energy together. - - The calculation formula is the same as operator DihedralForce() and DihedralEnergy(). - - Args: - dihedral_numbers(int32): the number of dihedral terms M. - - Inputs: - - **dihedral_numbers** (int32) - the number of dihedral terms M. - - **uint_crd_f** (Tensor, uint32) - [N, 3], the unsigned int coordinates - value of each atom. - - **scaler_f** (Tensor, float32) - [3,], the 3-D scale factor between - the real space float coordinates and the unsigned int coordinates. - - **atom_a** (Tensor, int32) - [M,], the 1st atom index of each dihedral. - - **atom_b** (Tensor, int32) - [M,], the 2nd atom index of each dihedral. - - **atom_c** (Tensor, int32) - [M,], the 3rd atom index of each dihedral. - - **atom_d** (Tensor, int32) - [M,], the 4th atom index of each dihedral. - 4 atoms are connected in the form a-b-c-d. - - **ipn** (Tensor, int32) - [M,], the period of dihedral angle of each dihedral. - - **pk** (Tensor, float32) - [M,], the force constant of each dihedral. - - **gamc** (Tensor, float32) - [M,], k*cos(phi_0) of each dihedral. - - **gams** (Tensor, float32) - [M,], k*sin(phi_0) of each dihedral. - - **pn** (Tensor, float32) - [M,], the floating point form of ipn. - - Outputs: - - **frc_f** (Tensor, float32) - [N, 3], same as operator DihedralForce(). - - **ene** (Tensor, float32) - [N,], same as operator DihedralAtomEnergy(). - - Supported Platforms: - ``GPU`` - """ - - @prim_attr_register - def __init__(self, dihedral_numbers): - validator.check_value_type('dihedral_numbers', dihedral_numbers, (int), self.name) - 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', - 'gamc', 'gams', 'pn'], - outputs=['frc_f', 'ene']) - self.add_prim_attr('dihedral_numbers', self.dihedral_numbers) - - 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): - cls_name = self.name - N = uint_crd_f_shape[0] - M = self.dihedral_numbers - validator.check_int(len(uint_crd_f_shape), 2, Rel.EQ, "uint_crd_f_dim", cls_name) - validator.check_int(len(scaler_f_shape), 1, Rel.EQ, "scaler_f_dim", cls_name) - validator.check_int(len(atom_a_shape), 1, Rel.EQ, "atom_a_dim", cls_name) - validator.check_int(len(atom_b_shape), 1, Rel.EQ, "atom_b_dim", cls_name) - validator.check_int(len(atom_c_shape), 1, Rel.EQ, "atom_c_dim", cls_name) - validator.check_int(len(atom_d_shape), 1, Rel.EQ, "atom_d_dim", cls_name) - validator.check_int(len(ipn_shape), 1, Rel.EQ, "ipn_dim", cls_name) - validator.check_int(len(pk_shape), 1, Rel.EQ, "pk_dim", cls_name) - validator.check_int(len(gamc_shape), 1, Rel.EQ, "gamc_dim", cls_name) - validator.check_int(len(gams_shape), 1, Rel.EQ, "gams_dim", cls_name) - validator.check_int(len(pn_shape), 1, Rel.EQ, "pn_dim", cls_name) - - validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f_shape[1]", cls_name) - validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f_shape", cls_name) - validator.check_int(atom_a_shape[0], M, Rel.EQ, "atom_a_shape", cls_name) - validator.check_int(atom_b_shape[0], M, Rel.EQ, "atom_b_shape", cls_name) - validator.check_int(atom_c_shape[0], M, Rel.EQ, "atom_c_shape", cls_name) - validator.check_int(atom_d_shape[0], M, Rel.EQ, "atom_d_shape", cls_name) - validator.check_int(ipn_shape[0], M, Rel.EQ, "ipn_shape", cls_name) - validator.check_int(pk_shape[0], M, Rel.EQ, "pk_shape", cls_name) - validator.check_int(gamc_shape[0], M, Rel.EQ, "gamc_shape", cls_name) - validator.check_int(gams_shape[0], M, Rel.EQ, "gams_shape", cls_name) - validator.check_int(pn_shape[0], M, Rel.EQ, "pn_shape", cls_name) - return uint_crd_f_shape, [N,] - - def infer_dtype(self, uint_crd_f_dtype, scaler_f_type, atom_a_type, atom_b_type, atom_c_type, atom_d_type, - ipn_type, pk_type, gamc_type, gams_type, pn_type): - validator.check_tensor_dtype_valid('uint_crd_f', uint_crd_f_dtype, [mstype.uint32], self.name) - validator.check_tensor_dtype_valid('scaler_f', scaler_f_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('atom_a', atom_a_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('atom_b', atom_b_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('atom_c', atom_c_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('atom_d', atom_d_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('ipn', ipn_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('pk', pk_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('gamc', gamc_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('gams', gams_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('pn', pn_type, [mstype.float32], self.name) - - return pn_type, pn_type - - -class AngleForce(PrimitiveWithInfer): - """ - Calculate the force exerted by angles made of 3 atoms on the - corresponding atoms. Assume the number of angles is M and the - number of atoms is N. - - .. math:: - 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) - .. math:: - 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} - - 1/|dr_{ab}|/|dr_{cb}|*dr_{cb}] - .. math:: - F_c = -2*k*(theta-theta_0)/sin(theta)*[cos(theta)/|dr_{cb}|^2*dr_{cb} - - 1/|dr_{cb}|/|dr_{ab}|*dr_{ab}] - .. math:: - F_b = -F_a - F_c - - Args: - angle_numbers(int32): the number of angles M. - - Inputs: - - **uint_crd_f** (Tensor, uint32) - [N, 3], the unsigned int coordinate value of each atom. - - **scaler_f** (Tensor, float32) - [3,], the 3-D scale factor between - the real space float coordinates and the unsigned int coordinates. - - **atom_a** (Tensor, int32) - [M,], the 1st atom index of each angle. - - **atom_b** (Tensor, int32) - [M,], the 2nd and the central atom index of each angle. - - **atom_c** (Tensor, int32) - [M,], the 3rd atom index of each angle. - - **angle_k** (Tensor, float32) - [M,], the force constant for each angle. - - **angle_theta0** (Tensor, float32) - [M,], the equilibrium position value for each angle. - - Outputs: - - **frc_f** (Tensor, float32) - [N, 3], the force felt by each atom. - - Supported Platforms: - ``GPU`` - """ - - @prim_attr_register - def __init__(self, angle_numbers): - validator.check_value_type('angle_numbers', angle_numbers, (int), self.name) - self.angle_numbers = angle_numbers - self.init_prim_io_names(inputs=['uint_crd_f', 'scaler_f', 'atom_a', 'atom_b', 'atom_c', 'angle_k', - 'angle_theta0'], - outputs=['frc_f']) - self.add_prim_attr('angle_numbers', self.angle_numbers) - - 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): - cls_name = self.name - M = self.angle_numbers - validator.check_int(len(uint_crd_f_shape), 2, Rel.EQ, "uint_crd_f_dim", cls_name) - validator.check_int(len(scaler_f_shape), 1, Rel.EQ, "scaler_f_dim", cls_name) - validator.check_int(len(atom_a_shape), 1, Rel.EQ, "atom_a_dim", cls_name) - validator.check_int(len(atom_b_shape), 1, Rel.EQ, "atom_b_dim", cls_name) - validator.check_int(len(atom_c_shape), 1, Rel.EQ, "atom_c_dim", cls_name) - validator.check_int(len(angle_k_shape), 1, Rel.EQ, "angle_k_dim", cls_name) - validator.check_int(len(angle_theta0_shape), 1, Rel.EQ, "angle_theta0_dim", cls_name) - - validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f_shape[1]", cls_name) - validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f_shape", cls_name) - validator.check_int(atom_a_shape[0], M, Rel.EQ, "atom_a_shape", cls_name) - validator.check_int(atom_b_shape[0], M, Rel.EQ, "atom_b_shape", cls_name) - validator.check_int(atom_c_shape[0], M, Rel.EQ, "atom_c_shape", cls_name) - validator.check_int(angle_k_shape[0], M, Rel.EQ, "angle_k_shape", cls_name) - validator.check_int(angle_theta0_shape[0], M, Rel.EQ, "angle_theta0_shape", cls_name) - return uint_crd_f_shape - - def infer_dtype(self, uint_crd_f_dtype, scaler_f_type, atom_a_type, atom_b_type, atom_c_type, angle_k_type, - angle_theta0_type): - validator.check_tensor_dtype_valid('uint_crd_f', uint_crd_f_dtype, [mstype.uint32], self.name) - validator.check_tensor_dtype_valid('scaler_f', scaler_f_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('atom_a', atom_a_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('atom_b', atom_b_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('atom_c', atom_c_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('angle_k', angle_k_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('angle_theta0', angle_theta0_type, [mstype.float32], self.name) - return angle_k_type - - -class AngleEnergy(PrimitiveWithInfer): - """ - Calculate the energy caused by 3-atoms angle term. Assume the number of angles is M and the - number of atoms is N. - - .. math:: - 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) - .. math:: - theta = arccos(inner_product(dr_{ab}, dr_{cb})/|dr_{ab}|/|dr_{cb}|) - .. math:: - E = k*(theta - theta_0)^2 - - Args: - angle_numbers(int32): the number of angles M. - - Inputs: - - **uint_crd_f** (Tensor, uint32) - [N, 3], the unsigned int coordinate value of each atom. - - **scaler_f** (Tensor, float32) - [3,], the 3-D scale factor between - the real space float coordinates and the unsigned int coordinates. - - **atom_a** (Tensor, int32) - [M,], the 1st atom index of each angle. - - **atom_b** (Tensor, int32) - [M,], the 2nd and the central atom index of each angle. - - **atom_c** (Tensor, int32) - [M,], the 3rd atom index of each angle. - - **angle_k** (Tensor, float32) - [M,], the force constant for each angle. - - **angle_theta0** (Tensor, float32) - [M,], the equilibrium position value for each angle. - - Outputs: - - **ene** (Tensor, float32) - [M,], the potential energy for each angle term. - - Supported Platforms: - ``GPU`` - """ - - @prim_attr_register - def __init__(self, angle_numbers): - validator.check_value_type('angle_numbers', angle_numbers, (int), self.name) - self.angle_numbers = angle_numbers - self.init_prim_io_names(inputs=['uint_crd_f', 'scaler_f', 'atom_a', 'atom_b', 'atom_c', 'angle_k', - 'angle_theta0'], - outputs=['ene']) - self.add_prim_attr('angle_numbers', self.angle_numbers) - - 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): - cls_name = self.name - M = self.angle_numbers - validator.check_int(len(uint_crd_f_shape), 2, Rel.EQ, "uint_crd_f_dim", cls_name) - validator.check_int(len(scaler_f_shape), 1, Rel.EQ, "scaler_f_dim", cls_name) - validator.check_int(len(atom_a_shape), 1, Rel.EQ, "atom_a_dim", cls_name) - validator.check_int(len(atom_b_shape), 1, Rel.EQ, "atom_b_dim", cls_name) - validator.check_int(len(atom_c_shape), 1, Rel.EQ, "atom_c_dim", cls_name) - validator.check_int(len(angle_k_shape), 1, Rel.EQ, "angle_k_dim", cls_name) - validator.check_int(len(angle_theta0_shape), 1, Rel.EQ, "angle_theta0_dim", cls_name) - - validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f_shape[1]", cls_name) - validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f_shape", cls_name) - validator.check_int(atom_a_shape[0], M, Rel.EQ, "atom_a_shape", cls_name) - validator.check_int(atom_b_shape[0], M, Rel.EQ, "atom_b_shape", cls_name) - validator.check_int(atom_c_shape[0], M, Rel.EQ, "atom_c_shape", cls_name) - validator.check_int(angle_k_shape[0], M, Rel.EQ, "angle_k_shape", cls_name) - validator.check_int(angle_theta0_shape[0], M, Rel.EQ, "angle_theta0_shape", cls_name) - return [M,] - - def infer_dtype(self, uint_crd_f_dtype, scaler_f_type, atom_a_type, atom_b_type, atom_c_type, angle_k_type, - angle_theta0_type): - validator.check_tensor_dtype_valid('uint_crd_f', uint_crd_f_dtype, [mstype.uint32], self.name) - validator.check_tensor_dtype_valid('scaler_f', scaler_f_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('atom_a', atom_a_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('atom_b', atom_b_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('atom_c', atom_c_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('angle_k', angle_k_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('angle_theta0', angle_theta0_type, [mstype.float32], self.name) - return angle_k_type - - -class AngleAtomEnergy(PrimitiveWithInfer): - """ - Add the potential energy caused by angle terms to the total potential - 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(). - - Args: - angle_numbers(int32): the number of angles M. - - Inputs: - - **uint_crd_f** (Tensor, uint32) - [N, 3], the unsigned int coordinate value of each atom. - - **scaler_f** (Tensor, float32) - [3,], the 3-D scale factor between - the real space float coordinates and the unsigned int coordinates. - - **atom_a** (Tensor, int32) - [M,], the 1st atom index of each angle. - - **atom_b** (Tensor, int32) - [M,], the 2nd and the central atom index of each angle. - - **atom_c** (Tensor, int32) - [M,], the 3rd atom index of each angle. - - **angle_k** (Tensor, float32) - [M,], the force constant for each angle. - - **angle_theta0** (Tensor, float32) - [M,], the equilibrium position value for each angle. - - Outputs: - - **ene** (Tensor, float32) - [N,], the accumulated potential energy for each atom. - - Supported Platforms: - ``GPU`` - """ - - @prim_attr_register - def __init__(self, angle_numbers): - validator.check_value_type('angle_numbers', angle_numbers, (int), self.name) - self.angle_numbers = angle_numbers - self.init_prim_io_names(inputs=['uint_crd_f', 'scaler_f', 'atom_a', 'atom_b', 'atom_c', 'angle_k', - 'angle_theta0'], - outputs=['ene']) - self.add_prim_attr('angle_numbers', self.angle_numbers) - - 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): - cls_name = self.name - N = uint_crd_f_shape[0] - M = self.angle_numbers - validator.check_int(len(uint_crd_f_shape), 2, Rel.EQ, "uint_crd_f_dim", cls_name) - validator.check_int(len(scaler_f_shape), 1, Rel.EQ, "scaler_f_dim", cls_name) - validator.check_int(len(atom_a_shape), 1, Rel.EQ, "atom_a_dim", cls_name) - validator.check_int(len(atom_b_shape), 1, Rel.EQ, "atom_b_dim", cls_name) - validator.check_int(len(atom_c_shape), 1, Rel.EQ, "atom_c_dim", cls_name) - validator.check_int(len(angle_k_shape), 1, Rel.EQ, "angle_k_dim", cls_name) - validator.check_int(len(angle_theta0_shape), 1, Rel.EQ, "angle_theta0_dim", cls_name) - - validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f_shape[1]", cls_name) - validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f_shape", cls_name) - validator.check_int(atom_a_shape[0], M, Rel.EQ, "atom_a_shape", cls_name) - validator.check_int(atom_b_shape[0], M, Rel.EQ, "atom_b_shape", cls_name) - validator.check_int(atom_c_shape[0], M, Rel.EQ, "atom_c_shape", cls_name) - validator.check_int(angle_k_shape[0], M, Rel.EQ, "angle_k_shape", cls_name) - validator.check_int(angle_theta0_shape[0], M, Rel.EQ, "angle_theta0_shape", cls_name) - return [N,] - - def infer_dtype(self, uint_crd_f_dtype, scaler_f_type, atom_a_type, atom_b_type, atom_c_type, angle_k_type, - angle_theta0_type): - validator.check_tensor_dtype_valid('uint_crd_f', uint_crd_f_dtype, [mstype.uint32], self.name) - validator.check_tensor_dtype_valid('scaler_f', scaler_f_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('atom_a', atom_a_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('atom_b', atom_b_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('atom_c', atom_c_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('angle_k', angle_k_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('angle_theta0', angle_theta0_type, [mstype.float32], self.name) - return angle_k_type - - -class AngleForceWithAtomEnergy(PrimitiveWithInfer): - """ - 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(). - - Args: - angle_numbers(int32): the number of angles M. - - Inputs: - - **uint_crd_f** (Tensor, uint32) - [N, 3], the unsigned int coordinate value of each atom. - - **scaler_f** (Tensor, float32) - [3,], the 3-D scale factor between - the real space float coordinates and the unsigned int coordinates. - - **atom_a** (Tensor, int32) - [M,], the 1st atom index of each angle. - - **atom_b** (Tensor, int32) - [M,], the 2nd and the central atom index of each angle. - - **atom_c** (Tensor, int32) - [M,], the 3rd atom index of each angle. - - **angle_k** (Tensor, float32) - [M,], the force constant for each angle. - - **angle_theta0** (Tensor, float32) - [M,], the equilibrium position value for each angle. - - Outputs: - - **frc_f** (Tensor, float32) - [N, 3], same as operator AngleForce(). - - **ene** (Tensor, float) - [N,], same as operator AngleAtomEnergy(). - - Supported Platforms: - ``GPU`` - """ - - @prim_attr_register - def __init__(self, angle_numbers): - validator.check_value_type('angle_numbers', angle_numbers, (int), self.name) - self.angle_numbers = angle_numbers - self.init_prim_io_names(inputs=['uint_crd_f', 'scaler_f', 'atom_a', 'atom_b', 'atom_c', 'angle_k', - 'angle_theta0'], - outputs=['frc_f', 'ene']) - self.add_prim_attr('angle_numbers', self.angle_numbers) - - 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): - cls_name = self.name - N = uint_crd_f_shape[0] - M = self.angle_numbers - validator.check_int(len(uint_crd_f_shape), 2, Rel.EQ, "uint_crd_f_dim", cls_name) - validator.check_int(len(scaler_f_shape), 1, Rel.EQ, "scaler_f_dim", cls_name) - validator.check_int(len(atom_a_shape), 1, Rel.EQ, "atom_a_dim", cls_name) - validator.check_int(len(atom_b_shape), 1, Rel.EQ, "atom_b_dim", cls_name) - validator.check_int(len(atom_c_shape), 1, Rel.EQ, "atom_c_dim", cls_name) - validator.check_int(len(angle_k_shape), 1, Rel.EQ, "angle_k_dim", cls_name) - validator.check_int(len(angle_theta0_shape), 1, Rel.EQ, "angle_theta0_dim", cls_name) - - validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f_shape[1]", cls_name) - validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f_shape", cls_name) - validator.check_int(atom_a_shape[0], M, Rel.EQ, "atom_a_shape", cls_name) - validator.check_int(atom_b_shape[0], M, Rel.EQ, "atom_b_shape", cls_name) - validator.check_int(atom_c_shape[0], M, Rel.EQ, "atom_c_shape", cls_name) - validator.check_int(angle_k_shape[0], M, Rel.EQ, "angle_k_shape", cls_name) - validator.check_int(angle_theta0_shape[0], M, Rel.EQ, "angle_theta0_shape", cls_name) - return uint_crd_f_shape, [N,] - - def infer_dtype(self, uint_crd_f_dtype, scaler_f_type, atom_a_type, atom_b_type, atom_c_type, angle_k_type, - angle_theta0_type): - validator.check_tensor_dtype_valid('uint_crd_f', uint_crd_f_dtype, [mstype.uint32], self.name) - validator.check_tensor_dtype_valid('scaler_f', scaler_f_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('atom_a', atom_a_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('atom_b', atom_b_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('atom_c', atom_c_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('angle_k', angle_k_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('angle_theta0', angle_theta0_type, [mstype.float32], self.name) - return angle_k_type, angle_k_type - - -class Dihedral14LJForce(PrimitiveWithInfer): - """ - Calculate the Lennard-Jones part of 1,4 dihedral force correction - for each necessary dihedral terms on the corresponding atoms. - - Assume the number of necessary dihedral 1,4 terms is M, the number of atoms is N, - and the number of Lennard-Jones types for all atoms is P, which means - there will be Q = P*(P+1)/2 types of possible Lennard-Jones interactions - for all kinds of atom pairs. - - .. math:: - 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 - - Args: - nb14_numbers (int32): the number of necessary dihedral 1,4 terms M. - atom_numbers (int32): the number of atoms N. - - Inputs: - - **uint_crd_f** (Tensor, uint32) - [N, 3], the unsigned int coordinate value of each atom. - - **LJ_type** (Tensor, int32) - [N,], the Lennard-Jones type of each atom. - - **charge** (Tensor, float32) - [N,], the charge of each atom. - - **boxlength_f** (Tensor, float32) - [3,], the length of molecular simulation box in 3 dimensions. - - **a_14** (Tensor, int32) - [M,], the first atom index of each dihedral 1,4 term. - - **b_14** (Tensor, int32) - [M,], the second atom index of each dihedral 1,4 term. - - **lj_scale_factor** (Tensor, float32) - [M,], the scale factor for the - 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. - 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. - Q is the number of atom pair. - - Outputs: - - **frc_f** (Tensor, float32) - [N, 3], the force felt by each atom. - - Supported Platforms: - ``GPU`` - """ - - @prim_attr_register - def __init__(self, nb14_numbers, atom_numbers): - validator.check_value_type('nb14_numbers', nb14_numbers, (int), self.name) - validator.check_value_type('atom_numbers', atom_numbers, (int), self.name) - self.dihedral_14_numbers = nb14_numbers - self.atom_numbers = atom_numbers - self.init_prim_io_names( - inputs=['uint_crd_f', 'LJtype', 'charge', 'boxlength_f', 'a_14', 'b_14', 'lj_scale_factor', - 'LJ_type_A', 'LJ_type_B'], - outputs=['frc_f']) - self.add_prim_attr('dihedral_14_numbers', self.dihedral_14_numbers) - self.add_prim_attr('atom_numbers', self.atom_numbers) - - 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): - cls_name = self.name - N = self.atom_numbers - Q = LJ_type_A_shape[0] - validator.check_int(len(uint_crd_f_shape), 2, Rel.EQ, "uint_crd_f_dim", cls_name) - validator.check_int(len(LJtype_shape), 1, Rel.EQ, "LJtype_dim", cls_name) - validator.check_int(len(charge_shape), 1, Rel.EQ, "charge_dim", cls_name) - validator.check_int(len(boxlength_f_shape), 1, Rel.EQ, "boxlength_f_dim", cls_name) - validator.check_int(len(a_14_shape), 1, Rel.EQ, "a_14_dim", cls_name) - validator.check_int(len(b_14_shape), 1, Rel.EQ, "b_14_dim", cls_name) - validator.check_int(len(lj_scale_factor_shape), 1, Rel.EQ, "lj_scale_factor_dim", cls_name) - validator.check_int(len(LJ_type_B_shape), 1, Rel.EQ, "LJ_type_B_dim", 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[1]", cls_name) - validator.check_int(LJtype_shape[0], N, Rel.EQ, "LJtype", cls_name) - validator.check_int(charge_shape[0], N, Rel.EQ, "charge", cls_name) - validator.check_int(boxlength_f_shape[0], 3, Rel.EQ, "boxlength_f", cls_name) - validator.check_int(LJ_type_B_shape[0], Q, Rel.EQ, "LJ_type_B", cls_name) - 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, - lj_scale_factor_type, LJ_type_A_type, LJ_type_B_type): - validator.check_tensor_dtype_valid('uint_crd_f', uint_crd_f_dtype, [mstype.uint32], self.name) - validator.check_tensor_dtype_valid('LJtype', LJtype_dtype, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('charge', charge_dtype, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('boxlength_f', boxlength_f_type, [mstype.float32], self.name) - - validator.check_tensor_dtype_valid('a_14', a_14_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('b_14', b_14_type, [mstype.int32], self.name) - - validator.check_tensor_dtype_valid('lj_scale_factor', lj_scale_factor_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('LJ_type_A', LJ_type_A_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('LJ_type_B', LJ_type_B_type, [mstype.float32], self.name) - return LJ_type_B_type - - -class Dihedral14LJEnergy(PrimitiveWithInfer): - """ - Calculate the Lennard-Jones part of 1,4 dihedral energy correction for - each necessary dihedral terms on the corresponding atoms. - - .. math:: - dr = (x_a-x_b, y_a-y_b, z_a-z-b) - .. math:: - E = k*(A/|dr|^{12} - B/|dr|^{6}) - - Args: - nb14_numbers (int32): the number of necessary dihedral 1,4 terms M. - atom_numbers (int32): the number of atoms N. - - Inputs: - - **uint_crd_f** (Tensor, uint32) - [N, 3], the unsigned int coordinate value of each atom. - - **LJ_type** (Tensor, int32) - [N,], the Lennard-Jones type of each atom. - - **charge** (Tensor, float32) - [N,], the charge of each atom. - - **boxlength_f** (Tensor, float32) - [3,], the length of molecular simulation box in 3 dimensions. - - **a_14** (Tensor, int32) - [M,], the first atom index of each dihedral 1,4 term. - - **b_14** (Tensor, int32) - [M,], the second atom index of each dihedral 1,4 term. - - **lj_scale_factor** (Tensor, float32) - [M,], the scale factor for the - 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. - 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. - Q is the number of atom pair. - - Outputs: - - **ene** (Tensor, float32) - [M,], the Lennard-Jones potential - energy correction for each necessary dihedral 1,4 term. - - Supported Platforms: - ``GPU`` - """ - - @prim_attr_register - def __init__(self, nb14_numbers, atom_numbers): - validator.check_value_type('nb14_numbers', nb14_numbers, (int), self.name) - validator.check_value_type('atom_numbers', atom_numbers, (int), self.name) - self.dihedral_14_numbers = nb14_numbers - self.atom_numbers = atom_numbers - - self.init_prim_io_names( - inputs=['uint_crd_f', 'LJtype', 'charge', 'boxlength_f', 'a_14', 'b_14', 'lj_scale_factor', - 'LJ_type_A', 'LJ_type_B'], - outputs=['ene']) - self.add_prim_attr('dihedral_14_numbers', self.dihedral_14_numbers) - self.add_prim_attr('atom_numbers', self.atom_numbers) - - 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): - cls_name = self.name - N = self.atom_numbers - Q = LJ_type_A_shape[0] - validator.check_int(len(uint_crd_f_shape), 2, Rel.EQ, "uint_crd_f_dim", cls_name) - validator.check_int(len(LJtype_shape), 1, Rel.EQ, "LJtype_dim", cls_name) - validator.check_int(len(charge_shape), 1, Rel.EQ, "charge_dim", cls_name) - validator.check_int(len(boxlength_f_shape), 1, Rel.EQ, "boxlength_f_dim", cls_name) - validator.check_int(len(a_14_shape), 1, Rel.EQ, "a_14_dim", cls_name) - validator.check_int(len(b_14_shape), 1, Rel.EQ, "b_14_dim", cls_name) - validator.check_int(len(lj_scale_factor_shape), 1, Rel.EQ, "lj_scale_factor_dim", cls_name) - validator.check_int(len(LJ_type_B_shape), 1, Rel.EQ, "LJ_type_B_dim", 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[1]", cls_name) - validator.check_int(LJtype_shape[0], N, Rel.EQ, "LJtype", cls_name) - validator.check_int(charge_shape[0], N, Rel.EQ, "charge", cls_name) - validator.check_int(boxlength_f_shape[0], 3, Rel.EQ, "boxlength_f", cls_name) - validator.check_int(LJ_type_B_shape[0], Q, Rel.EQ, "LJ_type_B", cls_name) - 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, - lj_scale_factor_type, LJ_type_A_type, LJ_type_B_type): - validator.check_tensor_dtype_valid('uint_crd_f', uint_crd_f_dtype, [mstype.uint32], self.name) - validator.check_tensor_dtype_valid('LJtype', LJtype_dtype, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('charge', charge_dtype, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('boxlength_f', boxlength_f_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('a_14', a_14_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('b_14', b_14_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('lj_scale_factor', lj_scale_factor_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('LJ_type_A', LJ_type_A_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('LJ_type_B', LJ_type_B_type, [mstype.float32], self.name) - - return LJ_type_A_type - - -class Dihedral14LJForceWithDirectCF(PrimitiveWithInfer): - """ - Calculate the Lennard-Jones part and the Coulomb part of force correction - for each necessary dihedral 1,4 terms. - - The calculation formula of the Lennard-Jones part is the same as operator - Dihedral14LJForce(), and the Coulomb part is as follows: - - .. math:: - dr = (x_a-x_b, y_a-y_b, z_a-z_b) - .. math:: - F = -k*q_a*q_b/|r|^3*dr - - Args: - nb14_numbers (int32): the number of necessary dihedral 1,4 terms M. - atom_numbers (int32): the number of atoms N. - - Inputs: - - **uint_crd_f** (Tensor, uint32) - [N, 3], the unsigned int coordinate value of each atom. - - **LJ_type** (Tensor, int32) - [N,], the Lennard-Jones type of each atom. - - **charge** (Tensor, float32) - [N,], the charge of each atom. - - **boxlength_f** (Tensor, float32) - [3,], the length of molecular simulation box in 3 dimensions. - - **a_14** (Tensor, int32) - [M,], the first atom index of each dihedral 1,4 term. - - **b_14** (Tensor, int32) - [M,], the second atom index of each dihedral 1,4 term. - - **lj_scale_factor** (Tensor, float32) - [M,], the scale factor for the - Lennard-Jones part of force correction of each dihedral 1,4 term. - - **cf_scale_factor** (Tensor, float) - [M,], the scale factor for the - 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. - 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. - Q is the number of atom pair. - - Outputs: - - **frc_f** (Tensor, float) - [N, 3], the force felt by each atom. - - Supported Platforms: - ``GPU`` - """ - - @prim_attr_register - def __init__(self, nb14_numbers, atom_numbers): - validator.check_value_type('nb14_numbers', nb14_numbers, (int), self.name) - validator.check_value_type('atom_numbers', atom_numbers, (int), self.name) - self.dihedral_14_numbers = nb14_numbers - self.atom_numbers = atom_numbers - - self.init_prim_io_names( - inputs=['uint_crd_f', 'LJtype', 'charge', 'boxlength_f', 'a_14', 'b_14', 'lj_scale_factor', - 'cf_scale_factor', - 'LJ_type_A', 'LJ_type_B'], - outputs=['frc_f']) - self.add_prim_attr('dihedral_14_numbers', self.dihedral_14_numbers) - self.add_prim_attr('atom_numbers', self.atom_numbers) - - 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): - cls_name = self.name - N = self.atom_numbers - Q = LJ_type_A_shape[0] - validator.check_int(len(uint_crd_f_shape), 2, Rel.EQ, "uint_crd_f_dim", cls_name) - validator.check_int(len(LJtype_shape), 1, Rel.EQ, "LJtype_dim", cls_name) - validator.check_int(len(charge_shape), 1, Rel.EQ, "charge_dim", cls_name) - validator.check_int(len(boxlength_f_shape), 1, Rel.EQ, "boxlength_f_dim", cls_name) - validator.check_int(len(a_14_shape), 1, Rel.EQ, "a_14_dim", cls_name) - validator.check_int(len(b_14_shape), 1, Rel.EQ, "b_14_dim", cls_name) - validator.check_int(len(lj_scale_factor_shape), 1, Rel.EQ, "lj_scale_factor_dim", cls_name) - validator.check_int(len(cf_scale_factor_shape), 1, Rel.EQ, "cf_scale_factor_dim", cls_name) - validator.check_int(len(LJ_type_B_shape), 1, Rel.EQ, "LJ_type_B_dim", cls_name) - - validator.check_int(uint_crd_f_shape[0], N, Rel.EQ, "uint_crd_f_shape[0]", cls_name) - validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f_shape[1]", cls_name) - validator.check_int(LJtype_shape[0], N, Rel.EQ, "LJtype_shape", cls_name) - validator.check_int(charge_shape[0], N, Rel.EQ, "charge_shape", cls_name) - validator.check_int(boxlength_f_shape[0], 3, Rel.EQ, "boxlength_f_shape", cls_name) - validator.check_int(LJ_type_B_shape[0], Q, Rel.EQ, "LJ_type_B_shape", cls_name) - 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, - lj_scale_factor_type, cf_scale_factor_type, LJ_type_A_type, LJ_type_B_type): - validator.check_tensor_dtype_valid('uint_crd_f', uint_crd_f_dtype, [mstype.uint32], self.name) - validator.check_tensor_dtype_valid('LJtype', LJtype_dtype, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('charge', charge_dtype, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('boxlength_f', boxlength_f_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('a_14', a_14_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('b_14', b_14_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('lj_scale_factor', lj_scale_factor_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('cf_scale_factor', cf_scale_factor_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('LJ_type_A', LJ_type_A_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('LJ_type_B', LJ_type_B_type, [mstype.float32], self.name) - - return LJ_type_A_type - - -class Dihedral14LJCFForceWithAtomEnergy(PrimitiveWithInfer): - """ - Calculate the Lennard-Jones and Coulumb energy correction and force correction - for each necessary dihedral 1,4 terms together and add them to the total force - and potential energy for each atom. - - The calculation formula of force correction is the same as operator - Dihedral14LJForceWithDirectCF(), and the energy correction part is the same - as operator Dihedral14LJEnergy() and Dihedral14CFEnergy(). - - Args: - nb14_numbers (int32): the number of necessary dihedral 1,4 terms M. - atom_numbers (int32): the number of atoms N. - - Inputs: - - **uint_crd_f** (Tensor, uint32) - [N, 3], the unsigned int coordinate value of each atom. - - **LJ_type** (Tensor, int32) - [N,], the Lennard-Jones type of each atom. - - **charge** (Tensor, float32) - [N,], the charge of each atom. - - **boxlength_f** (Tensor, float32) - [3,], the length of molecular simulation box in 3 dimensions. - - **a_14** (Tensor, int32) - [M,], the first atom index of each dihedral 1,4 term. - - **b_14** (Tensor, int32) - [M,], the second atom index of each dihedral 1,4 term. - - **lj_scale_factor** (Tensor, float32) - [M,], the scale factor for the - Lennard-Jones part of force correction of each dihedral 1,4 term. - - **cf_scale_factor** (Tensor, float) - [M,], the scale factor for the - 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. - 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. - Q is the number of atom pair. - - Outputs: - - **frc_f** (Tensor, float32) - [N, 3], the force felt by each atom. - - **atom_energy** (Tensor, float32) - [N,], the accumulated potential energy for each atom. - - Supported Platforms: - ``GPU`` - """ - - @prim_attr_register - def __init__(self, nb14_numbers, atom_numbers): - validator.check_value_type('nb14_numbers', nb14_numbers, (int), self.name) - validator.check_value_type('atom_numbers', atom_numbers, (int), self.name) - self.dihedral_14_numbers = nb14_numbers - self.atom_numbers = atom_numbers - - self.init_prim_io_names( - inputs=['uint_crd_f', 'LJtype', 'charge', 'boxlength_f', 'a_14', 'b_14', 'lj_scale_factor', - 'cf_scale_factor', - 'LJ_type_A', 'LJ_type_B'], - outputs=['frc_f', 'atom_energy']) - self.add_prim_attr('dihedral_14_numbers', self.dihedral_14_numbers) - self.add_prim_attr('atom_numbers', self.atom_numbers) - - 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): - cls_name = self.name - N = self.atom_numbers - Q = LJ_type_A_shape[0] - validator.check_int(len(uint_crd_f_shape), 2, Rel.EQ, "uint_crd_f_dim", cls_name) - validator.check_int(len(LJtype_shape), 1, Rel.EQ, "LJtype_dim", cls_name) - validator.check_int(len(charge_shape), 1, Rel.EQ, "charge_dim", cls_name) - validator.check_int(len(boxlength_f_shape), 1, Rel.EQ, "boxlength_f_dim", cls_name) - validator.check_int(len(a_14_shape), 1, Rel.EQ, "a_14_dim", cls_name) - validator.check_int(len(b_14_shape), 1, Rel.EQ, "b_14_dim", cls_name) - validator.check_int(len(lj_scale_factor_shape), 1, Rel.EQ, "lj_scale_factor_dim", cls_name) - validator.check_int(len(cf_scale_factor_shape), 1, Rel.EQ, "cf_scale_factor_dim", cls_name) - validator.check_int(len(LJ_type_B_shape), 1, Rel.EQ, "LJ_type_B_dim", cls_name) - - validator.check_int(uint_crd_f_shape[0], N, Rel.EQ, "uint_crd_f_shape[0]", cls_name) - validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f_shape[1]", cls_name) - validator.check_int(LJtype_shape[0], N, Rel.EQ, "LJtype_shape", cls_name) - validator.check_int(charge_shape[0], N, Rel.EQ, "charge_shape", cls_name) - validator.check_int(boxlength_f_shape[0], 3, Rel.EQ, "boxlength_f_shape", cls_name) - validator.check_int(LJ_type_B_shape[0], Q, Rel.EQ, "LJ_type_B_shape", cls_name) - 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, - lj_scale_factor_type, cf_scale_factor_type, LJ_type_A_type, LJ_type_B_type): - validator.check_tensor_dtype_valid('uint_crd_f', uint_crd_f_dtype, [mstype.uint32], self.name) - validator.check_tensor_dtype_valid('LJtype', LJtype_dtype, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('charge', charge_dtype, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('boxlength_f', boxlength_f_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('a_14', a_14_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('b_14', b_14_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('lj_scale_factor', lj_scale_factor_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('cf_scale_factor', cf_scale_factor_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('LJ_type_A', LJ_type_A_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('LJ_type_B', LJ_type_B_type, [mstype.float32], self.name) - - return charge_dtype, charge_dtype - - -class Dihedral14LJAtomEnergy(PrimitiveWithInfer): - """ - Add the potenrial energy caused by Lennard-Jones energy correction for each - necessary dihedral 1,4 terms to the total potential energy of each atom. - - The calculation formula is the same as operator Dihedral14LJEnergy(). - - Args: - nb14_numbers (int32): the number of necessary dihedral 1,4 terms M. - atom_numbers (int32): the number of atoms N. - - Inputs: - - **uint_crd_f** (Tensor, uint32) - [N, 3], the unsigned int coordinate value of each atom. - - **LJ_type** (Tensor, int32) - [N,], the Lennard-Jones type of each atom. - - **charge** (Tensor, float32) - [N,], the charge of each atom. - - **boxlength_f** (Tensor, float32) - [3,], the length of molecular simulation box in 3 dimensions. - - **a_14** (Tensor, int32) - [M,], the first atom index of each dihedral 1,4 term. - - **b_14** (Tensor, int32) - [M,], the second atom index of each dihedral 1,4 term. - - **lj_scale_factor** (Tensor, float32) - [M,], the scale factor for the - 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. - 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. - Q is the number of atom pair. - - Outputs: - - **ene** (Tensor, float32) - [N,], the accumulated potential energy of each atom. - - Supported Platforms: - ``GPU`` - """ - - @prim_attr_register - def __init__(self, nb14_numbers, atom_numbers): - validator.check_value_type('nb14_numbers', nb14_numbers, (int), self.name) - validator.check_value_type('atom_numbers', atom_numbers, (int), self.name) - self.dihedral_14_numbers = nb14_numbers - self.atom_numbers = atom_numbers - - self.init_prim_io_names( - inputs=['uint_crd_f', 'LJtype', 'charge', 'boxlength_f', 'a_14', 'b_14', 'lj_scale_factor', - 'LJ_type_A', 'LJ_type_B'], - outputs=['ene']) - self.add_prim_attr('dihedral_14_numbers', self.dihedral_14_numbers) - self.add_prim_attr('atom_numbers', self.atom_numbers) - - 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): - cls_name = self.name - N = self.atom_numbers - Q = LJ_type_A_shape[0] - validator.check_int(len(uint_crd_f_shape), 2, Rel.EQ, "uint_crd_f_dim", cls_name) - validator.check_int(len(LJtype_shape), 1, Rel.EQ, "LJtype_dim", cls_name) - validator.check_int(len(charge_shape), 1, Rel.EQ, "charge_dim", cls_name) - validator.check_int(len(boxlength_f_shape), 1, Rel.EQ, "boxlength_f_dim", cls_name) - validator.check_int(len(a_14_shape), 1, Rel.EQ, "a_14_dim", cls_name) - validator.check_int(len(b_14_shape), 1, Rel.EQ, "b_14_dim", cls_name) - validator.check_int(len(lj_scale_factor_shape), 1, Rel.EQ, "lj_scale_factor_dim", cls_name) - validator.check_int(len(LJ_type_B_shape), 1, Rel.EQ, "LJ_type_B_dim", cls_name) - - validator.check_int(uint_crd_f_shape[0], N, Rel.EQ, "uint_crd_f_shape[0]", cls_name) - validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f_shape[1]", cls_name) - validator.check_int(LJtype_shape[0], N, Rel.EQ, "LJtype_shape", cls_name) - validator.check_int(charge_shape[0], N, Rel.EQ, "charge_shape", cls_name) - validator.check_int(boxlength_f_shape[0], 3, Rel.EQ, "boxlength_f_shape", cls_name) - validator.check_int(LJ_type_B_shape[0], Q, Rel.EQ, "LJ_type_B_shape", cls_name) - return LJtype_shape - - def infer_dtype(self, uint_crd_f_dtype, LJtype_dtype, charge_dtype, boxlength_f_type, a_14_type, b_14_type, - lj_scale_factor_type, LJ_type_A_type, LJ_type_B_type): - validator.check_tensor_dtype_valid('uint_crd_f', uint_crd_f_dtype, [mstype.uint32], self.name) - validator.check_tensor_dtype_valid('LJtype', LJtype_dtype, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('charge', charge_dtype, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('boxlength_f', boxlength_f_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('a_14', a_14_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('b_14', b_14_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('lj_scale_factor', lj_scale_factor_type, [mstype.float32], - self.name) - validator.check_tensor_dtype_valid('LJ_type_A', LJ_type_A_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('LJ_type_B', LJ_type_B_type, [mstype.float32], self.name) - - return LJ_type_A_type - - -class Dihedral14CFEnergy(PrimitiveWithInfer): - """ - Calculate the Coulumb part of 1,4 dihedral energy correction for - each necessary dihedral terms on the corresponding atoms. - - .. math:: - - dr = (x_a-x_b, y_a-y_b, z_a-z_b) - - .. math:: - E = k*q_a*q_b/|dr| - - Args: - nb14_numbers (int32): the number of necessary dihedral 1,4 terms M. - atom_numbers (int32): the number of atoms N. - - Inputs: - - **uint_crd_f** (Tensor, uint32) - [N, 3], the unsigned int coordinate value of each atom. - - **LJ_type** (Tensor, int32) - [N,], the Lennard-Jones type of each atom. - - **charge** (Tensor, float32) - [N,], the charge of each atom. - - **boxlength_f** (Tensor, float32) - [3,], the length of molecular simulation box in 3 dimensions. - - **a_14** (Tensor, int32) - [M,], the first atom index of each dihedral 1,4 term. - - **b_14** (Tensor, int32) - [M,], the second atom index of each dihedral 1,4 term. - - **cf_scale_factor** (Tensor, float) - [M,], the scale factor for the - Coulomb part of force correction for each dihedral 1,4 terms. - - Outputs: - - **ene** (Tensor, float32) - [M,], the accumulated potential energy of each atom. - - Supported Platforms: - ``GPU`` - """ - - @prim_attr_register - def __init__(self, nb14_numbers, atom_numbers): - validator.check_value_type('nb14_numbers', nb14_numbers, (int), self.name) - validator.check_value_type('atom_numbers', atom_numbers, (int), self.name) - self.dihedral_14_numbers = nb14_numbers - self.atom_numbers = atom_numbers - - self.init_prim_io_names( - inputs=['uint_crd_f', 'LJtype', 'charge', 'boxlength_f', 'a_14', 'b_14', 'cj_scale_factor'], - outputs=['ene']) - self.add_prim_attr('dihedral_14_numbers', self.dihedral_14_numbers) - self.add_prim_attr('atom_numbers', self.atom_numbers) - - 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): - cls_name = self.name - N = self.atom_numbers - validator.check_int(len(uint_crd_f_shape), 2, Rel.EQ, "uint_crd_f_dim", cls_name) - validator.check_int(len(LJtype_shape), 1, Rel.EQ, "LJtype_dim", cls_name) - validator.check_int(len(charge_shape), 1, Rel.EQ, "charge_dim", cls_name) - validator.check_int(len(boxlength_f_shape), 1, Rel.EQ, "boxlength_f_dim", cls_name) - validator.check_int(len(a_14_shape), 1, Rel.EQ, "a_14_dim", cls_name) - validator.check_int(len(b_14_shape), 1, Rel.EQ, "b_14_dim", cls_name) - validator.check_int(len(cf_scale_factor_shape), 1, Rel.EQ, "cf_scale_factor_dim", cls_name) - - validator.check_int(uint_crd_f_shape[0], N, Rel.EQ, "uint_crd_f_shape[0]", cls_name) - validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f_shape[1]", cls_name) - validator.check_int(LJtype_shape[0], N, Rel.EQ, "LJtype_shape", cls_name) - validator.check_int(charge_shape[0], N, Rel.EQ, "charge_shape", cls_name) - validator.check_int(boxlength_f_shape[0], 3, Rel.EQ, "boxlength_f_shape", cls_name) - 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, - cf_scale_factor_type): - validator.check_tensor_dtype_valid('uint_crd_f', uint_crd_f_dtype, [mstype.uint32], self.name) - validator.check_tensor_dtype_valid('LJtype', LJtype_dtype, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('charge', charge_dtype, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('boxlength_f', boxlength_f_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('a_14', a_14_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('b_14', b_14_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('lj_scale_factor', cf_scale_factor_type, [mstype.float32], - self.name) - - return charge_dtype - - -class Dihedral14CFAtomEnergy(PrimitiveWithInfer): - """ - Add the potential energy caused by Coulumb energy correction for each - necessary dihedral 1,4 terms to the total potential energy of each atom. - - The calculation formula is the same as operator Dihedral14CFEnergy(). - - Args: - nb14_numbers (int32): the number of necessary dihedral 1,4 terms M. - atom_numbers (int32): the number of atoms N. - - Inputs: - - **uint_crd_f** (Tensor, uint32) - [N, 3], the unsigned int coordinate value of each atom. - - **LJ_type** (Tensor, int32) - [N,], the Lennard-Jones type of each atom. - - **charge** (Tensor, float32) - [N,], the charge of each atom. - - **boxlength_f** (Tensor, float32) - [3,], the length of molecular simulation box in 3 dimensions. - - **a_14** (Tensor, int32) - [M,], the first atom index of each dihedral 1,4 term. - - **b_14** (Tensor, int32) - [M,], the second atom index of each dihedral 1,4 term. - - **cf_scale_factor** (Tensor, float) - [M,], the scale factor for the - Coulomb part of force correction for each dihedral 1,4 terms. - - Outputs: - - **ene** (Tensor, float32) - [N,], the accumulated potential energy of each atom. - - - Supported Platforms: - ``GPU`` - """ - - @prim_attr_register - def __init__(self, nb14_numbers, atom_numbers): - validator.check_value_type('nb14_numbers', nb14_numbers, (int), self.name) - validator.check_value_type('atom_numbers', atom_numbers, (int), self.name) - self.dihedral_14_numbers = nb14_numbers - self.atom_numbers = atom_numbers - - self.init_prim_io_names( - inputs=['uint_crd_f', 'LJtype', 'charge', 'boxlength_f', 'a_14', 'b_14', 'cf_scale_factor'], - outputs=['ene']) - self.add_prim_attr('dihedral_14_numbers', self.dihedral_14_numbers) - self.add_prim_attr('atom_numbers', self.atom_numbers) - - 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): - cls_name = self.name - N = self.atom_numbers - validator.check_int(len(uint_crd_f_shape), 2, Rel.EQ, "uint_crd_f_dim", cls_name) - validator.check_int(len(LJtype_shape), 1, Rel.EQ, "LJtype_dim", cls_name) - validator.check_int(len(charge_shape), 1, Rel.EQ, "charge_dim", cls_name) - validator.check_int(len(boxlength_f_shape), 1, Rel.EQ, "boxlength_f_dim", cls_name) - validator.check_int(len(a_14_shape), 1, Rel.EQ, "a_14_dim", cls_name) - validator.check_int(len(b_14_shape), 1, Rel.EQ, "b_14_dim", cls_name) - validator.check_int(len(cf_scale_factor_shape), 1, Rel.EQ, "cf_scale_factor_dim", cls_name) - - validator.check_int(uint_crd_f_shape[0], N, Rel.EQ, "uint_crd_f_shape[0]", cls_name) - validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f_shape[1]", cls_name) - validator.check_int(LJtype_shape[0], N, Rel.EQ, "LJtype_shape", cls_name) - validator.check_int(charge_shape[0], N, Rel.EQ, "charge_shape", cls_name) - validator.check_int(boxlength_f_shape[0], 3, Rel.EQ, "boxlength_f_shape", cls_name) - return LJtype_shape - - def infer_dtype(self, uint_crd_f_dtype, LJtype_dtype, charge_dtype, boxlength_f_type, a_14_type, b_14_type, - cf_scale_factor_type): - validator.check_tensor_dtype_valid('uint_crd_f', uint_crd_f_dtype, [mstype.uint32], self.name) - validator.check_tensor_dtype_valid('LJtype', LJtype_dtype, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('charge', charge_dtype, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('boxlength_f', boxlength_f_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('a_14', a_14_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('b_14', b_14_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('cf_scale_factor', cf_scale_factor_type, [mstype.float32], - self.name) - - return charge_dtype - - -class MDIterationLeapFrog(PrimitiveWithInfer): - """ - One step of classical leap frog algorithm to solve the finite difference - Hamiltonian equations of motion for certain system, using Langevin dynamics - with Liu's thermostat scheme. Assume the number of atoms is N and the target - control temperature is T. - - Detailed iteration formula can be found in this paper: A unified thermostat - scheme for efficient configurational sampling for classical/quantum canonical - ensembles via molecular dynamics. DOI: 10.1063/1.4991621. - - Args: - float4_numbers(int32): total length to store random numbers. - atom_numbers(int32): the number of atoms N. - dt(float32): time step for finite difference. - half_dt(float32): half of time step for finite difference. - exp_gamma(float32): parameter in Liu's dynamic, equals exp(-gamma_ln * dt), - where gamma_ln is the firction factor in Langvin dynamics. - max_velocity(float32): the upper limit of velocity, when the veclocity overflows, - scale it to the upper limit. - is_max_velocity(int32): whether the max velocity control is open or not. - - Inputs: - - **mass_inverse** (Tensor, float32) - [N,], the inverse value of mass of each atom. - - **sqrt_mass** (Tensor, float32) - [N,], the inverse square root value - of effect mass in Liu's dynamics of each atom. - - Outputs: - - **vel** (Tensor, float32) - [N, 3], the velocity of each atom. - - **crd** (Tensor, float32) - [N, 3], the coordinate of each atom. - - **frc** (Tensor, float32) - [N, 3], the force felt by each atom. - - **acc** (Tensor, float32) - [N, 3], the acceleration of each atom. - - Supported Platforms: - ``GPU`` - """ - - @prim_attr_register - def __init__(self, float4_numbers, atom_numbers, half_dt, dt, exp_gamma, is_max_velocity, max_velocity): - validator.check_value_type('float4_numbers', float4_numbers, (int), self.name) - validator.check_value_type('atom_numbers', atom_numbers, (int), self.name) - validator.check_value_type('half_dt', half_dt, (float), self.name) - validator.check_value_type('dt', dt, (float), self.name) - validator.check_value_type('exp_gamma', exp_gamma, (float), self.name) - validator.check_value_type('is_max_velocity', is_max_velocity, (int), self.name) - validator.check_value_type('max_velocity', max_velocity, (float), self.name) - self.float4_numbers = float4_numbers - self.atom_numbers = atom_numbers - self.half_dt = half_dt - self.dt = dt - self.exp_gamma = exp_gamma - self.is_max_velocity = is_max_velocity - self.max_velocity = max_velocity - - self.init_prim_io_names( - inputs=['mass_inverse', 'sqrt_mass'], - outputs=['vel', 'crd', 'frc', 'acc']) - self.add_prim_attr('float4_numbers', self.float4_numbers) - self.add_prim_attr('atom_numbers', self.atom_numbers) - self.add_prim_attr('half_dt', self.half_dt) - self.add_prim_attr('dt', self.dt) - self.add_prim_attr('exp_gamma', self.exp_gamma) - self.add_prim_attr('is_max_velocity', self.is_max_velocity) - self.add_prim_attr('max_velocity', self.max_velocity) - - 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] - - def infer_dtype(self, mass_inverse_dtype, sqrt_mass_dtype): - validator.check_tensor_dtype_valid('mass_inverse', mass_inverse_dtype, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('sqrt_mass', sqrt_mass_dtype, [mstype.float32], self.name) - - return mass_inverse_dtype, mass_inverse_dtype, mass_inverse_dtype, mass_inverse_dtype - - -class PMEReciprocalForce(PrimitiveWithInfer): - """ - Calculate the reciprocal part of long-range Coulumb force using - PME(Particle Meshed Ewald) method. Assume the number of atoms is N. - - The detailed calculation formula of PME(Particle Meshed Ewald) method - can be found in this paper: A Smooth Particle Mesh Ewald Method. DOI: - 10.1063/1.470117. - - Args: - atom_numbers(int32): the number of atoms, N. - beta(float32): the PME beta parameter, determined by the - non-bond cutoff value and simulation precision tolerance. - fftx(int32): the number of points for Fourier transform in dimension X. - ffty(int32): the number of points for Fourier transform in dimension Y. - fftz(int32): the number of points for Fourier transform in dimension Z. - - Inputs: - - **boxlength** (Tensor, float32) - [3,], the length of simulation box in 3 dimensions. - - **uint_crd** (Tensor, uint32) - [N, 3], the unsigned int coordinates value of each atom. - - **charge** (Tensor, float32) - [N,], the charge carried by each atom. - - Outputs: - - **force** (Tensor, float32) - [N, 3], the force felt by each atom. - - Supported Platforms: - ``GPU`` - """ - - @prim_attr_register - def __init__(self, atom_numbers, beta, fftx, ffty, fftz): - validator.check_value_type('atom_numbers', atom_numbers, (int), self.name) - validator.check_value_type('beta', beta, (float), self.name) - validator.check_value_type('fftx', fftx, (int), self.name) - validator.check_value_type('ffty', ffty, (int), self.name) - validator.check_value_type('fftz', fftz, (int), self.name) - self.atom_numbers = atom_numbers - self.beta = beta - self.fftx = fftx - self.ffty = ffty - self.fftz = fftz - self.init_prim_io_names(inputs=['boxlength', 'uint_crd', 'charge'], - outputs=['force']) - self.add_prim_attr('atom_numbers', self.atom_numbers) - self.add_prim_attr('beta', self.beta) - self.add_prim_attr('fftx', self.fftx) - self.add_prim_attr('ffty', self.ffty) - self.add_prim_attr('fftz', self.fftz) - - def infer_shape(self, boxlength_shape, uint_crd_shape, charge_shape): - cls_name = self.name - N = self.atom_numbers - validator.check_int(len(uint_crd_shape), 2, Rel.EQ, "uint_crd_dim", cls_name) - validator.check_int(len(boxlength_shape), 1, Rel.EQ, "boxlength_dim", cls_name) - validator.check_int(len(charge_shape), 1, Rel.EQ, "charge_dim", cls_name) - - validator.check_int(uint_crd_shape[0], N, Rel.EQ, "uint_crd_shape[0]", cls_name) - validator.check_int(uint_crd_shape[1], 3, Rel.EQ, "uint_crd_shape[1]", cls_name) - validator.check_int(boxlength_shape[0], 3, Rel.EQ, "boxlength_shape", cls_name) - validator.check_int(charge_shape[0], N, Rel.EQ, "charge_shape", cls_name) - return uint_crd_shape - - def infer_dtype(self, boxlength_type, uint_crd_type, charge_type): - validator.check_tensor_dtype_valid('boxlength', boxlength_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('uint_crd', uint_crd_type, [mstype.uint32], self.name) - validator.check_tensor_dtype_valid('charge', charge_type, [mstype.float32], self.name) - return charge_type - - -class PMEExcludedForce(PrimitiveWithInfer): - """ - Calculate the excluded part of long-range Coulumb force using - PME(Particle Meshed Ewald) method. Assume the number of atoms is - N, and the length of excluded list is E. - - Args: - atom_numbers(int32): the number of atoms, N. - excluded_numbers(int32): the length of excluded list, E. - beta(float32): the PME beta parameter, determined by the - non-bond cutoff value and simulation precision tolerance. - - Inputs: - - **uint_crd** (Tensor, uint32) - [N, 3], the unsigned int coordinates value of each atom. - - **scaler** (Tensor, float32) - [3,], the scale factor between real space - coordinates and its unsigned int value. - - **charge** (Tensor, float32) - [N,], the charge carried by each atom. - - **excluded_list_start** (Tensor, int32) - [N,], the start excluded index - in excluded list for each atom. - - **excluded_list** (Tensor, int32) - [E,], the contiguous join of excluded - list of each atom. E is the number of excluded atoms. - - **excluded_atom_numbers** (Tensor, int32) - [N,], the number of atom excluded - in excluded list for each atom. - - Outputs: - - **force** (Tensor, float32) - [N, 3], the force felt by each atom. - - Supported Platforms: - ``GPU`` - """ - - @prim_attr_register - def __init__(self, atom_numbers, excluded_numbers, beta): - validator.check_value_type('atom_numbers', atom_numbers, (int), self.name) - validator.check_value_type('excluded_numbers', excluded_numbers, (int), self.name) - validator.check_value_type('beta', beta, (float), self.name) - self.atom_numbers = atom_numbers - self.excluded_numbers = excluded_numbers - self.beta = beta - self.init_prim_io_names( - inputs=['uint_crd', 'sacler', 'charge', 'excluded_list_start', 'excluded_list', 'excluded_atom_numbers'], - outputs=['force']) - self.add_prim_attr('atom_numbers', self.atom_numbers) - self.add_prim_attr('excluded_numbers', self.excluded_numbers) - self.add_prim_attr('beta', self.beta) - - def infer_shape(self, uint_crd_shape, sacler_shape, charge_shape, excluded_list_start_shape, excluded_list_shape, - excluded_atom_numbers_shape): - cls_name = self.name - N = self.atom_numbers - validator.check_int(len(uint_crd_shape), 2, Rel.EQ, "uint_crd_dim", cls_name) - validator.check_int(len(sacler_shape), 1, Rel.EQ, "sacler_dim", cls_name) - validator.check_int(len(charge_shape), 1, Rel.EQ, "charge_dim", cls_name) - validator.check_int(len(excluded_list_start_shape), 1, Rel.EQ, "excluded_list_start_dim", cls_name) - validator.check_int(len(excluded_atom_numbers_shape), 1, Rel.EQ, "excluded_atom_numbers_dim", cls_name) - - validator.check_int(uint_crd_shape[0], N, Rel.EQ, "uint_crd_shape[0]", cls_name) - validator.check_int(uint_crd_shape[1], 3, Rel.EQ, "uint_crd_shape[1]", cls_name) - validator.check_int(sacler_shape[0], 3, Rel.EQ, "sacler_shape", cls_name) - validator.check_int(charge_shape[0], N, Rel.EQ, "charge_shape", cls_name) - validator.check_int(excluded_list_start_shape[0], N, Rel.EQ, "excluded_list_start_shape", cls_name) - validator.check_int(excluded_atom_numbers_shape[0], N, Rel.EQ, "excluded_atom_numbers_shape", cls_name) - return uint_crd_shape - - def infer_dtype(self, uint_crd_type, sacler_type, charge_type, excluded_list_start_type, excluded_list_type, - excluded_atom_numbers_type): - validator.check_tensor_dtype_valid('sacler', sacler_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('uint_crd', uint_crd_type, [mstype.uint32], self.name) - validator.check_tensor_dtype_valid('charge', charge_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('excluded_list_start', excluded_list_start_type, [mstype.int32], - self.name) - validator.check_tensor_dtype_valid('excluded_list', excluded_list_type, [mstype.int32], - self.name) - validator.check_tensor_dtype_valid('excluded_atom_numbers', excluded_atom_numbers_type, [mstype.int32], - self.name) - return charge_type - - -class PMEEnergy(PrimitiveWithInfer): - """ - Calculate the Coulumb energy of the system using PME method. - - .. math:: - - E = sum_{ij} q_iq_j/r_{ij} - - Args: - atom_numbers(int32): the number of atoms, N. - excluded_numbers(int32): the length of excluded list, E. - beta(float32): the PME beta parameter, determined by the - non-bond cutoff value and simulation precision tolerance. - fftx(int32): the number of points for Fourier transform in dimension X. - ffty(int32): the number of points for Fourier transform in dimension Y. - fftz(int32): the number of points for Fourier transform in dimension Z. - - Inputs: - - **boxlength** (Tensor, float32) - [3,], the length of simulation box in 3 dimensions. - - **uint_crd** (Tensor, uint32) - [N, 3], the unsigned int coordinates value of each atom. - - **charge** (Tensor, float32) - [N,], the charge carried by 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. - - **scaler** (Tensor, float32) - [3,], the scale factor between real space - coordinates and its unsigned int value. - - **excluded_list_start** (Tensor, int32) - [N,], the start excluded index - in excluded list for each atom. - - **excluded_list** (Tensor, int32) - [E,], the contiguous join of excluded - list of each atom. E is the number of excluded atoms. - - **excluded_atom_numbers** (Tensor, int32) - [N,], the number of atom excluded - in excluded list for each atom. - - Outputs: - - **reciprocal_ene** (float32) - the reciprocal term of PME energy. - - **self_ene** (float32) - the self term of PME energy. - - **direct_ene** (float32) - the direct term of PME energy. - - **correction_ene** (float32) - the correction term of PME energy. - - Supported Platforms: - ``GPU`` - """ - - @prim_attr_register - def __init__(self, atom_numbers, excluded_numbers, beta, fftx, ffty, fftz): - validator.check_value_type('atom_numbers', atom_numbers, (int), self.name) - validator.check_value_type('excluded_numbers', excluded_numbers, (int), self.name) - validator.check_value_type('beta', beta, (float), self.name) - validator.check_value_type('fftx', fftx, (int), self.name) - validator.check_value_type('ffty', ffty, (int), self.name) - validator.check_value_type('fftz', fftz, (int), self.name) - self.atom_numbers = atom_numbers - self.excluded_numbers = excluded_numbers - self.beta = beta - self.fftx = fftx - self.ffty = ffty - self.fftz = fftz - self.init_prim_io_names( - inputs=['box_length', 'uint_crd', 'charge', 'nl_numbers', 'nl_serial', 'scaler', 'excluded_list_start', - 'excluded_list', 'excluded_atom_numbers'], - outputs=['reciprocal_ene', 'self_ene', 'direct_ene', 'correction_ene']) - self.add_prim_attr('atom_numbers', self.atom_numbers) - self.add_prim_attr('excluded_numbers', self.excluded_numbers) - self.add_prim_attr('beta', self.beta) - self.add_prim_attr('fftx', self.fftx) - self.add_prim_attr('ffty', self.ffty) - self.add_prim_attr('fftz', self.fftz) - - def infer_shape(self, box_length, uint_crd, charge, nl_numbers, nl_serial, scaler, excluded_list_start, - excluded_list, excluded_atom_numbers): - cls_name = self.name - N = self.atom_numbers - validator.check_int(len(uint_crd), 2, Rel.EQ, "uint_crd_dim", cls_name) - validator.check_int(len(box_length), 1, Rel.EQ, "sacler_dim", cls_name) - validator.check_int(len(charge), 1, Rel.EQ, "charge_dim", cls_name) - validator.check_int(len(nl_numbers), 1, Rel.EQ, "nl_numbers_dim", cls_name) - validator.check_int(len(nl_serial), 2, Rel.LE, "nl_serial_dim", cls_name) - validator.check_int(len(excluded_list_start), 1, Rel.EQ, "excluded_list_start_dim", cls_name) - validator.check_int(len(excluded_atom_numbers), 1, Rel.EQ, "excluded_atom_numbers_dim", cls_name) - validator.check_int(len(excluded_list), 1, Rel.GE, "excluded_list", cls_name) - - validator.check_int(uint_crd[0], N, Rel.EQ, "uint_crd_shape[0]", cls_name) - validator.check_int(uint_crd[1], 3, Rel.EQ, "uint_crd_shape[1]", cls_name) - validator.check_int(box_length[0], 3, Rel.EQ, "box_length_shape", cls_name) - validator.check_int(charge[0], N, Rel.EQ, "charge_shape", cls_name) - validator.check_int(nl_numbers[0], N, Rel.EQ, "nl_numbers_shape[0]", cls_name) - validator.check_int(nl_serial[0], N, Rel.LE, "nl_serial_shape[0]", cls_name) - validator.check_int(nl_serial[1], 800, Rel.LE, "nl_serial_shape[1]", cls_name) - validator.check_int(excluded_list_start[0], N, Rel.EQ, "excluded_list_start_shape", cls_name) - validator.check_int(excluded_atom_numbers[0], N, Rel.EQ, "excluded_atom_numbers_shape", cls_name) - validator.check_int(excluded_list[0], 0, Rel.GE, "excluded_list_shape", cls_name) - return (1,), (1,), (1,), (1,) - - def infer_dtype(self, box_length, uint_crd, charge, nl_numbers, nl_serial, scaler, excluded_list_start, - excluded_list, excluded_atom_numbers): - validator.check_tensor_dtype_valid('box_length', box_length, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('uint_crd', uint_crd, [mstype.uint32], self.name) - validator.check_tensor_dtype_valid('charge', charge, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('nl_numbers', nl_numbers, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('nl_serial', nl_serial, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('scaler', scaler, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('excluded_list_start', excluded_list_start, [mstype.int32], - self.name) - validator.check_tensor_dtype_valid('excluded_list', excluded_list, [mstype.int32], - self.name) - validator.check_tensor_dtype_valid('excluded_atom_numbers', excluded_atom_numbers, [mstype.int32], - self.name) - return charge, charge, charge, charge - - -class LJEnergy(PrimitiveWithInfer): - """ - Calculate the Van der Waals interaction energy described by Lennard-Jones - potential for each atom. Assume the number of atoms is N, and the number - of Lennard-Jones types for all atoms is P, which means there will be - Q = P*(P+1)/2 types of possible Lennard-Jones interactions for all kinds - of atom pairs. - - - .. math:: - - dr = (x_a-x_b, y_a-y_b, z_a-z_b) - - .. math:: - E = A/|dr|^{12} - B/|dr|^{6} - - Agrs: - atom_numbers(int32): the number of atoms, N. - cutoff_square(float32): the square value of cutoff. - - Inputs: - - **uint_crd** (Tensor, uint32) - [N, 3], the unsigned int coordinate value of each atom. - - **LJtype** (Tensor, int32) - [N,], the Lennard-Jones type of each atom. - - **charge** (Tensor, float32) - [N,], the charge carried by each atom. - - **scaler** (Tensor, float32) - [3,], the scale factor between real - space coordinate and its unsigned int value. - - **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. - - **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. - Q is the number of atom pair. - - Outputs: - - **d_LJ_energy_atom** (Tensor, float32) - [N,], the Lennard-Jones potential energy of each atom. - - **d_LJ_energy_sum** (float32), the sum of Lennard-Jones potential energy of each atom. - - Supported Platforms: - ``GPU`` - """ - - @prim_attr_register - def __init__(self, atom_numbers, cutoff_square): - validator.check_value_type('atom_numbers', atom_numbers, (int), self.name) - validator.check_value_type('cutoff_square', cutoff_square, (float), self.name) - self.atom_numbers = atom_numbers - self.cutoff_square = cutoff_square - self.init_prim_io_names( - inputs=['uint_crd', 'LJtype', 'charge', 'scaler', 'nl_numbers', 'nl_serial', 'd_LJ_A', 'd_LJ_B'], - outputs=['d_LJ_energy_atom']) - self.add_prim_attr('atom_numbers', self.atom_numbers) - 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): - cls_name = self.name - N = self.atom_numbers - Q = d_LJ_A[0] - validator.check_int(len(uint_crd), 2, Rel.EQ, "uint_crd_dim", cls_name) - validator.check_int(len(LJtype), 1, Rel.EQ, "LJtype_dim", cls_name) - validator.check_int(len(charge), 1, Rel.EQ, "charge_dim", cls_name) - validator.check_int(len(scaler), 1, Rel.EQ, "scaler_dim", cls_name) - validator.check_int(len(nl_numbers), 1, Rel.EQ, "nl_numbers_dim", cls_name) - validator.check_int(len(nl_serial), 2, Rel.EQ, "nl_serial_dim", cls_name) - validator.check_int(len(scaler), 1, Rel.EQ, "scaler_dim", cls_name) - validator.check_int(len(d_LJ_B), 1, Rel.EQ, "d_LJ_B_dim", cls_name) - - validator.check_int(uint_crd[0], N, Rel.EQ, "uint_crd_shape[0]", cls_name) - validator.check_int(uint_crd[1], 3, Rel.EQ, "uint_crd_shape[1]", cls_name) - validator.check_int(LJtype[0], N, Rel.EQ, "LJtype_shape", cls_name) - validator.check_int(charge[0], N, Rel.EQ, "charge_shape", cls_name) - validator.check_int(scaler[0], 3, Rel.EQ, "scaler_shape", cls_name) - validator.check_int(nl_numbers[0], N, Rel.EQ, "nl_numbers_shape", cls_name) - validator.check_int(nl_serial[0], N, Rel.EQ, "nl_serial_shape[0]", cls_name) - validator.check_int(nl_serial[1], 800, Rel.LE, "nl_serial_shape[1]", cls_name) - validator.check_int(scaler[0], 3, Rel.EQ, "scaler_shape", cls_name) - validator.check_int(d_LJ_B[0], Q, Rel.EQ, "d_LJ_B_shape[0]", cls_name) - return charge - - def infer_dtype(self, uint_crd, LJtype, charge, scaler, nl_numbers, nl_serial, d_LJ_A, d_LJ_B): - validator.check_tensor_dtype_valid('uint_crd', uint_crd, [mstype.uint32], self.name) - validator.check_tensor_dtype_valid('LJtype', LJtype, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('charge', charge, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('scaler', scaler, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('nl_numbers', nl_numbers, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('nl_serial', nl_serial, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('d_LJ_A', d_LJ_A, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('d_LJ_B', d_LJ_B, [mstype.float32], self.name) - return charge - - -class LJForce(PrimitiveWithInfer): - """ - Calculate the Van der Waals interaction force described by Lennard-Jones - potential energy for each atom. - - .. math:: - - dr = (x_a-x_b, y_a-y_b, z_a-z_b) - - .. math:: - - F = (-12*A/|dr|^{14} + 6*B/|dr|^{8}) * dr - - Agrs: - atom_numbers(int32): the number of atoms, N. - cutoff_square(float32): the square value of cutoff. - - Inputs: - - **uint_crd** (Tensor, uint32) - [N, 3], the unsigned int coordinate value of each atom. - - **LJtype** (Tensor, int32) - [N,], the Lennard-Jones type of each atom. - - **charge** (Tensor, float32) - [N,], the charge carried by each atom. - - **scaler** (Tensor, float32) - [3,], the scale factor between real - space coordinate and its unsigned int value. - - **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. - - **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. - Q is the number of atom pair. - - outputs: - - **frc** (Tensor, float32) - [N, 3], the force felt by each atom. - - Supported Platforms: - ``GPU`` - """ - - @prim_attr_register - def __init__(self, atom_numbers, cutoff_square): - validator.check_value_type('atom_numbers', atom_numbers, (int), self.name) - validator.check_value_type('cutoff_square', cutoff_square, (float), self.name) - self.atom_numbers = atom_numbers - self.cutoff_square = cutoff_square - self.init_prim_io_names( - inputs=['uint_crd', 'LJtype', 'charge', 'scaler', 'nl_numbers', 'nl_serial', 'd_LJ_A', 'd_LJ_B'], - outputs=['frc']) - self.add_prim_attr('atom_numbers', self.atom_numbers) - 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): - cls_name = self.name - N = self.atom_numbers - Q = d_LJ_A[0] - validator.check_int(len(uint_crd), 2, Rel.EQ, "uint_crd_dim", cls_name) - validator.check_int(len(LJtype), 1, Rel.EQ, "LJtype_dim", cls_name) - validator.check_int(len(charge), 1, Rel.EQ, "charge_dim", cls_name) - validator.check_int(len(scaler), 1, Rel.EQ, "scaler_dim", cls_name) - validator.check_int(len(nl_numbers), 1, Rel.EQ, "nl_numbers_dim", cls_name) - validator.check_int(len(nl_serial), 2, Rel.EQ, "nl_serial_dim", cls_name) - validator.check_int(len(scaler), 1, Rel.EQ, "scaler_dim", cls_name) - validator.check_int(len(d_LJ_B), 1, Rel.EQ, "d_LJ_B_dim", cls_name) - - validator.check_int(uint_crd[0], N, Rel.EQ, "uint_crd_shape[0]", cls_name) - validator.check_int(uint_crd[1], 3, Rel.EQ, "uint_crd_shape[1]", cls_name) - validator.check_int(LJtype[0], N, Rel.EQ, "LJtype_shape", cls_name) - validator.check_int(charge[0], N, Rel.EQ, "charge_shape", cls_name) - validator.check_int(scaler[0], 3, Rel.EQ, "scaler_shape", cls_name) - validator.check_int(nl_numbers[0], N, Rel.EQ, "nl_numbers_shape", cls_name) - validator.check_int(nl_serial[0], N, Rel.EQ, "nl_serial_shape[0]", cls_name) - validator.check_int(nl_serial[1], 800, Rel.LE, "nl_serial_shape[1]", cls_name) - validator.check_int(scaler[0], 3, Rel.EQ, "scaler_shape", cls_name) - validator.check_int(d_LJ_B[0], Q, Rel.EQ, "d_LJ_B_shape[0]", cls_name) - return uint_crd - - def infer_dtype(self, uint_crd, LJtype, charge, scaler, nl_numbers, nl_serial, d_LJ_A, d_LJ_B): - validator.check_tensor_dtype_valid('uint_crd', uint_crd, [mstype.uint32], self.name) - validator.check_tensor_dtype_valid('LJtype', LJtype, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('charge', charge, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('scaler', scaler, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('nl_numbers', nl_numbers, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('nl_serial', nl_serial, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('d_LJ_A', d_LJ_A, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('d_LJ_B', d_LJ_B, [mstype.float32], self.name) - return charge - - -class LJForceWithPMEDirectForce(PrimitiveWithInfer): - """ - Calculate the Lennard-Jones force and PME direct force together. - - The calculation formula of Lennard-Jones part is the same as operator - LJForce(), and the PME direct part is within PME method. - - Agrs: - atom_numbers(int32): the number of atoms, N. - cutoff_square(float32): the square value of cutoff. - pme_beta(float32): PME beta parameter, same as operator PMEReciprocalForce(). - - Inputs: - - **uint_crd** (Tensor, uint32) - [N, 3], the unsigned int coordinate value of each atom. - - **LJtype** (Tensor, int32) - [N,], the Lennard-Jones type of each atom. - - **charge** (Tensor, float32) - [N,], the charge carried by each atom. - - **scaler** (Tensor, float32) - [3,], the scale factor between real - space coordinate and its unsigned int value. - - **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. - - **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. - Q is the number of atom pair. - - Outputs: - - **frc** (Tensor, float32), [N, 3], the force felt by each atom. - - Supported Platforms: - ``GPU`` - """ - - @prim_attr_register - def __init__(self, atom_numbers, cutoff, pme_beta): - validator.check_value_type('atom_numbers', atom_numbers, (int), self.name) - validator.check_value_type('cutoff', cutoff, (float), self.name) - validator.check_value_type('pme_beta', pme_beta, (float), self.name) - self.atom_numbers = atom_numbers - self.cutoff = cutoff - self.pme_beta = pme_beta - self.init_prim_io_names( - inputs=['uint_crd', 'LJtype', 'charge', 'scaler', 'nl_numbers', 'nl_serial', 'd_LJ_A', 'd_LJ_B'], - outputs=['frc']) - self.add_prim_attr('atom_numbers', self.atom_numbers) - self.add_prim_attr('cutoff', self.cutoff) - 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): - cls_name = self.name - N = self.atom_numbers - Q = d_LJ_A[0] - validator.check_int(len(uint_crd), 2, Rel.EQ, "uint_crd_dim", cls_name) - validator.check_int(len(LJtype), 1, Rel.EQ, "LJtype_dim", cls_name) - validator.check_int(len(charge), 1, Rel.EQ, "charge_dim", cls_name) - validator.check_int(len(scaler), 1, Rel.EQ, "scaler_dim", cls_name) - validator.check_int(len(nl_numbers), 1, Rel.EQ, "nl_numbers_dim", cls_name) - validator.check_int(len(nl_serial), 2, Rel.EQ, "nl_serial_dim", cls_name) - validator.check_int(len(scaler), 1, Rel.EQ, "scaler_dim", cls_name) - validator.check_int(len(d_LJ_B), 1, Rel.EQ, "d_LJ_B_dim", cls_name) - - validator.check_int(uint_crd[0], N, Rel.EQ, "uint_crd_shape[0]", cls_name) - validator.check_int(uint_crd[1], 3, Rel.EQ, "uint_crd_shape[1]", cls_name) - validator.check_int(LJtype[0], N, Rel.EQ, "LJtype_shape", cls_name) - validator.check_int(charge[0], N, Rel.EQ, "charge_shape", cls_name) - validator.check_int(scaler[0], 3, Rel.EQ, "scaler_shape", cls_name) - validator.check_int(nl_numbers[0], N, Rel.EQ, "nl_numbers_shape", cls_name) - validator.check_int(nl_serial[0], N, Rel.EQ, "nl_serial_shape[0]", cls_name) - validator.check_int(nl_serial[1], 800, Rel.LE, "nl_serial_shape[1]", cls_name) - validator.check_int(scaler[0], 3, Rel.EQ, "scaler_shape", cls_name) - validator.check_int(d_LJ_B[0], Q, Rel.EQ, "d_LJ_B_shape[0]", cls_name) - return uint_crd - - def infer_dtype(self, uint_crd, LJtype, charge, scaler, nl_numbers, nl_serial, d_LJ_A, d_LJ_B): - validator.check_tensor_dtype_valid('uint_crd', uint_crd, [mstype.uint32], self.name) - validator.check_tensor_dtype_valid('LJtype', LJtype, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('charge', charge, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('scaler', scaler, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('nl_numbers', nl_numbers, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('nl_serial', nl_serial, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('d_LJ_A', d_LJ_A, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('d_LJ_B', d_LJ_B, [mstype.float32], self.name) - return charge - - -class GetCenterOfGeometry(PrimitiveWithInfer): - """ - Get Center Of Geometry. - - Supported Platforms: - ``GPU`` - """ - - @prim_attr_register - 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_inverse = center_numbers_inverse - self.add_prim_attr('center_numbers', self.center_numbers) - self.add_prim_attr('center_numbers_inverse', self.center_numbers_inverse) - self.init_prim_io_names( - inputs=['center_atoms', 'crd_f'], - outputs=['center_of_geometry_f']) - - def infer_shape(self, center_atoms_shape, crd_f_shape): - cls_name = self.name - N = self.center_numbers - 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[0]", cls_name) - validator.check_int(crd_f_shape[1], 3, Rel.EQ, "crd_f[1]", cls_name) - return [3,] - - def infer_dtype(self, center_atoms_dtype, crd_f_dtype): - validator.check_tensor_dtype_valid('center_atoms', center_atoms_dtype, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('crd_f', crd_f_dtype, [mstype.float32], self.name) - - return crd_f_dtype - - -class MDTemperature(PrimitiveWithInfer): - """ - Calculate the MD Temperature. - - Calculate the temperature. - - Supported Platforms: - ``GPU`` - """ - - @prim_attr_register - def __init__(self, residue_numbers, atom_numbers): - assert isinstance(residue_numbers, int) - assert isinstance(atom_numbers, int) - self.residue_numbers = residue_numbers - self.atom_numbers = atom_numbers - self.add_prim_attr('residue_numbers', self.residue_numbers) - self.add_prim_attr('atom_numbers', self.atom_numbers) - self.init_prim_io_names( - inputs=['start', 'end', 'atom_vel_f', 'atom_mass'], - outputs=['ek']) - - def infer_shape(self, start_shape, end_shape, atom_vel_f_shape, atom_mass_shape): - cls_name = self.name - N = self.residue_numbers - M = self.atom_numbers - validator.check_int(len(start_shape), 1, Rel.EQ, "start", cls_name) - validator.check_int(start_shape[0], N, Rel.EQ, "end", cls_name) - validator.check_int(len(end_shape), 1, Rel.EQ, "start", cls_name) - validator.check_int(end_shape[0], N, Rel.EQ, "end", cls_name) - validator.check_int(atom_vel_f_shape[0], M, Rel.EQ, "atom_vel_f", cls_name) - validator.check_int(atom_vel_f_shape[1], 3, Rel.EQ, "atom_vel_f", cls_name) - validator.check_int(len(atom_mass_shape), 1, Rel.EQ, "atom_mass", cls_name) - validator.check_int(atom_mass_shape[0], M, Rel.EQ, "atom_mass", cls_name) - return [N,] - - def infer_dtype(self, start_dtype, end_dtype, atom_vel_f_dtype, atom_mass_dtype): - validator.check_tensor_dtype_valid('start', start_dtype, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('end', end_dtype, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('atom_vel_f', atom_vel_f_dtype, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('atom_mass', atom_mass_dtype, [mstype.float32], self.name) - return atom_mass_dtype - - -class NeighborListUpdate(PrimitiveWithInfer): - """ - Update (or construct if first time) the Verlet neighbor list for the - calculation of short-ranged force. Assume the number of atoms is N, - the number of grids divided is G, the maximum number of atoms in one - grid is M, the maximum number of atoms in single atom's neighbor list - is L, and the number of total atom in excluded list is E. - - Args: - grid_numbers(int32): the total number of grids divided. - refresh_count(int32): the counter which counts how many - iteration steps have passed since last update. - not_first_time(int32): whether to construct the neighbor - list first time or not. - Nxy(int32): the total number of grids divided in xy plane. - excluded_atom_numbers(int32): the total atom numbers in the excluded list. - cutoff(float32): the cutoff distance for short-range force calculation. - skin(float32): the overflow value of cutoff to maintain a neighbor list. - cutoff_square(float32): the suqare value of cutoff. - half_skin_square(float32): skin*skin/4, indicates the maximum - square value of the distance atom allowed to move between two updates. - cutoff_with_skin(float32): cutoff + skin, indicates the - radius of the neighbor list for each atom. - half_cutoff_with_skin(float32): cutoff_with_skin/2. - cutoff_with_skin_square(float32): the square value of cutoff_with_skin. - refresh_interval(int32): the number of iteration steps between two updates of neighbor list. - max_atom_in_grid_numbers(int32): the maximum number of atoms in one grid. - - Inputs: - - **atom_numbers_in_grid_bucket** (Tensor, int32) - [G,], the number of atoms in each grid bucket. - - **bucket** (Tensor, int32) - (Tensor,int32) - [G, M], the atom indices in each grid bucket. - - **crd** (Tensor, float32) - [N,], the coordinates of each atom. - - **box_length** (Tensor, float32) - [3,], the length of 3 dimensions of the simulation box. - - **grid_N** (Tensor, int32) - [3,], the number of grids divided of 3 dimensions of the simulation box. - - **grid_length_inverse** (float32) - the inverse value of grid length. - - **atom_in_grid_serial** (Tensor, int32) - [N,], the grid index for each atom. - - **old_crd** (Tensor, float32) - [N, 3], the coordinates before update of each atom. - - **crd_to_uint_crd_cof** (Tensor, float32) - [3,], the scale factor - between the unsigned int value and the real space coordinates. - - **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. - 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_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 - the real space coordinates and the unsigned int value. - - **excluded_list_start** (Tensor, int32) - [N,], the start excluded index in excluded list for each atom. - - **excluded_numbers** (Tensor, int32) - [N,], the number of atom excluded in excluded list for each atom. - - **excluded_list** (Tensor, int32) - [E,], the contiguous join of excluded list of each atom. - - **need_refresh_flag** (Tensor, int32) - [N,], whether the neighbor list of each atom need update or not. - - Outputs: - - **res** (float32) - - Supported Platforms: - ``GPU`` - """ - - @prim_attr_register - def __init__(self, grid_numbers, atom_numbers, refresh_count, not_first_time, Nxy, excluded_atom_numbers, - cutoff_square, half_skin_square, cutoff_with_skin, half_cutoff_with_skin, cutoff_with_skin_square, - refresh_interval=20, cutoff=10.0, skin=2.0, max_atom_in_grid_numbers=64, max_neighbor_numbers=800): - self.grid_numbers = grid_numbers - self.atom_numbers = atom_numbers - self.refresh_count = refresh_count - self.refresh_interval = refresh_interval - self.not_first_time = not_first_time - self.cutoff = cutoff - self.skin = skin - self.max_atom_in_grid_numbers = max_atom_in_grid_numbers - self.Nxy = Nxy - self.excluded_atom_numbers = excluded_atom_numbers - self.cutoff_square = cutoff_square - self.half_skin_square = half_skin_square - self.cutoff_with_skin = cutoff_with_skin - self.half_cutoff_with_skin = half_cutoff_with_skin - self.cutoff_with_skin_square = cutoff_with_skin_square - self.max_neighbor_numbers = max_neighbor_numbers - self.init_prim_io_names( - inputs=['atom_numbers_in_grid_bucket', 'bucket', 'crd', 'box_length', 'grid_N', 'grid_length_inverse', - 'atom_in_grid_serial', 'old_crd', 'crd_to_uint_crd_cof', 'uint_crd', 'gpointer', 'nl_atom_numbers', - 'nl_atom_serial', 'uint_dr_to_dr_cof', 'excluded_list_start', 'excluded_list', 'excluded_numbers', - 'need_refresh_flag'], outputs=['res']) - - self.add_prim_attr('grid_numbers', self.grid_numbers) - self.add_prim_attr('atom_numbers', self.atom_numbers) - self.add_prim_attr('refresh_count', self.refresh_count) - self.add_prim_attr('refresh_interval', self.refresh_interval) - self.add_prim_attr('not_first_time', self.not_first_time) - self.add_prim_attr('cutoff', self.cutoff) - self.add_prim_attr('skin', self.skin) - self.add_prim_attr('max_atom_in_grid_numbers', self.max_atom_in_grid_numbers) - self.add_prim_attr('Nxy', self.Nxy) - self.add_prim_attr('excluded_atom_numbers', self.excluded_atom_numbers) - self.add_prim_attr('cutoff_square', self.cutoff_square) - self.add_prim_attr('half_skin_square', self.half_skin_square) - self.add_prim_attr('cutoff_with_skin', self.cutoff_with_skin) - self.add_prim_attr('half_cutoff_with_skin', self.half_cutoff_with_skin) - self.add_prim_attr('cutoff_with_skin_square', self.cutoff_with_skin_square) - - def infer_shape(self, atom_numbers_in_grid_bucket_shape, bucket_shape, crd_shape, box_length_shape, grid_N_shape, - grid_length_inverse_shape, atom_in_grid_serial_shape, old_crd_shape, crd_to_uint_crd_cof_shape, - uint_crd_shape, gpointer_shape, nl_atom_numbers_shape, nl_atom_serial_shape, - uint_dr_to_dr_cof_shape, excluded_list_start_shape, excluded_list_shape, excluded_numbers_shape, - need_refresh_flag_shape): - assert len(atom_numbers_in_grid_bucket_shape) == 1 - assert len(bucket_shape) == 2 - assert len(crd_shape) == 2 - assert len(box_length_shape) == 1 - assert len(grid_N_shape) == 1 - assert len(grid_length_inverse_shape) == 1 - assert len(atom_in_grid_serial_shape) == 1 - assert len(old_crd_shape) == 2 - assert len(crd_to_uint_crd_cof_shape) == 1 - assert len(uint_crd_shape) == 2 - assert len(gpointer_shape) == 2 - assert len(nl_atom_numbers_shape) == 1 - assert len(nl_atom_serial_shape) == 2 - assert len(uint_dr_to_dr_cof_shape) == 1 - assert len(excluded_list_start_shape) == 1 - assert len(excluded_list_shape) == 1 - assert len(excluded_numbers_shape) == 1 - assert len(need_refresh_flag_shape) == 1 - - validator.check_int(atom_numbers_in_grid_bucket_shape[0], self.grid_numbers, Rel.EQ, - "atom_numbers_in_grid_bucket", self.name) - validator.check_int(bucket_shape[0], self.grid_numbers, Rel.EQ, "bucket", self.name) - validator.check_int(bucket_shape[1], self.max_atom_in_grid_numbers, Rel.EQ, "bucket", self.name) - validator.check_int(crd_shape[0], self.atom_numbers, Rel.EQ, "crd", self.name) - validator.check_int(crd_shape[1], 3, Rel.EQ, "crd", self.name) - validator.check_int(box_length_shape[0], 3, Rel.EQ, "box_length", self.name) - validator.check_int(grid_N_shape[0], 3, Rel.EQ, "grid_N", self.name) - validator.check_int(grid_length_inverse_shape[0], 3, Rel.EQ, "grid_length_inverse", self.name) - validator.check_int(atom_in_grid_serial_shape[0], self.atom_numbers, Rel.EQ, "atom_in_grid_serial", - self.name) - validator.check_int(old_crd_shape[0], self.atom_numbers, Rel.EQ, "old_crd", self.name) - validator.check_int(old_crd_shape[1], 3, Rel.EQ, "old_crd", self.name) - validator.check_int(crd_to_uint_crd_cof_shape[0], 3, Rel.EQ, "crd_to_uint_crd_cof", self.name) - validator.check_int(uint_crd_shape[0], self.atom_numbers, Rel.EQ, "uint_crd", self.name) - validator.check_int(uint_crd_shape[1], 3, Rel.EQ, "uint_crd", self.name) - validator.check_int(gpointer_shape[0], self.grid_numbers, Rel.EQ, "gpointer", self.name) - validator.check_int(gpointer_shape[1], 125, Rel.EQ, "gpointer", self.name) - validator.check_int(nl_atom_numbers_shape[0], self.atom_numbers, Rel.EQ, "nl_atom_numbers", self.name) - validator.check_int(nl_atom_serial_shape[0], self.atom_numbers, Rel.EQ, "nl_atom_serial", self.name) - validator.check_int(nl_atom_serial_shape[1], self.max_neighbor_numbers, Rel.EQ, "nl_atom_serial", - self.name) - validator.check_int(uint_dr_to_dr_cof_shape[0], 3, Rel.EQ, "uint_dr_to_dr_cof", self.name) - validator.check_int(excluded_list_start_shape[0], self.atom_numbers, Rel.EQ, "excluded_list_start", - self.name) - validator.check_int(excluded_list_shape[0], self.excluded_atom_numbers, Rel.EQ, "excluded_list", - self.name) - validator.check_int(excluded_numbers_shape[0], self.atom_numbers, Rel.EQ, "excluded_numbers", self.name) - validator.check_int(need_refresh_flag_shape[0], 1, Rel.EQ, "need_refresh_flag", self.name) - - return [1,] - - def infer_dtype(self, atom_numbers_in_grid_bucket_dtype, bucket_dtype, crd_dtype, box_length_dtype, grid_N_dtype, - grid_length_inverse_dtype, atom_in_grid_serial_dtype, old_crd_dtype, crd_to_uint_crd_cof_dtype, - uint_crd_dtype, gpointer_dtype, nl_atom_numbers_dtype, nl_atom_serial_dtype, - uint_dr_to_dr_cof_dtype, excluded_list_start_dtype, excluded_list_dtype, excluded_numbers_dtype, - need_refresh_flag_dtype): - validator.check_tensor_dtype_valid('atom_numbers_in_grid_bucket', atom_numbers_in_grid_bucket_dtype, - [mstype.int32], self.name) - validator.check_tensor_dtype_valid('bucket', bucket_dtype, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('crd', crd_dtype, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('box_length', box_length_dtype, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('grid_N', grid_N_dtype, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('grid_length_inverse', grid_length_inverse_dtype, [mstype.float32], - self.name) - validator.check_tensor_dtype_valid('atom_in_grid_serial', atom_in_grid_serial_dtype, [mstype.int32], - self.name) - validator.check_tensor_dtype_valid('old_crd', old_crd_dtype, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('crd_to_uint_crd_cof', crd_to_uint_crd_cof_dtype, [mstype.float32], - self.name) - validator.check_tensor_dtype_valid('uint_crd', uint_crd_dtype, [mstype.uint32], self.name) - validator.check_tensor_dtype_valid('gpointer', gpointer_dtype, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('nl_atom_numbers', nl_atom_numbers_dtype, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('nl_atom_serial', nl_atom_serial_dtype, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('uint_dr_to_dr_cof', uint_dr_to_dr_cof_dtype, [mstype.float32], - self.name) - validator.check_tensor_dtype_valid('excluded_list_start', excluded_list_start_dtype, [mstype.int32], - self.name) - validator.check_tensor_dtype_valid('excluded_list', excluded_list_dtype, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('excluded_numbers', excluded_numbers_dtype, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('need_refresh_flag', need_refresh_flag_dtype, [mstype.int32], - self.name) - - return mstype.float32 - - -class MDIterationLeapFrogWithRF(PrimitiveWithInfer): - """ - One step of classical leap frog algorithm to solve the finite difference - Hamiltonian equations of motion for certain system, using Langevin dynamics - with Liu's thermostat scheme. Assume the number of atoms is N and the target - control temperature is T. - - Detailed iteration formula can be found in this paper: A unified thermostat - scheme for efficient configurational sampling for classical/quantum canonical - ensembles via molecular dynamics. DOI: 10.1063/1.4991621. - - Inputs: - - **float4_numbers** (int32) - total length to store random numbers. - - **atom_numbers** (int32) - the number of atoms N. - - **dt** (float32) - time step for finite difference. - - **half_dt** (float32) - half of time step for finite difference. - - **exp_gamma** (float32) - parameter in Liu's dynamic, equals - exp(-gamma_ln * dt), where gamma_ln is the firction factor in Langvin - dynamics. - - **max_velocity** (float32) - the upper limit of velocity, when the - veclocity overflows, scale it to the upper limit. - - **is_max_velocity** (int32) - whether the max velocity control is - open or not. - - - **mass_inverse** (Tensor, float32) - [N,], the inverse value of - mass of each atom. - - **sqrt_mass** (Tensor, float32) - [N,], the inverse square root value - of effect mass in Liu's dynamics of each atom. - - **vel** (Tensor, float32) - [N, 3], the velocity of each atom. - - **crd** (Tensor, float32) - [N, 3], the coordinate of each atom. - - **frc** (Tensor, float32) - [N, 3], the force felt by each atom. - - **acc** (Tensor, float32) - [N, 3], the acceleration of each atom. - - **random force** (Tensor, float32) - [N, 3], the random forces. - - Outputs: - - **res** (float32) - - Supported Platforms: - ``GPU`` - Examples: - """ - - @prim_attr_register - def __init__(self, float4_numbers, atom_numbers, half_dt, dt, exp_gamma, is_max_velocity, max_velocity): - validator.check_value_type('float4_numbers', float4_numbers, (int), self.name) - validator.check_value_type('atom_numbers', atom_numbers, (int), self.name) - validator.check_value_type('half_dt', half_dt, (float), self.name) - validator.check_value_type('dt', dt, (float), self.name) - validator.check_value_type('exp_gamma', exp_gamma, (float), self.name) - validator.check_value_type('is_max_velocity', is_max_velocity, (int), self.name) - validator.check_value_type('max_velocity', max_velocity, (float), self.name) - self.float4_numbers = float4_numbers - self.atom_numbers = atom_numbers - self.half_dt = half_dt - self.dt = dt - self.exp_gamma = exp_gamma - self.is_max_velocity = is_max_velocity - self.max_velocity = max_velocity - - self.init_prim_io_names( - inputs=['mass_inverse', 'sqrt_mass', 'vel_in', 'crd_in', 'frc_in', 'acc_in', 'random_force'], - outputs=['res']) - self.add_prim_attr('float4_numbers', self.float4_numbers) - self.add_prim_attr('atom_numbers', self.atom_numbers) - self.add_prim_attr('half_dt', self.half_dt) - self.add_prim_attr('dt', self.dt) - self.add_prim_attr('exp_gamma', self.exp_gamma) - self.add_prim_attr('is_max_velocity', self.is_max_velocity) - self.add_prim_attr('max_velocity', self.max_velocity) - - def infer_shape(self, mass_inverse_shape, sqrt_mass_shape, vel_in_shape, crd_in_shape, frc_in_shape, acc_in_shape, - random_force_shape): - N = self.atom_numbers - validator.check_int(len(mass_inverse_shape), 1, Rel.EQ, "mass_inverse", self.name) - validator.check_int(len(sqrt_mass_shape), 1, Rel.EQ, "mass_inverse", self.name) - validator.check_int(mass_inverse_shape[0], N, Rel.EQ, "mass_inverse", self.name) - validator.check_int(sqrt_mass_shape[0], N, Rel.EQ, "mass_inverse", self.name) - validator.check_int(vel_in_shape[0], N, Rel.EQ, "vel_in", self.name) - validator.check_int(vel_in_shape[1], 3, Rel.EQ, "vel_in", self.name) - validator.check_int(crd_in_shape[0], N, Rel.EQ, "crd_in", self.name) - validator.check_int(crd_in_shape[1], 3, Rel.EQ, "crd_in", self.name) - validator.check_int(frc_in_shape[0], N, Rel.EQ, "frc_in", self.name) - validator.check_int(frc_in_shape[1], 3, Rel.EQ, "frc_in", self.name) - validator.check_int(acc_in_shape[0], N, Rel.EQ, "acc_in", self.name) - validator.check_int(acc_in_shape[1], 3, Rel.EQ, "acc_in", self.name) - validator.check_int(random_force_shape[0], N, Rel.EQ, "random_force", self.name) - validator.check_int(random_force_shape[1], 3, Rel.EQ, "random_force", self.name) - - return [1,] - - def infer_dtype(self, mass_inverse_dtype, sqrt_mass_dtype, vel_in_dtype, crd_in_dtype, frc_in_dtype, acc_in_dtype, - rf_dtype): - validator.check_tensor_dtype_valid('mass_inverse', mass_inverse_dtype, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('sqrt_mass', sqrt_mass_dtype, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('vel_in', vel_in_dtype, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('crd_in', crd_in_dtype, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('frc_in', frc_in_dtype, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('acc_in', acc_in_dtype, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('rf', rf_dtype, [mstype.float32], self.name) - return mstype.float32 +# Copyright 2021 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================ + +"""Operators for sponge.""" + +import math + +from ..primitive import PrimitiveWithInfer, prim_attr_register +from ..._checkparam import Rel +from ..._checkparam import Validator as validator +from ...common import dtype as mstype + + +class BondForce(PrimitiveWithInfer): + """ + Calculate the force exerted by the simple harmonic bond on the corresponding atoms. + Assume the number of harmonic bonds is M and the number of atoms is N. + + .. math:: + + 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 + + Args: + atom_numbers(int32): the number of atoms N. + bond_numbers(int32): the number of harmonic bonds M. + + Inputs: + - **uint_crd_f** (Tensor, uint32 ) - [N, 3], the unsigned int coordinate value of each atom. + - **scaler_f** (Tensor, float32) - [3,], the 3-D scale factor (x, y, z), + between the real space float coordinates and the unsigned int coordinates. + - **atom_a** (Tensor, int32) - [M,], the first atom index of each bond. + - **atom_b** (Tensor, int32) - [M,], the second atom index of each bond. + - **bond_k** (Tensor, float32) - [M,], the force constant of each bond. + - **bond_r0** (Tensor, float32) - [M,], the equlibrium length of each bond. + + Outputs: + - **frc_f** (float32 Tensor) - [N, 3], the force felt by each atom. + + Supported Platforms: + ``GPU`` + """ + + @prim_attr_register + def __init__(self, bond_numbers, atom_numbers): + validator.check_value_type('bond_numbers', bond_numbers, (int), self.name) + validator.check_value_type('atom_numbers', atom_numbers, (int), self.name) + self.bond_numbers = bond_numbers + self.atom_numbers = atom_numbers + self.add_prim_attr('bond_numbers', self.bond_numbers) + self.add_prim_attr('atom_numbers', self.atom_numbers) + self.init_prim_io_names(inputs=['uint_crd_f', 'scaler_f', 'atom_a', 'atom_b', 'bond_k', 'bond_r0'], + outputs=['frc_f']) + + def infer_shape(self, uint_crd_f_shape, scaler_f_shape, atom_a_shape, atom_b_shape, bond_k_shape, bond_r0_shape): + cls_name = self.name + N = self.atom_numbers + M = self.bond_numbers + validator.check_int(len(uint_crd_f_shape), 2, Rel.EQ, "uint_crd_f_dim", cls_name) + validator.check_int(len(scaler_f_shape), 1, Rel.EQ, "scaler_f_dim", cls_name) + validator.check_int(len(atom_a_shape), 1, Rel.EQ, "atom_a_dim", cls_name) + validator.check_int(len(atom_b_shape), 1, Rel.EQ, "atom_b_dim", cls_name) + validator.check_int(len(bond_k_shape), 1, Rel.EQ, "bond_k_dim", cls_name) + validator.check_int(len(bond_r0_shape), 1, Rel.EQ, "bond_r0_dim", cls_name) + validator.check_int(uint_crd_f_shape[0], N, Rel.EQ, "uint_crd_f_shape[0]", cls_name) + validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f_shape[1]", cls_name) + validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f_shape", cls_name) + validator.check_int(atom_a_shape[0], M, Rel.EQ, "uint_crd_f_shape", cls_name) + validator.check_int(atom_b_shape[0], M, Rel.EQ, "atom_b_shape", cls_name) + validator.check_int(bond_k_shape[0], M, Rel.EQ, "bond_k_shape", cls_name) + validator.check_int(bond_r0_shape[0], M, Rel.EQ, "bond_r0_shape", cls_name) + return uint_crd_f_shape + + def infer_dtype(self, uint_crd_f_dtype, scaler_f_type, atom_a_type, atom_b_type, bond_k_type, bond_r0_type): + validator.check_tensor_dtype_valid('uint_crd_f', uint_crd_f_dtype, [mstype.uint32], self.name) + validator.check_tensor_dtype_valid('scaler_f', scaler_f_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('atom_a', atom_a_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('atom_b', atom_b_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('bond_k', bond_k_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('bond_r0', bond_r0_type, [mstype.float32], self.name) + return bond_r0_type + + +class BondEnergy(PrimitiveWithInfer): + """ + Calculate the harmonic potential energy between each bonded atom pair. + Assume our system has N atoms and M harmonic bonds. + + .. math:: + + dr = (x_1-x_2, y_1-y_2, z_1-z_2) + + .. math:: + + E = k*(|dr| - r_0)^2 + + Args: + atom_numbers(int32): the number of atoms N. + bond_numbers(int32): the number of harmonic bonds M. + + Inputs: + - **uint_crd_f** (Tensor, uint32 ) - [N, 3], the unsigned int coordinate value of each atom. + - **scaler_f** (Tensor, float32) - [3,], the 3-D scale factor (x, y, z), + between the real space float coordinates and the unsigned int coordinates. + - **atom_a** (Tensor, int32) - [M,], the first atom index of each bond. + - **atom_b** (Tensor, int32) - [M,], the second atom index of each bond. + - **bond_k** (Tensor, float32) - [M,], the force constant of each bond. + - **bond_r0** (Tensor, float32) - [M,], the equlibrium length of each bond. + + Outputs: + - **bond_ene** (Tensor, float32) - [M,], the harmonic potential energy for each bond. + + Supported Platforms: + ``GPU`` + """ + + @prim_attr_register + def __init__(self, bond_numbers, atom_numbers): + validator.check_value_type('bond_numbers', bond_numbers, (int), self.name) + validator.check_value_type('atom_numbers', atom_numbers, (int), self.name) + self.bond_numbers = bond_numbers + self.atom_numbers = atom_numbers + self.add_prim_attr('bond_numbers', self.bond_numbers) + self.add_prim_attr('atom_numbers', self.atom_numbers) + self.init_prim_io_names(inputs=['uint_crd_f', 'scaler_f', 'atom_a', 'atom_b', 'bond_k', 'bond_r0'], + outputs=['bond_ene']) + + def infer_shape(self, uint_crd_f_shape, scaler_f_shape, atom_a_shape, atom_b_shape, bond_k_shape, bond_r0_shape): + cls_name = self.name + N = self.atom_numbers + M = self.bond_numbers + validator.check_int(len(uint_crd_f_shape), 2, Rel.EQ, "uint_crd_f_dim", cls_name) + validator.check_int(len(scaler_f_shape), 1, Rel.EQ, "scaler_f_dim", cls_name) + validator.check_int(len(atom_a_shape), 1, Rel.EQ, "atom_a_dim", cls_name) + validator.check_int(len(atom_b_shape), 1, Rel.EQ, "atom_b_dim", cls_name) + validator.check_int(len(bond_k_shape), 1, Rel.EQ, "bond_k_dim", cls_name) + validator.check_int(len(bond_r0_shape), 1, Rel.EQ, "bond_r0_dim", cls_name) + validator.check_int(uint_crd_f_shape[0], N, Rel.EQ, "uint_crd_f_shape[0]", cls_name) + validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f_shape[1]", cls_name) + validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f_shape", cls_name) + validator.check_int(atom_a_shape[0], M, Rel.EQ, "uint_crd_f_shape", cls_name) + validator.check_int(atom_b_shape[0], M, Rel.EQ, "atom_b_shape", cls_name) + validator.check_int(bond_k_shape[0], M, Rel.EQ, "bond_k_shape", cls_name) + validator.check_int(bond_r0_shape[0], M, Rel.EQ, "bond_r0_shape", cls_name) + + return bond_k_shape + + def infer_dtype(self, uint_crd_f_dtype, scaler_f_type, atom_a_type, atom_b_type, bond_k_type, bond_r0_type): + validator.check_tensor_dtype_valid('uint_crd_f', uint_crd_f_dtype, [mstype.uint32], self.name) + validator.check_tensor_dtype_valid('scaler_f', scaler_f_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('atom_a', atom_a_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('atom_b', atom_b_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('bond_k', bond_k_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('bond_r0', bond_r0_type, [mstype.float32], self.name) + return bond_r0_type + + +class BondAtomEnergy(PrimitiveWithInfer): + """ + Add the potential energy caused by simple harmonic bonds to the total + potential energy of each atom. + + The calculation formula is the same as operator BondEnergy(). + + Args: + atom_numbers(int32): the number of atoms N. + bond_numbers(int32): the number of harmonic bonds M. + + Inputs: + - **uint_crd_f** (Tensor, uint32 ) - [N, 3], the unsigned int coordinate value of each atom. + - **scaler_f** (Tensor, float32) - [3,], the 3-D scale factor (x, y, z), + between the real space float coordinates and the unsigned int coordinates. + - **atom_a** (Tensor, int32) - [M,], the first atom index of each bond. + - **atom_b** (Tensor, int32) - [M,], the second atom index of each bond. + - **bond_k** (Tensor, float32) - [M,], the force constant of each bond. + - **bond_r0** (Tensor, float32) - [M,], the equlibrium length of each bond. + + Outputs: + - **atom_ene** (Tensor, float32) - [N,], the accumulated potential energy for each atom. + + Supported Platforms: + ``GPU`` + """ + + @prim_attr_register + def __init__(self, bond_numbers, atom_numbers): + validator.check_value_type('bond_numbers', bond_numbers, (int), self.name) + validator.check_value_type('atom_numbers', atom_numbers, (int), self.name) + self.bond_numbers = bond_numbers + self.atom_numbers = atom_numbers + self.add_prim_attr('bond_numbers', self.bond_numbers) + self.add_prim_attr('atom_numbers', self.atom_numbers) + self.init_prim_io_names(inputs=['uint_crd_f', 'scaler_f', 'atom_a', 'atom_b', 'bond_k', 'bond_r0'], + outputs=['atom_ene']) + + def infer_shape(self, uint_crd_f_shape, scaler_f_shape, atom_a_shape, atom_b_shape, bond_k_shape, bond_r0_shape): + cls_name = self.name + N = self.atom_numbers + M = self.bond_numbers + validator.check_int(len(uint_crd_f_shape), 2, Rel.EQ, "uint_crd_f_dim", cls_name) + validator.check_int(len(scaler_f_shape), 1, Rel.EQ, "scaler_f_dim", cls_name) + validator.check_int(len(atom_a_shape), 1, Rel.EQ, "atom_a_dim", cls_name) + validator.check_int(len(atom_b_shape), 1, Rel.EQ, "atom_b_dim", cls_name) + validator.check_int(len(bond_k_shape), 1, Rel.EQ, "bond_k_dim", cls_name) + validator.check_int(len(bond_r0_shape), 1, Rel.EQ, "bond_r0_dim", cls_name) + validator.check_int(uint_crd_f_shape[0], N, Rel.EQ, "uint_crd_f_shape[0]", cls_name) + validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f_shape[1]", cls_name) + validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f_shape", cls_name) + validator.check_int(atom_a_shape[0], M, Rel.EQ, "uint_crd_f_shape", cls_name) + validator.check_int(atom_b_shape[0], M, Rel.EQ, "atom_b_shape", cls_name) + validator.check_int(bond_k_shape[0], M, Rel.EQ, "bond_k_shape", cls_name) + validator.check_int(bond_r0_shape[0], M, Rel.EQ, "bond_r0_shape", cls_name) + return [N,] + + def infer_dtype(self, uint_crd_f_dtype, scaler_f_type, atom_a_type, atom_b_type, bond_k_type, bond_r0_type): + validator.check_tensor_dtype_valid('uint_crd_f', uint_crd_f_dtype, [mstype.uint32], self.name) + validator.check_tensor_dtype_valid('scaler_f', scaler_f_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('atom_a', atom_a_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('atom_b', atom_b_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('bond_k', bond_k_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('bond_r0', bond_r0_type, [mstype.float32], self.name) + return bond_r0_type + + +class BondForceWithAtomEnergy(PrimitiveWithInfer): + """ + Calculate bond force and harmonic potential energy together. + + The calculation formula is the same as operator BondForce() and BondEnergy(). + + Args: + atom_numbers(int32): the number of atoms N. + bond_numbers(int32): the number of harmonic bonds M. + + Inputs: + - **uint_crd_f** (Tensor, uint32 ) - [N, 3], the unsigned int coordinate value of each atom. + - **scaler_f** (Tensor, float32) - [3,], the 3-D scale factor (x, y, z), + between the real space float coordinates and the unsigned int coordinates. + - **atom_a** (Tensor, int32) - [M,], the first atom index of each bond. + - **atom_b** (Tensor, int32) - [M,], the second atom index of each bond. + - **bond_k** (Tensor, float32) - [M,], the force constant of each bond. + - **bond_r0** (Tensor, float32) - [M,], the equlibrium length of each bond. + + Outputs: + - **frc_f** (Tensor, float32) - [N, 3], same as operator BondForce(). + - **atom_e** (Tensor, float32) - [N,], same as atom_ene in operator BondAtomEnergy(). + + Supported Platforms: + ``GPU`` + """ + + @prim_attr_register + def __init__(self, bond_numbers, atom_numbers): + validator.check_value_type('bond_numbers', bond_numbers, (int), self.name) + validator.check_value_type('atom_numbers', atom_numbers, (int), self.name) + self.bond_numbers = bond_numbers + self.atom_numbers = atom_numbers + self.add_prim_attr('bond_numbers', self.bond_numbers) + self.add_prim_attr('atom_numbers', self.atom_numbers) + self.init_prim_io_names(inputs=['uint_crd_f', 'scaler_f', 'atom_a', 'atom_b', 'bond_k', 'bond_r0'], + outputs=['frc_f', 'atom_e']) + + def infer_shape(self, uint_crd_f_shape, scaler_f_shape, atom_a_shape, atom_b_shape, bond_k_shape, bond_r0_shape): + cls_name = self.name + N = self.atom_numbers + M = self.bond_numbers + validator.check_int(len(uint_crd_f_shape), 2, Rel.EQ, "uint_crd_f_dim", cls_name) + validator.check_int(len(scaler_f_shape), 1, Rel.EQ, "scaler_f_dim", cls_name) + validator.check_int(len(atom_a_shape), 1, Rel.EQ, "atom_a_dim", cls_name) + validator.check_int(len(atom_b_shape), 1, Rel.EQ, "atom_b_dim", cls_name) + validator.check_int(len(bond_k_shape), 1, Rel.EQ, "bond_k_dim", cls_name) + validator.check_int(len(bond_r0_shape), 1, Rel.EQ, "bond_r0_dim", cls_name) + validator.check_int(uint_crd_f_shape[0], N, Rel.EQ, "uint_crd_f_shape[0]", cls_name) + validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f_shape[1]", cls_name) + validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f_shape", cls_name) + validator.check_int(atom_a_shape[0], M, Rel.EQ, "uint_crd_f_shape", cls_name) + validator.check_int(atom_b_shape[0], M, Rel.EQ, "atom_b_shape", cls_name) + validator.check_int(bond_k_shape[0], M, Rel.EQ, "bond_k_shape", cls_name) + validator.check_int(bond_r0_shape[0], M, Rel.EQ, "bond_r0_shape", cls_name) + + return uint_crd_f_shape, [N,] + + def infer_dtype(self, uint_crd_f_dtype, scaler_f_type, atom_a_type, atom_b_type, bond_k_type, bond_r0_type): + validator.check_tensor_dtype_valid('uint_crd_f', uint_crd_f_dtype, [mstype.uint32], self.name) + validator.check_tensor_dtype_valid('scaler_f', scaler_f_type, [mstype.float32], self.name) + + validator.check_tensor_dtype_valid('atom_a', atom_a_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('atom_b', atom_b_type, [mstype.int32], self.name) + + validator.check_tensor_dtype_valid('bond_k', bond_k_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('bond_r0', bond_r0_type, [mstype.float32], self.name) + return bond_r0_type, bond_r0_type + + +class BondForceWithAtomVirial(PrimitiveWithInfer): + """ + Calculate bond force and the virial coefficient caused by simple harmonic + bond for each atom together. + + The calculation formula of the force part is the same as operator BondForce(). + The Virial part is as follows: + + .. math:: + + dr = (x_1-x_2, y_1-y_2, z_1-z_2) + + .. math:: + + virial = |dr|*(|dr| - r_0)*k + + Args: + atom_numbers(int32): the number of atoms N. + bond_numbers(int32): the number of harmonic bonds M. + + Inputs: + - **uint_crd_f** (Tensor, uint32 ) - [N, 3], the unsigned int coordinate value of each atom. + - **scaler_f** (Tensor, float32) - [3,], the 3-D scale factor (x, y, z), + between the real space float coordinates and the unsigned int coordinates. + - **atom_a** (Tensor, int32) - [M,], the first atom index of each bond. + - **atom_b** (Tensor, int32) - [M,], the second atom index of each bond. + - **bond_k** (Tensor, float32) - [M,], the force constant of each bond. + - **bond_r0** (Tensor, float32) - [M,], the equlibrium length of each bond. + + Outputs: + - **frc_f** (Tensor, float32) - [N, 3], same as operator BondForce(). + - **atom_v** (Tensor, float32) - [N,], the accumulated virial coefficient for each atom. + + Supported Platforms: + ``GPU`` + """ + + @prim_attr_register + def __init__(self, bond_numbers, atom_numbers): + validator.check_value_type('bond_numbers', bond_numbers, (int), self.name) + validator.check_value_type('atom_numbers', atom_numbers, (int), self.name) + self.bond_numbers = bond_numbers + self.atom_numbers = atom_numbers + self.add_prim_attr('bond_numbers', self.bond_numbers) + self.add_prim_attr('atom_numbers', self.atom_numbers) + self.init_prim_io_names(inputs=['uint_crd_f', 'scaler_f', 'atom_a', 'atom_b', 'bond_k', 'bond_r0'], + outputs=['frc_f', 'atom_v']) + + def infer_shape(self, uint_crd_f_shape, scaler_f_shape, atom_a_shape, atom_b_shape, bond_k_shape, bond_r0_shape): + cls_name = self.name + N = self.atom_numbers + M = self.bond_numbers + validator.check_int(len(uint_crd_f_shape), 2, Rel.EQ, "uint_crd_f_dim", cls_name) + validator.check_int(len(scaler_f_shape), 1, Rel.EQ, "scaler_f_dim", cls_name) + validator.check_int(len(atom_a_shape), 1, Rel.EQ, "atom_a_dim", cls_name) + validator.check_int(len(atom_b_shape), 1, Rel.EQ, "atom_b_dim", cls_name) + validator.check_int(len(bond_k_shape), 1, Rel.EQ, "bond_k_dim", cls_name) + validator.check_int(len(bond_r0_shape), 1, Rel.EQ, "bond_r0_dim", cls_name) + validator.check_int(uint_crd_f_shape[0], N, Rel.EQ, "uint_crd_f_shape[0]", cls_name) + validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f_shape[1]", cls_name) + validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f_shape", cls_name) + validator.check_int(atom_a_shape[0], M, Rel.EQ, "uint_crd_f_shape", cls_name) + validator.check_int(atom_b_shape[0], M, Rel.EQ, "atom_b_shape", cls_name) + validator.check_int(bond_k_shape[0], M, Rel.EQ, "bond_k_shape", cls_name) + validator.check_int(bond_r0_shape[0], M, Rel.EQ, "bond_r0_shape", cls_name) + + return uint_crd_f_shape, [N,] + + def infer_dtype(self, uint_crd_f_dtype, scaler_f_type, atom_a_type, atom_b_type, bond_k_type, bond_r0_type): + validator.check_tensor_dtype_valid('uint_crd_f', uint_crd_f_dtype, [mstype.uint32], self.name) + validator.check_tensor_dtype_valid('scaler_f', scaler_f_type, [mstype.float32], self.name) + + validator.check_tensor_dtype_valid('atom_a', atom_a_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('atom_b', atom_b_type, [mstype.int32], self.name) + + validator.check_tensor_dtype_valid('bond_k', bond_k_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('bond_r0', bond_r0_type, [mstype.float32], self.name) + return bond_r0_type, bond_r0_type + + +class DihedralForce(PrimitiveWithInfer): + """ + Calculate the force exerted by the dihedral term which made of 4-atoms + on the corresponding atoms. Assume the number of dihedral terms is M and + the number of atoms is N. + + .. math:: + + 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) + .. math:: + dr_{cd} = (x_d-x_c, y_d-y_c, z_d-z_c) + .. math:: + r1 = dr_{ab}*dr_{cb} + .. math:: + r2 = dr_{cd}*dr_{cb} + .. math:: + phi = pi - sign(inner_product(r1*r2), dr_{cb}) + * 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) + .. math:: + dphidr1 = r2/|r1|/|r2| + cos(phi)/|r1|^2*r1 + .. math:: + dphidr2 = r1/|r1|/|r2| + cos(phi)/|r2|^2*r2 + .. math:: + dEdra = dEdphi * dr_{cb} * dphidr1 + .. math:: + dEdrd = dEdphi * dphi_dr2 * dr_{cb} + .. math:: + dEdrjpart = dEdphi * ((dr_{ab} * dphidr1) + (dr_{cd} * dphidr2)) + .. math:: + F_a = dEdri + .. math:: + F_b = dEdrjpart - dEdri + .. math:: + F_c = - dEdrl - dEdrjpart + .. math:: + F_d = dEdrl + + Args: + dihedral_numbers(int32): the number of dihedral terms M. + + Inputs: + - **dihedral_numbers** (int32) - the number of dihedral terms M. + - **uint_crd_f** (Tensor, uint32) - [N, 3], the unsigned int coordinates + value of each atom. + - **scaler_f** (Tensor, float32) - [3,], the 3-D scale factor between + the real space float coordinates and the unsigned int coordinates. + - **atom_a** (Tensor, int32) - [M,], the 1st atom index of each dihedral. + - **atom_b** (Tensor, int32) - [M,], the 2nd atom index of each dihedral. + - **atom_c** (Tensor, int32) - [M,], the 3rd atom index of each dihedral. + - **atom_d** (Tensor, int32) - [M,], the 4th atom index of each dihedral. + 4 atoms are connected in the form a-b-c-d. + - **ipn** (Tensor, int32) - [M,], the period of dihedral angle of each dihedral. + - **pk** (Tensor, float32) - [M,], the force constant of each dihedral. + - **gamc** (Tensor, float32) - [M,], k*cos(phi_0) of each dihedral. + - **gams** (Tensor, float32) - [M,], k*sin(phi_0) of each dihedral. + - **pn** (Tensor, float32) - [M,], the floating point form of ipn. + + Outputs: + - **frc_f** (Tensor, float32) - [N, 3], the force felt by each atom. + + Supported Platforms: + ``GPU`` + """ + + @prim_attr_register + def __init__(self, dihedral_numbers): + validator.check_value_type('dihedral_numbers', dihedral_numbers, (int), self.name) + 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', + 'gamc', 'gams', 'pn'], + outputs=['frc_f']) + self.add_prim_attr('dihedral_numbers', self.dihedral_numbers) + + 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): + cls_name = self.name + M = self.dihedral_numbers + validator.check_int(len(uint_crd_f_shape), 2, Rel.EQ, "uint_crd_f_dim", cls_name) + validator.check_int(len(scaler_f_shape), 1, Rel.EQ, "scaler_f_dim", cls_name) + validator.check_int(len(atom_a_shape), 1, Rel.EQ, "atom_a_dim", cls_name) + validator.check_int(len(atom_b_shape), 1, Rel.EQ, "atom_b_dim", cls_name) + validator.check_int(len(atom_c_shape), 1, Rel.EQ, "atom_c_dim", cls_name) + validator.check_int(len(atom_d_shape), 1, Rel.EQ, "atom_d_dim", cls_name) + validator.check_int(len(ipn_shape), 1, Rel.EQ, "ipn_dim", cls_name) + validator.check_int(len(pk_shape), 1, Rel.EQ, "pk_dim", cls_name) + validator.check_int(len(gamc_shape), 1, Rel.EQ, "gamc_dim", cls_name) + validator.check_int(len(gams_shape), 1, Rel.EQ, "gams_dim", cls_name) + validator.check_int(len(pn_shape), 1, Rel.EQ, "pn_dim", cls_name) + + validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f_shape[1]", cls_name) + validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f_shape", cls_name) + validator.check_int(atom_a_shape[0], M, Rel.EQ, "atom_a_shape", cls_name) + validator.check_int(atom_b_shape[0], M, Rel.EQ, "atom_b_shape", cls_name) + validator.check_int(atom_c_shape[0], M, Rel.EQ, "atom_c_shape", cls_name) + validator.check_int(atom_d_shape[0], M, Rel.EQ, "atom_d_shape", cls_name) + validator.check_int(ipn_shape[0], M, Rel.EQ, "ipn_shape", cls_name) + validator.check_int(pk_shape[0], M, Rel.EQ, "pk_shape", cls_name) + validator.check_int(gamc_shape[0], M, Rel.EQ, "gamc_shape", cls_name) + validator.check_int(gams_shape[0], M, Rel.EQ, "gams_shape", cls_name) + validator.check_int(pn_shape[0], M, Rel.EQ, "pn_shape", cls_name) + return uint_crd_f_shape + + def infer_dtype(self, uint_crd_f_dtype, scaler_f_type, atom_a_type, atom_b_type, atom_c_type, atom_d_type, + ipn_type, pk_type, gamc_type, gams_type, pn_type): + validator.check_tensor_dtype_valid('uint_crd_f', uint_crd_f_dtype, [mstype.uint32], self.name) + validator.check_tensor_dtype_valid('scaler_f', scaler_f_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('atom_a', atom_a_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('atom_b', atom_b_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('atom_c', atom_c_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('atom_d', atom_d_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('ipn', ipn_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('pk', pk_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('gamc', gamc_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('gams', gams_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('pn', pn_type, [mstype.float32], self.name) + + return pn_type + + +class DihedralEnergy(PrimitiveWithInfer): + """ + Calculate the potential energy caused by dihedral terms for each 4-atom pair. + Assume our system has N atoms and M dihedral terms. + + .. math:: + + E = k(1 + cos(n*phi - phi_0)) + + Args: + dihedral_numbers(int32): the number of dihedral terms M. + + Inputs: + - **dihedral_numbers** (int32) - the number of dihedral terms M. + - **uint_crd_f** (Tensor, uint32) - [N, 3], the unsigned int coordinates + value of each atom. + - **scaler_f** (Tensor, float32) - [3,], the 3-D scale factor between + the real space float coordinates and the unsigned int coordinates. + - **atom_a** (Tensor, int32) - [M,], the 1st atom index of each dihedral. + - **atom_b** (Tensor, int32) - [M,], the 2nd atom index of each dihedral. + - **atom_c** (Tensor, int32) - [M,], the 3rd atom index of each dihedral. + - **atom_d** (Tensor, int32) - [M,], the 4th atom index of each dihedral. + 4 atoms are connected in the form a-b-c-d. + - **ipn** (Tensor, int32) - [M,], the period of dihedral angle of each dihedral. + - **pk** (Tensor, float32) - [M,], the force constant of each dihedral. + - **gamc** (Tensor, float32) - [M,], k*cos(phi_0) of each dihedral. + - **gams** (Tensor, float32) - [M,], k*sin(phi_0) of each dihedral. + - **pn** (Tensor, float32) - [M,], the floating point form of ipn. + + Outputs: + - **ene** (Tensor, float32) - [M,], the potential energy for each + dihedral term. + + Supported Platforms: + ``GPU`` + """ + + @prim_attr_register + def __init__(self, dihedral_numbers): + validator.check_value_type('dihedral_numbers', dihedral_numbers, (int), self.name) + 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', + 'gamc', 'gams', 'pn'], + outputs=['ene']) + self.add_prim_attr('dihedral_numbers', self.dihedral_numbers) + + 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): + cls_name = self.name + M = self.dihedral_numbers + validator.check_int(len(uint_crd_f_shape), 2, Rel.EQ, "uint_crd_f_dim", cls_name) + validator.check_int(len(scaler_f_shape), 1, Rel.EQ, "scaler_f_dim", cls_name) + validator.check_int(len(atom_a_shape), 1, Rel.EQ, "atom_a_dim", cls_name) + validator.check_int(len(atom_b_shape), 1, Rel.EQ, "atom_b_dim", cls_name) + validator.check_int(len(atom_c_shape), 1, Rel.EQ, "atom_c_dim", cls_name) + validator.check_int(len(atom_d_shape), 1, Rel.EQ, "atom_d_dim", cls_name) + validator.check_int(len(ipn_shape), 1, Rel.EQ, "ipn_dim", cls_name) + validator.check_int(len(pk_shape), 1, Rel.EQ, "pk_dim", cls_name) + validator.check_int(len(gamc_shape), 1, Rel.EQ, "gamc_dim", cls_name) + validator.check_int(len(gams_shape), 1, Rel.EQ, "gams_dim", cls_name) + validator.check_int(len(pn_shape), 1, Rel.EQ, "pn_dim", cls_name) + + validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f_shape[1]", cls_name) + validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f_shape", cls_name) + validator.check_int(atom_a_shape[0], M, Rel.EQ, "atom_a_shape", cls_name) + validator.check_int(atom_b_shape[0], M, Rel.EQ, "atom_b_shape", cls_name) + validator.check_int(atom_c_shape[0], M, Rel.EQ, "atom_c_shape", cls_name) + validator.check_int(atom_d_shape[0], M, Rel.EQ, "atom_d_shape", cls_name) + validator.check_int(ipn_shape[0], M, Rel.EQ, "ipn_shape", cls_name) + validator.check_int(pk_shape[0], M, Rel.EQ, "pk_shape", cls_name) + validator.check_int(gamc_shape[0], M, Rel.EQ, "gamc_shape", cls_name) + validator.check_int(gams_shape[0], M, Rel.EQ, "gams_shape", cls_name) + validator.check_int(pn_shape[0], M, Rel.EQ, "pn_shape", cls_name) + return [M,] + + def infer_dtype(self, uint_crd_f_dtype, scaler_f_type, atom_a_type, atom_b_type, atom_c_type, atom_d_type, + ipn_type, pk_type, gamc_type, gams_type, pn_type): + validator.check_tensor_dtype_valid('uint_crd_f', uint_crd_f_dtype, [mstype.uint32], self.name) + validator.check_tensor_dtype_valid('scaler_f', scaler_f_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('atom_a', atom_a_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('atom_b', atom_b_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('atom_c', atom_c_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('atom_d', atom_d_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('ipn', ipn_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('pk', pk_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('gamc', gamc_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('gams', gams_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('pn', pn_type, [mstype.float32], self.name) + + return pn_type + + +class DihedralAtomEnergy(PrimitiveWithInfer): + """ + Add the potential energy caused by dihedral terms to the total potential + energy of each atom. + + The calculation formula is the same as operator DihedralEnergy(). + + Args: + dihedral_numbers(int32): the number of dihedral terms M. + + Inputs: + - **dihedral_numbers** (int32) - the number of dihedral terms M. + - **uint_crd_f** (Tensor, uint32) - [N, 3], the unsigned int coordinates + value of each atom. + - **scaler_f** (Tensor, float32) - [3,], the 3-D scale factor between + the real space float coordinates and the unsigned int coordinates. + - **atom_a** (Tensor, int32) - [M,], the 1st atom index of each dihedral. + - **atom_b** (Tensor, int32) - [M,], the 2nd atom index of each dihedral. + - **atom_c** (Tensor, int32) - [M,], the 3rd atom index of each dihedral. + - **atom_d** (Tensor, int32) - [M,], the 4th atom index of each dihedral. + 4 atoms are connected in the form a-b-c-d. + - **ipn** (Tensor, int32) - [M,], the period of dihedral angle of each dihedral. + - **pk** (Tensor, float32) - [M,], the force constant of each dihedral. + - **gamc** (Tensor, float32) - [M,], k*cos(phi_0) of each dihedral. + - **gams** (Tensor, float32) - [M,], k*sin(phi_0) of each dihedral. + - **pn** (Tensor, float32) - [M,], the floating point form of ipn. + + Outputs: + - **ene** (Tensor, float32) - [N,], the accumulated potential + energy for each atom. + + Supported Platforms: + ``GPU`` + """ + + @prim_attr_register + def __init__(self, dihedral_numbers): + validator.check_value_type('dihedral_numbers', dihedral_numbers, (int), self.name) + 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', + 'gamc', 'gams', 'pn'], + outputs=['ene']) + self.add_prim_attr('dihedral_numbers', self.dihedral_numbers) + + 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): + cls_name = self.name + N = uint_crd_f_shape[0] + M = self.dihedral_numbers + validator.check_int(len(uint_crd_f_shape), 2, Rel.EQ, "uint_crd_f_dim", cls_name) + validator.check_int(len(scaler_f_shape), 1, Rel.EQ, "scaler_f_dim", cls_name) + validator.check_int(len(atom_a_shape), 1, Rel.EQ, "atom_a_dim", cls_name) + validator.check_int(len(atom_b_shape), 1, Rel.EQ, "atom_b_dim", cls_name) + validator.check_int(len(atom_c_shape), 1, Rel.EQ, "atom_c_dim", cls_name) + validator.check_int(len(atom_d_shape), 1, Rel.EQ, "atom_d_dim", cls_name) + validator.check_int(len(ipn_shape), 1, Rel.EQ, "ipn_dim", cls_name) + validator.check_int(len(pk_shape), 1, Rel.EQ, "pk_dim", cls_name) + validator.check_int(len(gamc_shape), 1, Rel.EQ, "gamc_dim", cls_name) + validator.check_int(len(gams_shape), 1, Rel.EQ, "gams_dim", cls_name) + validator.check_int(len(pn_shape), 1, Rel.EQ, "pn_dim", cls_name) + + validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f_shape[1]", cls_name) + validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f_shape", cls_name) + validator.check_int(atom_a_shape[0], M, Rel.EQ, "atom_a_shape", cls_name) + validator.check_int(atom_b_shape[0], M, Rel.EQ, "atom_b_shape", cls_name) + validator.check_int(atom_c_shape[0], M, Rel.EQ, "atom_c_shape", cls_name) + validator.check_int(atom_d_shape[0], M, Rel.EQ, "atom_d_shape", cls_name) + validator.check_int(ipn_shape[0], M, Rel.EQ, "ipn_shape", cls_name) + validator.check_int(pk_shape[0], M, Rel.EQ, "pk_shape", cls_name) + validator.check_int(gamc_shape[0], M, Rel.EQ, "gamc_shape", cls_name) + validator.check_int(gams_shape[0], M, Rel.EQ, "gams_shape", cls_name) + validator.check_int(pn_shape[0], M, Rel.EQ, "pn_shape", cls_name) + return [N,] + + def infer_dtype(self, uint_crd_f_dtype, scaler_f_type, atom_a_type, atom_b_type, atom_c_type, atom_d_type, + ipn_type, pk_type, gamc_type, gams_type, pn_type): + validator.check_tensor_dtype_valid('uint_crd_f', uint_crd_f_dtype, [mstype.uint32], self.name) + validator.check_tensor_dtype_valid('scaler_f', scaler_f_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('atom_a', atom_a_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('atom_b', atom_b_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('atom_c', atom_c_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('atom_d', atom_d_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('ipn', ipn_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('pk', pk_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('gamc', gamc_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('gams', gams_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('pn', pn_type, [mstype.float32], self.name) + + return pn_type + + +class DihedralForceWithAtomEnergy(PrimitiveWithInfer): + """ + Calculate dihedral force and potential energy together. + + The calculation formula is the same as operator DihedralForce() and DihedralEnergy(). + + Args: + dihedral_numbers(int32): the number of dihedral terms M. + + Inputs: + - **dihedral_numbers** (int32) - the number of dihedral terms M. + - **uint_crd_f** (Tensor, uint32) - [N, 3], the unsigned int coordinates + value of each atom. + - **scaler_f** (Tensor, float32) - [3,], the 3-D scale factor between + the real space float coordinates and the unsigned int coordinates. + - **atom_a** (Tensor, int32) - [M,], the 1st atom index of each dihedral. + - **atom_b** (Tensor, int32) - [M,], the 2nd atom index of each dihedral. + - **atom_c** (Tensor, int32) - [M,], the 3rd atom index of each dihedral. + - **atom_d** (Tensor, int32) - [M,], the 4th atom index of each dihedral. + 4 atoms are connected in the form a-b-c-d. + - **ipn** (Tensor, int32) - [M,], the period of dihedral angle of each dihedral. + - **pk** (Tensor, float32) - [M,], the force constant of each dihedral. + - **gamc** (Tensor, float32) - [M,], k*cos(phi_0) of each dihedral. + - **gams** (Tensor, float32) - [M,], k*sin(phi_0) of each dihedral. + - **pn** (Tensor, float32) - [M,], the floating point form of ipn. + + Outputs: + - **frc_f** (Tensor, float32) - [N, 3], same as operator DihedralForce(). + - **ene** (Tensor, float32) - [N,], same as operator DihedralAtomEnergy(). + + Supported Platforms: + ``GPU`` + """ + + @prim_attr_register + def __init__(self, dihedral_numbers): + validator.check_value_type('dihedral_numbers', dihedral_numbers, (int), self.name) + 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', + 'gamc', 'gams', 'pn'], + outputs=['frc_f', 'ene']) + self.add_prim_attr('dihedral_numbers', self.dihedral_numbers) + + 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): + cls_name = self.name + N = uint_crd_f_shape[0] + M = self.dihedral_numbers + validator.check_int(len(uint_crd_f_shape), 2, Rel.EQ, "uint_crd_f_dim", cls_name) + validator.check_int(len(scaler_f_shape), 1, Rel.EQ, "scaler_f_dim", cls_name) + validator.check_int(len(atom_a_shape), 1, Rel.EQ, "atom_a_dim", cls_name) + validator.check_int(len(atom_b_shape), 1, Rel.EQ, "atom_b_dim", cls_name) + validator.check_int(len(atom_c_shape), 1, Rel.EQ, "atom_c_dim", cls_name) + validator.check_int(len(atom_d_shape), 1, Rel.EQ, "atom_d_dim", cls_name) + validator.check_int(len(ipn_shape), 1, Rel.EQ, "ipn_dim", cls_name) + validator.check_int(len(pk_shape), 1, Rel.EQ, "pk_dim", cls_name) + validator.check_int(len(gamc_shape), 1, Rel.EQ, "gamc_dim", cls_name) + validator.check_int(len(gams_shape), 1, Rel.EQ, "gams_dim", cls_name) + validator.check_int(len(pn_shape), 1, Rel.EQ, "pn_dim", cls_name) + + validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f_shape[1]", cls_name) + validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f_shape", cls_name) + validator.check_int(atom_a_shape[0], M, Rel.EQ, "atom_a_shape", cls_name) + validator.check_int(atom_b_shape[0], M, Rel.EQ, "atom_b_shape", cls_name) + validator.check_int(atom_c_shape[0], M, Rel.EQ, "atom_c_shape", cls_name) + validator.check_int(atom_d_shape[0], M, Rel.EQ, "atom_d_shape", cls_name) + validator.check_int(ipn_shape[0], M, Rel.EQ, "ipn_shape", cls_name) + validator.check_int(pk_shape[0], M, Rel.EQ, "pk_shape", cls_name) + validator.check_int(gamc_shape[0], M, Rel.EQ, "gamc_shape", cls_name) + validator.check_int(gams_shape[0], M, Rel.EQ, "gams_shape", cls_name) + validator.check_int(pn_shape[0], M, Rel.EQ, "pn_shape", cls_name) + return uint_crd_f_shape, [N,] + + def infer_dtype(self, uint_crd_f_dtype, scaler_f_type, atom_a_type, atom_b_type, atom_c_type, atom_d_type, + ipn_type, pk_type, gamc_type, gams_type, pn_type): + validator.check_tensor_dtype_valid('uint_crd_f', uint_crd_f_dtype, [mstype.uint32], self.name) + validator.check_tensor_dtype_valid('scaler_f', scaler_f_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('atom_a', atom_a_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('atom_b', atom_b_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('atom_c', atom_c_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('atom_d', atom_d_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('ipn', ipn_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('pk', pk_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('gamc', gamc_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('gams', gams_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('pn', pn_type, [mstype.float32], self.name) + + return pn_type, pn_type + + +class AngleForce(PrimitiveWithInfer): + """ + Calculate the force exerted by angles made of 3 atoms on the + corresponding atoms. Assume the number of angles is M and the + number of atoms is N. + + .. math:: + 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) + .. math:: + 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} + - 1/|dr_{ab}|/|dr_{cb}|*dr_{cb}] + .. math:: + F_c = -2*k*(theta-theta_0)/sin(theta)*[cos(theta)/|dr_{cb}|^2*dr_{cb} + - 1/|dr_{cb}|/|dr_{ab}|*dr_{ab}] + .. math:: + F_b = -F_a - F_c + + Args: + angle_numbers(int32): the number of angles M. + + Inputs: + - **uint_crd_f** (Tensor, uint32) - [N, 3], the unsigned int coordinate value of each atom. + - **scaler_f** (Tensor, float32) - [3,], the 3-D scale factor between + the real space float coordinates and the unsigned int coordinates. + - **atom_a** (Tensor, int32) - [M,], the 1st atom index of each angle. + - **atom_b** (Tensor, int32) - [M,], the 2nd and the central atom index of each angle. + - **atom_c** (Tensor, int32) - [M,], the 3rd atom index of each angle. + - **angle_k** (Tensor, float32) - [M,], the force constant for each angle. + - **angle_theta0** (Tensor, float32) - [M,], the equilibrium position value for each angle. + + Outputs: + - **frc_f** (Tensor, float32) - [N, 3], the force felt by each atom. + + Supported Platforms: + ``GPU`` + """ + + @prim_attr_register + def __init__(self, angle_numbers): + validator.check_value_type('angle_numbers', angle_numbers, (int), self.name) + self.angle_numbers = angle_numbers + self.init_prim_io_names(inputs=['uint_crd_f', 'scaler_f', 'atom_a', 'atom_b', 'atom_c', 'angle_k', + 'angle_theta0'], + outputs=['frc_f']) + self.add_prim_attr('angle_numbers', self.angle_numbers) + + 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): + cls_name = self.name + M = self.angle_numbers + validator.check_int(len(uint_crd_f_shape), 2, Rel.EQ, "uint_crd_f_dim", cls_name) + validator.check_int(len(scaler_f_shape), 1, Rel.EQ, "scaler_f_dim", cls_name) + validator.check_int(len(atom_a_shape), 1, Rel.EQ, "atom_a_dim", cls_name) + validator.check_int(len(atom_b_shape), 1, Rel.EQ, "atom_b_dim", cls_name) + validator.check_int(len(atom_c_shape), 1, Rel.EQ, "atom_c_dim", cls_name) + validator.check_int(len(angle_k_shape), 1, Rel.EQ, "angle_k_dim", cls_name) + validator.check_int(len(angle_theta0_shape), 1, Rel.EQ, "angle_theta0_dim", cls_name) + + validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f_shape[1]", cls_name) + validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f_shape", cls_name) + validator.check_int(atom_a_shape[0], M, Rel.EQ, "atom_a_shape", cls_name) + validator.check_int(atom_b_shape[0], M, Rel.EQ, "atom_b_shape", cls_name) + validator.check_int(atom_c_shape[0], M, Rel.EQ, "atom_c_shape", cls_name) + validator.check_int(angle_k_shape[0], M, Rel.EQ, "angle_k_shape", cls_name) + validator.check_int(angle_theta0_shape[0], M, Rel.EQ, "angle_theta0_shape", cls_name) + return uint_crd_f_shape + + def infer_dtype(self, uint_crd_f_dtype, scaler_f_type, atom_a_type, atom_b_type, atom_c_type, angle_k_type, + angle_theta0_type): + validator.check_tensor_dtype_valid('uint_crd_f', uint_crd_f_dtype, [mstype.uint32], self.name) + validator.check_tensor_dtype_valid('scaler_f', scaler_f_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('atom_a', atom_a_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('atom_b', atom_b_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('atom_c', atom_c_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('angle_k', angle_k_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('angle_theta0', angle_theta0_type, [mstype.float32], self.name) + return angle_k_type + + +class AngleEnergy(PrimitiveWithInfer): + """ + Calculate the energy caused by 3-atoms angle term. Assume the number of angles is M and the + number of atoms is N. + + .. math:: + 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) + .. math:: + theta = arccos(inner_product(dr_{ab}, dr_{cb})/|dr_{ab}|/|dr_{cb}|) + .. math:: + E = k*(theta - theta_0)^2 + + Args: + angle_numbers(int32): the number of angles M. + + Inputs: + - **uint_crd_f** (Tensor, uint32) - [N, 3], the unsigned int coordinate value of each atom. + - **scaler_f** (Tensor, float32) - [3,], the 3-D scale factor between + the real space float coordinates and the unsigned int coordinates. + - **atom_a** (Tensor, int32) - [M,], the 1st atom index of each angle. + - **atom_b** (Tensor, int32) - [M,], the 2nd and the central atom index of each angle. + - **atom_c** (Tensor, int32) - [M,], the 3rd atom index of each angle. + - **angle_k** (Tensor, float32) - [M,], the force constant for each angle. + - **angle_theta0** (Tensor, float32) - [M,], the equilibrium position value for each angle. + + Outputs: + - **ene** (Tensor, float32) - [M,], the potential energy for each angle term. + + Supported Platforms: + ``GPU`` + """ + + @prim_attr_register + def __init__(self, angle_numbers): + validator.check_value_type('angle_numbers', angle_numbers, (int), self.name) + self.angle_numbers = angle_numbers + self.init_prim_io_names(inputs=['uint_crd_f', 'scaler_f', 'atom_a', 'atom_b', 'atom_c', 'angle_k', + 'angle_theta0'], + outputs=['ene']) + self.add_prim_attr('angle_numbers', self.angle_numbers) + + 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): + cls_name = self.name + M = self.angle_numbers + validator.check_int(len(uint_crd_f_shape), 2, Rel.EQ, "uint_crd_f_dim", cls_name) + validator.check_int(len(scaler_f_shape), 1, Rel.EQ, "scaler_f_dim", cls_name) + validator.check_int(len(atom_a_shape), 1, Rel.EQ, "atom_a_dim", cls_name) + validator.check_int(len(atom_b_shape), 1, Rel.EQ, "atom_b_dim", cls_name) + validator.check_int(len(atom_c_shape), 1, Rel.EQ, "atom_c_dim", cls_name) + validator.check_int(len(angle_k_shape), 1, Rel.EQ, "angle_k_dim", cls_name) + validator.check_int(len(angle_theta0_shape), 1, Rel.EQ, "angle_theta0_dim", cls_name) + + validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f_shape[1]", cls_name) + validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f_shape", cls_name) + validator.check_int(atom_a_shape[0], M, Rel.EQ, "atom_a_shape", cls_name) + validator.check_int(atom_b_shape[0], M, Rel.EQ, "atom_b_shape", cls_name) + validator.check_int(atom_c_shape[0], M, Rel.EQ, "atom_c_shape", cls_name) + validator.check_int(angle_k_shape[0], M, Rel.EQ, "angle_k_shape", cls_name) + validator.check_int(angle_theta0_shape[0], M, Rel.EQ, "angle_theta0_shape", cls_name) + return [M,] + + def infer_dtype(self, uint_crd_f_dtype, scaler_f_type, atom_a_type, atom_b_type, atom_c_type, angle_k_type, + angle_theta0_type): + validator.check_tensor_dtype_valid('uint_crd_f', uint_crd_f_dtype, [mstype.uint32], self.name) + validator.check_tensor_dtype_valid('scaler_f', scaler_f_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('atom_a', atom_a_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('atom_b', atom_b_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('atom_c', atom_c_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('angle_k', angle_k_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('angle_theta0', angle_theta0_type, [mstype.float32], self.name) + return angle_k_type + + +class AngleAtomEnergy(PrimitiveWithInfer): + """ + Add the potential energy caused by angle terms to the total potential + 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(). + + Args: + angle_numbers(int32): the number of angles M. + + Inputs: + - **uint_crd_f** (Tensor, uint32) - [N, 3], the unsigned int coordinate value of each atom. + - **scaler_f** (Tensor, float32) - [3,], the 3-D scale factor between + the real space float coordinates and the unsigned int coordinates. + - **atom_a** (Tensor, int32) - [M,], the 1st atom index of each angle. + - **atom_b** (Tensor, int32) - [M,], the 2nd and the central atom index of each angle. + - **atom_c** (Tensor, int32) - [M,], the 3rd atom index of each angle. + - **angle_k** (Tensor, float32) - [M,], the force constant for each angle. + - **angle_theta0** (Tensor, float32) - [M,], the equilibrium position value for each angle. + + Outputs: + - **ene** (Tensor, float32) - [N,], the accumulated potential energy for each atom. + + Supported Platforms: + ``GPU`` + """ + + @prim_attr_register + def __init__(self, angle_numbers): + validator.check_value_type('angle_numbers', angle_numbers, (int), self.name) + self.angle_numbers = angle_numbers + self.init_prim_io_names(inputs=['uint_crd_f', 'scaler_f', 'atom_a', 'atom_b', 'atom_c', 'angle_k', + 'angle_theta0'], + outputs=['ene']) + self.add_prim_attr('angle_numbers', self.angle_numbers) + + 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): + cls_name = self.name + N = uint_crd_f_shape[0] + M = self.angle_numbers + validator.check_int(len(uint_crd_f_shape), 2, Rel.EQ, "uint_crd_f_dim", cls_name) + validator.check_int(len(scaler_f_shape), 1, Rel.EQ, "scaler_f_dim", cls_name) + validator.check_int(len(atom_a_shape), 1, Rel.EQ, "atom_a_dim", cls_name) + validator.check_int(len(atom_b_shape), 1, Rel.EQ, "atom_b_dim", cls_name) + validator.check_int(len(atom_c_shape), 1, Rel.EQ, "atom_c_dim", cls_name) + validator.check_int(len(angle_k_shape), 1, Rel.EQ, "angle_k_dim", cls_name) + validator.check_int(len(angle_theta0_shape), 1, Rel.EQ, "angle_theta0_dim", cls_name) + + validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f_shape[1]", cls_name) + validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f_shape", cls_name) + validator.check_int(atom_a_shape[0], M, Rel.EQ, "atom_a_shape", cls_name) + validator.check_int(atom_b_shape[0], M, Rel.EQ, "atom_b_shape", cls_name) + validator.check_int(atom_c_shape[0], M, Rel.EQ, "atom_c_shape", cls_name) + validator.check_int(angle_k_shape[0], M, Rel.EQ, "angle_k_shape", cls_name) + validator.check_int(angle_theta0_shape[0], M, Rel.EQ, "angle_theta0_shape", cls_name) + return [N,] + + def infer_dtype(self, uint_crd_f_dtype, scaler_f_type, atom_a_type, atom_b_type, atom_c_type, angle_k_type, + angle_theta0_type): + validator.check_tensor_dtype_valid('uint_crd_f', uint_crd_f_dtype, [mstype.uint32], self.name) + validator.check_tensor_dtype_valid('scaler_f', scaler_f_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('atom_a', atom_a_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('atom_b', atom_b_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('atom_c', atom_c_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('angle_k', angle_k_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('angle_theta0', angle_theta0_type, [mstype.float32], self.name) + return angle_k_type + + +class AngleForceWithAtomEnergy(PrimitiveWithInfer): + """ + 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(). + + Args: + angle_numbers(int32): the number of angles M. + + Inputs: + - **uint_crd_f** (Tensor, uint32) - [N, 3], the unsigned int coordinate value of each atom. + - **scaler_f** (Tensor, float32) - [3,], the 3-D scale factor between + the real space float coordinates and the unsigned int coordinates. + - **atom_a** (Tensor, int32) - [M,], the 1st atom index of each angle. + - **atom_b** (Tensor, int32) - [M,], the 2nd and the central atom index of each angle. + - **atom_c** (Tensor, int32) - [M,], the 3rd atom index of each angle. + - **angle_k** (Tensor, float32) - [M,], the force constant for each angle. + - **angle_theta0** (Tensor, float32) - [M,], the equilibrium position value for each angle. + + Outputs: + - **frc_f** (Tensor, float32) - [N, 3], same as operator AngleForce(). + - **ene** (Tensor, float) - [N,], same as operator AngleAtomEnergy(). + + Supported Platforms: + ``GPU`` + """ + + @prim_attr_register + def __init__(self, angle_numbers): + validator.check_value_type('angle_numbers', angle_numbers, (int), self.name) + self.angle_numbers = angle_numbers + self.init_prim_io_names(inputs=['uint_crd_f', 'scaler_f', 'atom_a', 'atom_b', 'atom_c', 'angle_k', + 'angle_theta0'], + outputs=['frc_f', 'ene']) + self.add_prim_attr('angle_numbers', self.angle_numbers) + + 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): + cls_name = self.name + N = uint_crd_f_shape[0] + M = self.angle_numbers + validator.check_int(len(uint_crd_f_shape), 2, Rel.EQ, "uint_crd_f_dim", cls_name) + validator.check_int(len(scaler_f_shape), 1, Rel.EQ, "scaler_f_dim", cls_name) + validator.check_int(len(atom_a_shape), 1, Rel.EQ, "atom_a_dim", cls_name) + validator.check_int(len(atom_b_shape), 1, Rel.EQ, "atom_b_dim", cls_name) + validator.check_int(len(atom_c_shape), 1, Rel.EQ, "atom_c_dim", cls_name) + validator.check_int(len(angle_k_shape), 1, Rel.EQ, "angle_k_dim", cls_name) + validator.check_int(len(angle_theta0_shape), 1, Rel.EQ, "angle_theta0_dim", cls_name) + + validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f_shape[1]", cls_name) + validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f_shape", cls_name) + validator.check_int(atom_a_shape[0], M, Rel.EQ, "atom_a_shape", cls_name) + validator.check_int(atom_b_shape[0], M, Rel.EQ, "atom_b_shape", cls_name) + validator.check_int(atom_c_shape[0], M, Rel.EQ, "atom_c_shape", cls_name) + validator.check_int(angle_k_shape[0], M, Rel.EQ, "angle_k_shape", cls_name) + validator.check_int(angle_theta0_shape[0], M, Rel.EQ, "angle_theta0_shape", cls_name) + return uint_crd_f_shape, [N,] + + def infer_dtype(self, uint_crd_f_dtype, scaler_f_type, atom_a_type, atom_b_type, atom_c_type, angle_k_type, + angle_theta0_type): + validator.check_tensor_dtype_valid('uint_crd_f', uint_crd_f_dtype, [mstype.uint32], self.name) + validator.check_tensor_dtype_valid('scaler_f', scaler_f_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('atom_a', atom_a_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('atom_b', atom_b_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('atom_c', atom_c_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('angle_k', angle_k_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('angle_theta0', angle_theta0_type, [mstype.float32], self.name) + return angle_k_type, angle_k_type + + +class Dihedral14LJForce(PrimitiveWithInfer): + """ + Calculate the Lennard-Jones part of 1,4 dihedral force correction + for each necessary dihedral terms on the corresponding atoms. + + Assume the number of necessary dihedral 1,4 terms is M, the number of atoms is N, + and the number of Lennard-Jones types for all atoms is P, which means + there will be Q = P*(P+1)/2 types of possible Lennard-Jones interactions + for all kinds of atom pairs. + + .. math:: + 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 + + Args: + nb14_numbers (int32): the number of necessary dihedral 1,4 terms M. + atom_numbers (int32): the number of atoms N. + + Inputs: + - **uint_crd_f** (Tensor, uint32) - [N, 3], the unsigned int coordinate value of each atom. + - **LJ_type** (Tensor, int32) - [N,], the Lennard-Jones type of each atom. + - **charge** (Tensor, float32) - [N,], the charge of each atom. + - **boxlength_f** (Tensor, float32) - [3,], the length of molecular simulation box in 3 dimensions. + - **a_14** (Tensor, int32) - [M,], the first atom index of each dihedral 1,4 term. + - **b_14** (Tensor, int32) - [M,], the second atom index of each dihedral 1,4 term. + - **lj_scale_factor** (Tensor, float32) - [M,], the scale factor for the + 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. + 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. + Q is the number of atom pair. + + Outputs: + - **frc_f** (Tensor, float32) - [N, 3], the force felt by each atom. + + Supported Platforms: + ``GPU`` + """ + + @prim_attr_register + def __init__(self, nb14_numbers, atom_numbers): + validator.check_value_type('nb14_numbers', nb14_numbers, (int), self.name) + validator.check_value_type('atom_numbers', atom_numbers, (int), self.name) + self.dihedral_14_numbers = nb14_numbers + self.atom_numbers = atom_numbers + self.init_prim_io_names( + inputs=['uint_crd_f', 'LJtype', 'charge', 'boxlength_f', 'a_14', 'b_14', 'lj_scale_factor', + 'LJ_type_A', 'LJ_type_B'], + outputs=['frc_f']) + self.add_prim_attr('dihedral_14_numbers', self.dihedral_14_numbers) + self.add_prim_attr('atom_numbers', self.atom_numbers) + + 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): + cls_name = self.name + N = self.atom_numbers + Q = LJ_type_A_shape[0] + validator.check_int(len(uint_crd_f_shape), 2, Rel.EQ, "uint_crd_f_dim", cls_name) + validator.check_int(len(LJtype_shape), 1, Rel.EQ, "LJtype_dim", cls_name) + validator.check_int(len(charge_shape), 1, Rel.EQ, "charge_dim", cls_name) + validator.check_int(len(boxlength_f_shape), 1, Rel.EQ, "boxlength_f_dim", cls_name) + validator.check_int(len(a_14_shape), 1, Rel.EQ, "a_14_dim", cls_name) + validator.check_int(len(b_14_shape), 1, Rel.EQ, "b_14_dim", cls_name) + validator.check_int(len(lj_scale_factor_shape), 1, Rel.EQ, "lj_scale_factor_dim", cls_name) + validator.check_int(len(LJ_type_B_shape), 1, Rel.EQ, "LJ_type_B_dim", 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[1]", cls_name) + validator.check_int(LJtype_shape[0], N, Rel.EQ, "LJtype", cls_name) + validator.check_int(charge_shape[0], N, Rel.EQ, "charge", cls_name) + validator.check_int(boxlength_f_shape[0], 3, Rel.EQ, "boxlength_f", cls_name) + validator.check_int(LJ_type_B_shape[0], Q, Rel.EQ, "LJ_type_B", cls_name) + 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, + lj_scale_factor_type, LJ_type_A_type, LJ_type_B_type): + validator.check_tensor_dtype_valid('uint_crd_f', uint_crd_f_dtype, [mstype.uint32], self.name) + validator.check_tensor_dtype_valid('LJtype', LJtype_dtype, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('charge', charge_dtype, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('boxlength_f', boxlength_f_type, [mstype.float32], self.name) + + validator.check_tensor_dtype_valid('a_14', a_14_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('b_14', b_14_type, [mstype.int32], self.name) + + validator.check_tensor_dtype_valid('lj_scale_factor', lj_scale_factor_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('LJ_type_A', LJ_type_A_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('LJ_type_B', LJ_type_B_type, [mstype.float32], self.name) + return LJ_type_B_type + + +class Dihedral14LJEnergy(PrimitiveWithInfer): + """ + Calculate the Lennard-Jones part of 1,4 dihedral energy correction for + each necessary dihedral terms on the corresponding atoms. + + .. math:: + dr = (x_a-x_b, y_a-y_b, z_a-z-b) + .. math:: + E = k*(A/|dr|^{12} - B/|dr|^{6}) + + Args: + nb14_numbers (int32): the number of necessary dihedral 1,4 terms M. + atom_numbers (int32): the number of atoms N. + + Inputs: + - **uint_crd_f** (Tensor, uint32) - [N, 3], the unsigned int coordinate value of each atom. + - **LJ_type** (Tensor, int32) - [N,], the Lennard-Jones type of each atom. + - **charge** (Tensor, float32) - [N,], the charge of each atom. + - **boxlength_f** (Tensor, float32) - [3,], the length of molecular simulation box in 3 dimensions. + - **a_14** (Tensor, int32) - [M,], the first atom index of each dihedral 1,4 term. + - **b_14** (Tensor, int32) - [M,], the second atom index of each dihedral 1,4 term. + - **lj_scale_factor** (Tensor, float32) - [M,], the scale factor for the + 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. + 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. + Q is the number of atom pair. + + Outputs: + - **ene** (Tensor, float32) - [M,], the Lennard-Jones potential + energy correction for each necessary dihedral 1,4 term. + + Supported Platforms: + ``GPU`` + """ + + @prim_attr_register + def __init__(self, nb14_numbers, atom_numbers): + validator.check_value_type('nb14_numbers', nb14_numbers, (int), self.name) + validator.check_value_type('atom_numbers', atom_numbers, (int), self.name) + self.dihedral_14_numbers = nb14_numbers + self.atom_numbers = atom_numbers + + self.init_prim_io_names( + inputs=['uint_crd_f', 'LJtype', 'charge', 'boxlength_f', 'a_14', 'b_14', 'lj_scale_factor', + 'LJ_type_A', 'LJ_type_B'], + outputs=['ene']) + self.add_prim_attr('dihedral_14_numbers', self.dihedral_14_numbers) + self.add_prim_attr('atom_numbers', self.atom_numbers) + + 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): + cls_name = self.name + N = self.atom_numbers + Q = LJ_type_A_shape[0] + validator.check_int(len(uint_crd_f_shape), 2, Rel.EQ, "uint_crd_f_dim", cls_name) + validator.check_int(len(LJtype_shape), 1, Rel.EQ, "LJtype_dim", cls_name) + validator.check_int(len(charge_shape), 1, Rel.EQ, "charge_dim", cls_name) + validator.check_int(len(boxlength_f_shape), 1, Rel.EQ, "boxlength_f_dim", cls_name) + validator.check_int(len(a_14_shape), 1, Rel.EQ, "a_14_dim", cls_name) + validator.check_int(len(b_14_shape), 1, Rel.EQ, "b_14_dim", cls_name) + validator.check_int(len(lj_scale_factor_shape), 1, Rel.EQ, "lj_scale_factor_dim", cls_name) + validator.check_int(len(LJ_type_B_shape), 1, Rel.EQ, "LJ_type_B_dim", 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[1]", cls_name) + validator.check_int(LJtype_shape[0], N, Rel.EQ, "LJtype", cls_name) + validator.check_int(charge_shape[0], N, Rel.EQ, "charge", cls_name) + validator.check_int(boxlength_f_shape[0], 3, Rel.EQ, "boxlength_f", cls_name) + validator.check_int(LJ_type_B_shape[0], Q, Rel.EQ, "LJ_type_B", cls_name) + 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, + lj_scale_factor_type, LJ_type_A_type, LJ_type_B_type): + validator.check_tensor_dtype_valid('uint_crd_f', uint_crd_f_dtype, [mstype.uint32], self.name) + validator.check_tensor_dtype_valid('LJtype', LJtype_dtype, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('charge', charge_dtype, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('boxlength_f', boxlength_f_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('a_14', a_14_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('b_14', b_14_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('lj_scale_factor', lj_scale_factor_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('LJ_type_A', LJ_type_A_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('LJ_type_B', LJ_type_B_type, [mstype.float32], self.name) + + return LJ_type_A_type + + +class Dihedral14LJForceWithDirectCF(PrimitiveWithInfer): + """ + Calculate the Lennard-Jones part and the Coulomb part of force correction + for each necessary dihedral 1,4 terms. + + The calculation formula of the Lennard-Jones part is the same as operator + Dihedral14LJForce(), and the Coulomb part is as follows: + + .. math:: + dr = (x_a-x_b, y_a-y_b, z_a-z_b) + .. math:: + F = -k*q_a*q_b/|r|^3*dr + + Args: + nb14_numbers (int32): the number of necessary dihedral 1,4 terms M. + atom_numbers (int32): the number of atoms N. + + Inputs: + - **uint_crd_f** (Tensor, uint32) - [N, 3], the unsigned int coordinate value of each atom. + - **LJ_type** (Tensor, int32) - [N,], the Lennard-Jones type of each atom. + - **charge** (Tensor, float32) - [N,], the charge of each atom. + - **boxlength_f** (Tensor, float32) - [3,], the length of molecular simulation box in 3 dimensions. + - **a_14** (Tensor, int32) - [M,], the first atom index of each dihedral 1,4 term. + - **b_14** (Tensor, int32) - [M,], the second atom index of each dihedral 1,4 term. + - **lj_scale_factor** (Tensor, float32) - [M,], the scale factor for the + Lennard-Jones part of force correction of each dihedral 1,4 term. + - **cf_scale_factor** (Tensor, float) - [M,], the scale factor for the + 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. + 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. + Q is the number of atom pair. + + Outputs: + - **frc_f** (Tensor, float) - [N, 3], the force felt by each atom. + + Supported Platforms: + ``GPU`` + """ + + @prim_attr_register + def __init__(self, nb14_numbers, atom_numbers): + validator.check_value_type('nb14_numbers', nb14_numbers, (int), self.name) + validator.check_value_type('atom_numbers', atom_numbers, (int), self.name) + self.dihedral_14_numbers = nb14_numbers + self.atom_numbers = atom_numbers + + self.init_prim_io_names( + inputs=['uint_crd_f', 'LJtype', 'charge', 'boxlength_f', 'a_14', 'b_14', 'lj_scale_factor', + 'cf_scale_factor', + 'LJ_type_A', 'LJ_type_B'], + outputs=['frc_f']) + self.add_prim_attr('dihedral_14_numbers', self.dihedral_14_numbers) + self.add_prim_attr('atom_numbers', self.atom_numbers) + + 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): + cls_name = self.name + N = self.atom_numbers + M = self.dihedral_14_numbers + Q = LJ_type_A_shape[0] + validator.check_int(len(uint_crd_f_shape), 2, Rel.EQ, "uint_crd_f_dim", cls_name) + validator.check_int(len(LJtype_shape), 1, Rel.EQ, "LJtype_dim", cls_name) + validator.check_int(len(charge_shape), 1, Rel.EQ, "charge_dim", cls_name) + validator.check_int(len(boxlength_f_shape), 1, Rel.EQ, "boxlength_f_dim", cls_name) + validator.check_int(len(a_14_shape), 1, Rel.EQ, "a_14_dim", cls_name) + validator.check_int(len(b_14_shape), 1, Rel.EQ, "b_14_dim", cls_name) + validator.check_int(len(lj_scale_factor_shape), 1, Rel.EQ, "lj_scale_factor_dim", cls_name) + validator.check_int(len(cf_scale_factor_shape), 1, Rel.EQ, "cf_scale_factor_dim", cls_name) + validator.check_int(len(LJ_type_B_shape), 1, Rel.EQ, "LJ_type_B_dim", cls_name) + + validator.check_int(uint_crd_f_shape[0], N, Rel.EQ, "uint_crd_f_shape[0]", cls_name) + validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f_shape[1]", cls_name) + validator.check_int(LJtype_shape[0], N, Rel.EQ, "LJtype_shape", cls_name) + validator.check_int(charge_shape[0], M, Rel.EQ, "charge_shape", cls_name) + validator.check_int(boxlength_f_shape[0], 3, Rel.EQ, "boxlength_f_shape", cls_name) + validator.check_int(LJ_type_B_shape[0], Q, Rel.EQ, "LJ_type_B_shape", cls_name) + 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, + lj_scale_factor_type, cf_scale_factor_type, LJ_type_A_type, LJ_type_B_type): + validator.check_tensor_dtype_valid('uint_crd_f', uint_crd_f_dtype, [mstype.uint32], self.name) + validator.check_tensor_dtype_valid('LJtype', LJtype_dtype, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('charge', charge_dtype, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('boxlength_f', boxlength_f_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('a_14', a_14_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('b_14', b_14_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('lj_scale_factor', lj_scale_factor_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('cf_scale_factor', cf_scale_factor_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('LJ_type_A', LJ_type_A_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('LJ_type_B', LJ_type_B_type, [mstype.float32], self.name) + + return LJ_type_A_type + + +class Dihedral14LJCFForceWithAtomEnergy(PrimitiveWithInfer): + """ + Calculate the Lennard-Jones and Coulumb energy correction and force correction + for each necessary dihedral 1,4 terms together and add them to the total force + and potential energy for each atom. + + The calculation formula of force correction is the same as operator + Dihedral14LJForceWithDirectCF(), and the energy correction part is the same + as operator Dihedral14LJEnergy() and Dihedral14CFEnergy(). + + Args: + nb14_numbers (int32): the number of necessary dihedral 1,4 terms M. + atom_numbers (int32): the number of atoms N. + + Inputs: + - **uint_crd_f** (Tensor, uint32) - [N, 3], the unsigned int coordinate value of each atom. + - **LJ_type** (Tensor, int32) - [N,], the Lennard-Jones type of each atom. + - **charge** (Tensor, float32) - [N,], the charge of each atom. + - **boxlength_f** (Tensor, float32) - [3,], the length of molecular simulation box in 3 dimensions. + - **a_14** (Tensor, int32) - [M,], the first atom index of each dihedral 1,4 term. + - **b_14** (Tensor, int32) - [M,], the second atom index of each dihedral 1,4 term. + - **lj_scale_factor** (Tensor, float32) - [M,], the scale factor for the + Lennard-Jones part of force correction of each dihedral 1,4 term. + - **cf_scale_factor** (Tensor, float) - [M,], the scale factor for the + 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. + 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. + Q is the number of atom pair. + + Outputs: + - **frc_f** (Tensor, float32) - [N, 3], the force felt by each atom. + - **atom_energy** (Tensor, float32) - [N,], the accumulated potential energy for each atom. + + Supported Platforms: + ``GPU`` + """ + + @prim_attr_register + def __init__(self, nb14_numbers, atom_numbers): + validator.check_value_type('nb14_numbers', nb14_numbers, (int), self.name) + validator.check_value_type('atom_numbers', atom_numbers, (int), self.name) + self.dihedral_14_numbers = nb14_numbers + self.atom_numbers = atom_numbers + + self.init_prim_io_names( + inputs=['uint_crd_f', 'LJtype', 'charge', 'boxlength_f', 'a_14', 'b_14', 'lj_scale_factor', + 'cf_scale_factor', + 'LJ_type_A', 'LJ_type_B'], + outputs=['frc_f', 'atom_energy']) + self.add_prim_attr('dihedral_14_numbers', self.dihedral_14_numbers) + self.add_prim_attr('atom_numbers', self.atom_numbers) + + 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): + cls_name = self.name + N = self.atom_numbers + Q = LJ_type_A_shape[0] + validator.check_int(len(uint_crd_f_shape), 2, Rel.EQ, "uint_crd_f_dim", cls_name) + validator.check_int(len(LJtype_shape), 1, Rel.EQ, "LJtype_dim", cls_name) + validator.check_int(len(charge_shape), 1, Rel.EQ, "charge_dim", cls_name) + validator.check_int(len(boxlength_f_shape), 1, Rel.EQ, "boxlength_f_dim", cls_name) + validator.check_int(len(a_14_shape), 1, Rel.EQ, "a_14_dim", cls_name) + validator.check_int(len(b_14_shape), 1, Rel.EQ, "b_14_dim", cls_name) + validator.check_int(len(lj_scale_factor_shape), 1, Rel.EQ, "lj_scale_factor_dim", cls_name) + validator.check_int(len(cf_scale_factor_shape), 1, Rel.EQ, "cf_scale_factor_dim", cls_name) + validator.check_int(len(LJ_type_B_shape), 1, Rel.EQ, "LJ_type_B_dim", cls_name) + + validator.check_int(uint_crd_f_shape[0], N, Rel.EQ, "uint_crd_f_shape[0]", cls_name) + validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f_shape[1]", cls_name) + validator.check_int(LJtype_shape[0], N, Rel.EQ, "LJtype_shape", cls_name) + validator.check_int(charge_shape[0], N, Rel.EQ, "charge_shape", cls_name) + validator.check_int(boxlength_f_shape[0], 3, Rel.EQ, "boxlength_f_shape", cls_name) + validator.check_int(LJ_type_B_shape[0], Q, Rel.EQ, "LJ_type_B_shape", cls_name) + 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, + lj_scale_factor_type, cf_scale_factor_type, LJ_type_A_type, LJ_type_B_type): + validator.check_tensor_dtype_valid('uint_crd_f', uint_crd_f_dtype, [mstype.uint32], self.name) + validator.check_tensor_dtype_valid('LJtype', LJtype_dtype, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('charge', charge_dtype, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('boxlength_f', boxlength_f_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('a_14', a_14_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('b_14', b_14_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('lj_scale_factor', lj_scale_factor_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('cf_scale_factor', cf_scale_factor_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('LJ_type_A', LJ_type_A_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('LJ_type_B', LJ_type_B_type, [mstype.float32], self.name) + + return charge_dtype, charge_dtype + + +class Dihedral14LJAtomEnergy(PrimitiveWithInfer): + """ + Add the potenrial energy caused by Lennard-Jones energy correction for each + necessary dihedral 1,4 terms to the total potential energy of each atom. + + The calculation formula is the same as operator Dihedral14LJEnergy(). + + Args: + nb14_numbers (int32): the number of necessary dihedral 1,4 terms M. + atom_numbers (int32): the number of atoms N. + + Inputs: + - **uint_crd_f** (Tensor, uint32) - [N, 3], the unsigned int coordinate value of each atom. + - **LJ_type** (Tensor, int32) - [N,], the Lennard-Jones type of each atom. + - **charge** (Tensor, float32) - [N,], the charge of each atom. + - **boxlength_f** (Tensor, float32) - [3,], the length of molecular simulation box in 3 dimensions. + - **a_14** (Tensor, int32) - [M,], the first atom index of each dihedral 1,4 term. + - **b_14** (Tensor, int32) - [M,], the second atom index of each dihedral 1,4 term. + - **lj_scale_factor** (Tensor, float32) - [M,], the scale factor for the + 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. + 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. + Q is the number of atom pair. + + Outputs: + - **ene** (Tensor, float32) - [N,], the accumulated potential energy of each atom. + + Supported Platforms: + ``GPU`` + """ + + @prim_attr_register + def __init__(self, nb14_numbers, atom_numbers): + validator.check_value_type('nb14_numbers', nb14_numbers, (int), self.name) + validator.check_value_type('atom_numbers', atom_numbers, (int), self.name) + self.dihedral_14_numbers = nb14_numbers + self.atom_numbers = atom_numbers + + self.init_prim_io_names( + inputs=['uint_crd_f', 'LJtype', 'charge', 'boxlength_f', 'a_14', 'b_14', 'lj_scale_factor', + 'LJ_type_A', 'LJ_type_B'], + outputs=['ene']) + self.add_prim_attr('dihedral_14_numbers', self.dihedral_14_numbers) + self.add_prim_attr('atom_numbers', self.atom_numbers) + + 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): + cls_name = self.name + N = self.atom_numbers + Q = LJ_type_A_shape[0] + validator.check_int(len(uint_crd_f_shape), 2, Rel.EQ, "uint_crd_f_dim", cls_name) + validator.check_int(len(LJtype_shape), 1, Rel.EQ, "LJtype_dim", cls_name) + validator.check_int(len(charge_shape), 1, Rel.EQ, "charge_dim", cls_name) + validator.check_int(len(boxlength_f_shape), 1, Rel.EQ, "boxlength_f_dim", cls_name) + validator.check_int(len(a_14_shape), 1, Rel.EQ, "a_14_dim", cls_name) + validator.check_int(len(b_14_shape), 1, Rel.EQ, "b_14_dim", cls_name) + validator.check_int(len(lj_scale_factor_shape), 1, Rel.EQ, "lj_scale_factor_dim", cls_name) + validator.check_int(len(LJ_type_B_shape), 1, Rel.EQ, "LJ_type_B_dim", cls_name) + + validator.check_int(uint_crd_f_shape[0], N, Rel.EQ, "uint_crd_f_shape[0]", cls_name) + validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f_shape[1]", cls_name) + validator.check_int(LJtype_shape[0], N, Rel.EQ, "LJtype_shape", cls_name) + validator.check_int(charge_shape[0], N, Rel.EQ, "charge_shape", cls_name) + validator.check_int(boxlength_f_shape[0], 3, Rel.EQ, "boxlength_f_shape", cls_name) + validator.check_int(LJ_type_B_shape[0], Q, Rel.EQ, "LJ_type_B_shape", cls_name) + return LJtype_shape + + def infer_dtype(self, uint_crd_f_dtype, LJtype_dtype, charge_dtype, boxlength_f_type, a_14_type, b_14_type, + lj_scale_factor_type, LJ_type_A_type, LJ_type_B_type): + validator.check_tensor_dtype_valid('uint_crd_f', uint_crd_f_dtype, [mstype.uint32], self.name) + validator.check_tensor_dtype_valid('LJtype', LJtype_dtype, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('charge', charge_dtype, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('boxlength_f', boxlength_f_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('a_14', a_14_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('b_14', b_14_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('lj_scale_factor', lj_scale_factor_type, [mstype.float32], + self.name) + validator.check_tensor_dtype_valid('LJ_type_A', LJ_type_A_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('LJ_type_B', LJ_type_B_type, [mstype.float32], self.name) + + return LJ_type_A_type + + +class Dihedral14CFEnergy(PrimitiveWithInfer): + """ + Calculate the Coulumb part of 1,4 dihedral energy correction for + each necessary dihedral terms on the corresponding atoms. + + .. math:: + + dr = (x_a-x_b, y_a-y_b, z_a-z_b) + + .. math:: + E = k*q_a*q_b/|dr| + + Args: + nb14_numbers (int32): the number of necessary dihedral 1,4 terms M. + atom_numbers (int32): the number of atoms N. + + Inputs: + - **uint_crd_f** (Tensor, uint32) - [N, 3], the unsigned int coordinate value of each atom. + - **LJ_type** (Tensor, int32) - [N,], the Lennard-Jones type of each atom. + - **charge** (Tensor, float32) - [N,], the charge of each atom. + - **boxlength_f** (Tensor, float32) - [3,], the length of molecular simulation box in 3 dimensions. + - **a_14** (Tensor, int32) - [M,], the first atom index of each dihedral 1,4 term. + - **b_14** (Tensor, int32) - [M,], the second atom index of each dihedral 1,4 term. + - **cf_scale_factor** (Tensor, float) - [M,], the scale factor for the + Coulomb part of force correction for each dihedral 1,4 terms. + + Outputs: + - **ene** (Tensor, float32) - [M,], the accumulated potential energy of each atom. + + Supported Platforms: + ``GPU`` + """ + + @prim_attr_register + def __init__(self, nb14_numbers, atom_numbers): + validator.check_value_type('nb14_numbers', nb14_numbers, (int), self.name) + validator.check_value_type('atom_numbers', atom_numbers, (int), self.name) + self.dihedral_14_numbers = nb14_numbers + self.atom_numbers = atom_numbers + + self.init_prim_io_names( + inputs=['uint_crd_f', 'LJtype', 'charge', 'boxlength_f', 'a_14', 'b_14', 'cj_scale_factor'], + outputs=['ene']) + self.add_prim_attr('dihedral_14_numbers', self.dihedral_14_numbers) + self.add_prim_attr('atom_numbers', self.atom_numbers) + + 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): + cls_name = self.name + N = self.atom_numbers + validator.check_int(len(uint_crd_f_shape), 2, Rel.EQ, "uint_crd_f_dim", cls_name) + validator.check_int(len(LJtype_shape), 1, Rel.EQ, "LJtype_dim", cls_name) + validator.check_int(len(charge_shape), 1, Rel.EQ, "charge_dim", cls_name) + validator.check_int(len(boxlength_f_shape), 1, Rel.EQ, "boxlength_f_dim", cls_name) + validator.check_int(len(a_14_shape), 1, Rel.EQ, "a_14_dim", cls_name) + validator.check_int(len(b_14_shape), 1, Rel.EQ, "b_14_dim", cls_name) + validator.check_int(len(cf_scale_factor_shape), 1, Rel.EQ, "cf_scale_factor_dim", cls_name) + + validator.check_int(uint_crd_f_shape[0], N, Rel.EQ, "uint_crd_f_shape[0]", cls_name) + validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f_shape[1]", cls_name) + validator.check_int(LJtype_shape[0], N, Rel.EQ, "LJtype_shape", cls_name) + validator.check_int(charge_shape[0], N, Rel.EQ, "charge_shape", cls_name) + validator.check_int(boxlength_f_shape[0], 3, Rel.EQ, "boxlength_f_shape", cls_name) + 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, + cf_scale_factor_type): + validator.check_tensor_dtype_valid('uint_crd_f', uint_crd_f_dtype, [mstype.uint32], self.name) + validator.check_tensor_dtype_valid('LJtype', LJtype_dtype, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('charge', charge_dtype, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('boxlength_f', boxlength_f_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('a_14', a_14_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('b_14', b_14_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('lj_scale_factor', cf_scale_factor_type, [mstype.float32], + self.name) + + return charge_dtype + + +class Dihedral14CFAtomEnergy(PrimitiveWithInfer): + """ + Add the potential energy caused by Coulumb energy correction for each + necessary dihedral 1,4 terms to the total potential energy of each atom. + + The calculation formula is the same as operator Dihedral14CFEnergy(). + + Args: + nb14_numbers (int32): the number of necessary dihedral 1,4 terms M. + atom_numbers (int32): the number of atoms N. + + Inputs: + - **uint_crd_f** (Tensor, uint32) - [N, 3], the unsigned int coordinate value of each atom. + - **LJ_type** (Tensor, int32) - [N,], the Lennard-Jones type of each atom. + - **charge** (Tensor, float32) - [N,], the charge of each atom. + - **boxlength_f** (Tensor, float32) - [3,], the length of molecular simulation box in 3 dimensions. + - **a_14** (Tensor, int32) - [M,], the first atom index of each dihedral 1,4 term. + - **b_14** (Tensor, int32) - [M,], the second atom index of each dihedral 1,4 term. + - **cf_scale_factor** (Tensor, float) - [M,], the scale factor for the + Coulomb part of force correction for each dihedral 1,4 terms. + + Outputs: + - **ene** (Tensor, float32) - [N,], the accumulated potential energy of each atom. + + + Supported Platforms: + ``GPU`` + """ + + @prim_attr_register + def __init__(self, nb14_numbers, atom_numbers): + validator.check_value_type('nb14_numbers', nb14_numbers, (int), self.name) + validator.check_value_type('atom_numbers', atom_numbers, (int), self.name) + self.dihedral_14_numbers = nb14_numbers + self.atom_numbers = atom_numbers + + self.init_prim_io_names( + inputs=['uint_crd_f', 'LJtype', 'charge', 'boxlength_f', 'a_14', 'b_14', 'cf_scale_factor'], + outputs=['ene']) + self.add_prim_attr('dihedral_14_numbers', self.dihedral_14_numbers) + self.add_prim_attr('atom_numbers', self.atom_numbers) + + 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): + cls_name = self.name + N = self.atom_numbers + validator.check_int(len(uint_crd_f_shape), 2, Rel.EQ, "uint_crd_f_dim", cls_name) + validator.check_int(len(LJtype_shape), 1, Rel.EQ, "LJtype_dim", cls_name) + validator.check_int(len(charge_shape), 1, Rel.EQ, "charge_dim", cls_name) + validator.check_int(len(boxlength_f_shape), 1, Rel.EQ, "boxlength_f_dim", cls_name) + validator.check_int(len(a_14_shape), 1, Rel.EQ, "a_14_dim", cls_name) + validator.check_int(len(b_14_shape), 1, Rel.EQ, "b_14_dim", cls_name) + validator.check_int(len(cf_scale_factor_shape), 1, Rel.EQ, "cf_scale_factor_dim", cls_name) + + validator.check_int(uint_crd_f_shape[0], N, Rel.EQ, "uint_crd_f_shape[0]", cls_name) + validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f_shape[1]", cls_name) + validator.check_int(LJtype_shape[0], N, Rel.EQ, "LJtype_shape", cls_name) + validator.check_int(charge_shape[0], N, Rel.EQ, "charge_shape", cls_name) + validator.check_int(boxlength_f_shape[0], 3, Rel.EQ, "boxlength_f_shape", cls_name) + return LJtype_shape + + def infer_dtype(self, uint_crd_f_dtype, LJtype_dtype, charge_dtype, boxlength_f_type, a_14_type, b_14_type, + cf_scale_factor_type): + validator.check_tensor_dtype_valid('uint_crd_f', uint_crd_f_dtype, [mstype.uint32], self.name) + validator.check_tensor_dtype_valid('LJtype', LJtype_dtype, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('charge', charge_dtype, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('boxlength_f', boxlength_f_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('a_14', a_14_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('b_14', b_14_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('cf_scale_factor', cf_scale_factor_type, [mstype.float32], + self.name) + + return charge_dtype + + +class MDIterationLeapFrog(PrimitiveWithInfer): + """ + One step of classical leap frog algorithm to solve the finite difference + Hamiltonian equations of motion for certain system, using Langevin dynamics + with Liu's thermostat scheme. Assume the number of atoms is N and the target + control temperature is T. + + Detailed iteration formula can be found in this paper: A unified thermostat + scheme for efficient configurational sampling for classical/quantum canonical + ensembles via molecular dynamics. DOI: 10.1063/1.4991621. + + Args: + float4_numbers(int32): total length to store random numbers. + atom_numbers(int32): the number of atoms N. + dt(float32): time step for finite difference. + half_dt(float32): half of time step for finite difference. + exp_gamma(float32): parameter in Liu's dynamic, equals exp(-gamma_ln * dt), + where gamma_ln is the firction factor in Langvin dynamics. + max_velocity(float32): the upper limit of velocity, when the veclocity overflows, + scale it to the upper limit. + is_max_velocity(int32): whether the max velocity control is open or not. + + Inputs: + - **mass_inverse** (Tensor, float32) - [N,], the inverse value of mass of each atom. + - **sqrt_mass** (Tensor, float32) - [N,], the inverse square root value + of effect mass in Liu's dynamics of each atom. + + Outputs: + - **vel** (Tensor, float32) - [N, 3], the velocity of each atom. + - **crd** (Tensor, float32) - [N, 3], the coordinate of each atom. + - **frc** (Tensor, float32) - [N, 3], the force felt by each atom. + - **acc** (Tensor, float32) - [N, 3], the acceleration of each atom. + + Supported Platforms: + ``GPU`` + """ + + @prim_attr_register + def __init__(self, float4_numbers, atom_numbers, half_dt, dt, exp_gamma, is_max_velocity, max_velocity): + validator.check_value_type('float4_numbers', float4_numbers, (int), self.name) + validator.check_value_type('atom_numbers', atom_numbers, (int), self.name) + validator.check_value_type('half_dt', half_dt, (float), self.name) + validator.check_value_type('dt', dt, (float), self.name) + validator.check_value_type('exp_gamma', exp_gamma, (float), self.name) + validator.check_value_type('is_max_velocity', is_max_velocity, (int), self.name) + validator.check_value_type('max_velocity', max_velocity, (float), self.name) + self.float4_numbers = float4_numbers + self.atom_numbers = atom_numbers + self.half_dt = half_dt + self.dt = dt + self.exp_gamma = exp_gamma + self.is_max_velocity = is_max_velocity + self.max_velocity = max_velocity + + self.init_prim_io_names( + inputs=['mass_inverse', 'sqrt_mass'], + outputs=['vel', 'crd', 'frc', 'acc']) + self.add_prim_attr('float4_numbers', self.float4_numbers) + self.add_prim_attr('atom_numbers', self.atom_numbers) + self.add_prim_attr('half_dt', self.half_dt) + self.add_prim_attr('dt', self.dt) + self.add_prim_attr('exp_gamma', self.exp_gamma) + self.add_prim_attr('is_max_velocity', self.is_max_velocity) + self.add_prim_attr('max_velocity', self.max_velocity) + + 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] + + def infer_dtype(self, mass_inverse_dtype, sqrt_mass_dtype): + validator.check_tensor_dtype_valid('mass_inverse', mass_inverse_dtype, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('sqrt_mass', sqrt_mass_dtype, [mstype.float32], self.name) + + return mass_inverse_dtype, mass_inverse_dtype, mass_inverse_dtype, mass_inverse_dtype + + +class PMEReciprocalForce(PrimitiveWithInfer): + """ + Calculate the reciprocal part of long-range Coulumb force using + PME(Particle Meshed Ewald) method. Assume the number of atoms is N. + + The detailed calculation formula of PME(Particle Meshed Ewald) method + can be found in this paper: A Smooth Particle Mesh Ewald Method. DOI: + 10.1063/1.470117. + + Args: + atom_numbers(int32): the number of atoms, N. + beta(float32): the PME beta parameter, determined by the + non-bond cutoff value and simulation precision tolerance. + fftx(int32): the number of points for Fourier transform in dimension X. + ffty(int32): the number of points for Fourier transform in dimension Y. + fftz(int32): the number of points for Fourier transform in dimension Z. + box_length_0(float32): the value of boxlength idx 0 + box_length_1(float32): the value of boxlength idx 1 + box_length_2(float32): the value of boxlength idx 2 + + Inputs: + - **uint_crd** (Tensor, uint32) - [N, 3], the unsigned int coordinates value of each atom. + - **charge** (Tensor, float32) - [N,], the charge carried by each atom. + + Outputs: + - **force** (Tensor, float32) - [N, 3], the force felt by each atom. + + Supported Platforms: + ``GPU`` + """ + + @prim_attr_register + def __init__(self, atom_numbers, beta, fftx, ffty, fftz, box_length_0, box_length_1, box_length_2): + validator.check_value_type('atom_numbers', atom_numbers, (int), self.name) + validator.check_value_type('beta', beta, (float), self.name) + validator.check_value_type('fftx', fftx, (int), self.name) + validator.check_value_type('ffty', ffty, (int), self.name) + validator.check_value_type('fftz', fftz, (int), self.name) + self.atom_numbers = atom_numbers + self.beta = beta + self.fftx = fftx + self.ffty = ffty + self.fftz = fftz + self.box_length_0 = box_length_0 + self.box_length_1 = box_length_1 + self.box_length_2 = box_length_2 + + self.init_prim_io_names(inputs=['boxlength', 'uint_crd', 'charge'], + outputs=['force']) + self.add_prim_attr('atom_numbers', self.atom_numbers) + self.add_prim_attr('beta', self.beta) + self.add_prim_attr('fftx', self.fftx) + self.add_prim_attr('ffty', self.ffty) + self.add_prim_attr('fftz', self.fftz) + self.add_prim_attr('box_length_0', self.box_length_0) + self.add_prim_attr('box_length_1', self.box_length_1) + self.add_prim_attr('box_length_2', self.box_length_2) + + def infer_shape(self, uint_crd_shape, charge_shape): + cls_name = self.name + N = self.atom_numbers + validator.check_int(len(uint_crd_shape), 2, Rel.EQ, "uint_crd_dim", cls_name) + validator.check_int(len(charge_shape), 1, Rel.EQ, "charge_dim", cls_name) + + validator.check_int(uint_crd_shape[0], N, Rel.EQ, "uint_crd_shape[0]", cls_name) + validator.check_int(uint_crd_shape[1], 3, Rel.EQ, "uint_crd_shape[1]", cls_name) + validator.check_int(charge_shape[0], N, Rel.EQ, "charge_shape", cls_name) + return uint_crd_shape + + def infer_dtype(self, uint_crd_type, charge_type): + validator.check_tensor_dtype_valid('uint_crd', uint_crd_type, [mstype.uint32], self.name) + validator.check_tensor_dtype_valid('charge', charge_type, [mstype.float32], self.name) + return charge_type + + +class PMEExcludedForce(PrimitiveWithInfer): + """ + Calculate the excluded part of long-range Coulumb force using + PME(Particle Meshed Ewald) method. Assume the number of atoms is + N, and the length of excluded list is E. + + Args: + atom_numbers(int32): the number of atoms, N. + excluded_numbers(int32): the length of excluded list, E. + beta(float32): the PME beta parameter, determined by the + non-bond cutoff value and simulation precision tolerance. + + Inputs: + - **uint_crd** (Tensor, uint32) - [N, 3], the unsigned int coordinates value of each atom. + - **scaler** (Tensor, float32) - [3,], the scale factor between real space + coordinates and its unsigned int value. + - **charge** (Tensor, float32) - [N,], the charge carried by each atom. + - **excluded_list_start** (Tensor, int32) - [N,], the start excluded index + in excluded list for each atom. + - **excluded_list** (Tensor, int32) - [E,], the contiguous join of excluded + list of each atom. E is the number of excluded atoms. + - **excluded_atom_numbers** (Tensor, int32) - [N,], the number of atom excluded + in excluded list for each atom. + + Outputs: + - **force** (Tensor, float32) - [N, 3], the force felt by each atom. + + Supported Platforms: + ``GPU`` + """ + + @prim_attr_register + def __init__(self, atom_numbers, excluded_numbers, beta): + validator.check_value_type('atom_numbers', atom_numbers, (int), self.name) + validator.check_value_type('excluded_numbers', excluded_numbers, (int), self.name) + validator.check_value_type('beta', beta, (float), self.name) + self.atom_numbers = atom_numbers + self.excluded_numbers = excluded_numbers + self.beta = beta + self.init_prim_io_names( + inputs=['uint_crd', 'sacler', 'charge', 'excluded_list_start', 'excluded_list', 'excluded_atom_numbers'], + outputs=['force']) + self.add_prim_attr('atom_numbers', self.atom_numbers) + self.add_prim_attr('excluded_numbers', self.excluded_numbers) + self.add_prim_attr('beta', self.beta) + + def infer_shape(self, uint_crd_shape, sacler_shape, charge_shape, excluded_list_start_shape, excluded_list_shape, + excluded_atom_numbers_shape): + cls_name = self.name + N = self.atom_numbers + validator.check_int(len(uint_crd_shape), 2, Rel.EQ, "uint_crd_dim", cls_name) + validator.check_int(len(sacler_shape), 1, Rel.EQ, "sacler_dim", cls_name) + validator.check_int(len(charge_shape), 1, Rel.EQ, "charge_dim", cls_name) + validator.check_int(len(excluded_list_start_shape), 1, Rel.EQ, "excluded_list_start_dim", cls_name) + validator.check_int(len(excluded_atom_numbers_shape), 1, Rel.EQ, "excluded_atom_numbers_dim", cls_name) + + validator.check_int(uint_crd_shape[0], N, Rel.EQ, "uint_crd_shape[0]", cls_name) + validator.check_int(uint_crd_shape[1], 3, Rel.EQ, "uint_crd_shape[1]", cls_name) + validator.check_int(sacler_shape[0], 3, Rel.EQ, "sacler_shape", cls_name) + validator.check_int(charge_shape[0], N, Rel.EQ, "charge_shape", cls_name) + validator.check_int(excluded_list_start_shape[0], N, Rel.EQ, "excluded_list_start_shape", cls_name) + validator.check_int(excluded_atom_numbers_shape[0], N, Rel.EQ, "excluded_atom_numbers_shape", cls_name) + return uint_crd_shape + + def infer_dtype(self, uint_crd_type, sacler_type, charge_type, excluded_list_start_type, excluded_list_type, + excluded_atom_numbers_type): + validator.check_tensor_dtype_valid('sacler', sacler_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('uint_crd', uint_crd_type, [mstype.uint32], self.name) + validator.check_tensor_dtype_valid('charge', charge_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('excluded_list_start', excluded_list_start_type, [mstype.int32], + self.name) + validator.check_tensor_dtype_valid('excluded_list', excluded_list_type, [mstype.int32], + self.name) + validator.check_tensor_dtype_valid('excluded_atom_numbers', excluded_atom_numbers_type, [mstype.int32], + self.name) + return charge_type + + +class PMEEnergy(PrimitiveWithInfer): + """ + Calculate the Coulumb energy of the system using PME method. + + .. math:: + + E = sum_{ij} q_iq_j/r_{ij} + + Args: + atom_numbers(int32): the number of atoms, N. + excluded_numbers(int32): the length of excluded list, E. + beta(float32): the PME beta parameter, determined by the + non-bond cutoff value and simulation precision tolerance. + fftx(int32): the number of points for Fourier transform in dimension X. + ffty(int32): the number of points for Fourier transform in dimension Y. + fftz(int32): the number of points for Fourier transform in dimension Z. + box_length_0(float32): the value of boxlength idx 0 + box_length_1(float32): the value of boxlength idx 1 + box_length_2(float32): the value of boxlength idx 2 + + + Inputs: + - **uint_crd** (Tensor, uint32) - [N, 3], the unsigned int coordinates value of each atom. + - **charge** (Tensor, float32) - [N,], the charge carried by 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. + - **scaler** (Tensor, float32) - [3,], the scale factor between real space + coordinates and its unsigned int value. + - **excluded_list_start** (Tensor, int32) - [N,], the start excluded index + in excluded list for each atom. + - **excluded_list** (Tensor, int32) - [E,], the contiguous join of excluded + list of each atom. E is the number of excluded atoms. + - **excluded_atom_numbers** (Tensor, int32) - [N,], the number of atom excluded + in excluded list for each atom. + + Outputs: + - **reciprocal_ene** (float32) - the reciprocal term of PME energy. + - **self_ene** (float32) - the self term of PME energy. + - **direct_ene** (float32) - the direct term of PME energy. + - **correction_ene** (float32) - the correction term of PME energy. + + Supported Platforms: + ``GPU`` + """ + + @prim_attr_register + def __init__(self, atom_numbers, excluded_numbers, beta, fftx, ffty, fftz, box_length_0, box_length_1, + box_length_2): + validator.check_value_type('atom_numbers', atom_numbers, (int), self.name) + validator.check_value_type('excluded_numbers', excluded_numbers, (int), self.name) + validator.check_value_type('beta', beta, (float), self.name) + validator.check_value_type('fftx', fftx, (int), self.name) + validator.check_value_type('ffty', ffty, (int), self.name) + validator.check_value_type('fftz', fftz, (int), self.name) + self.atom_numbers = atom_numbers + self.excluded_numbers = excluded_numbers + self.beta = beta + self.fftx = fftx + self.ffty = ffty + self.fftz = fftz + self.box_length_0 = box_length_0 + self.box_length_1 = box_length_1 + self.box_length_2 = box_length_2 + self.init_prim_io_names( + inputs=['box_length', 'uint_crd', 'charge', 'nl_numbers', 'nl_serial', 'scaler', 'excluded_list_start', + 'excluded_list', 'excluded_atom_numbers'], + outputs=['reciprocal_ene', 'self_ene', 'direct_ene', 'correction_ene']) + self.add_prim_attr('atom_numbers', self.atom_numbers) + self.add_prim_attr('excluded_numbers', self.excluded_numbers) + self.add_prim_attr('beta', self.beta) + self.add_prim_attr('fftx', self.fftx) + self.add_prim_attr('ffty', self.ffty) + self.add_prim_attr('fftz', self.fftz) + self.add_prim_attr('box_length_0', self.box_length_0) + self.add_prim_attr('box_length_1', self.box_length_1) + self.add_prim_attr('box_length_2', self.box_length_2) + + def infer_shape(self, uint_crd, charge, nl_numbers, nl_serial, scaler, excluded_list_start, + excluded_list, excluded_atom_numbers): + cls_name = self.name + N = self.atom_numbers + validator.check_int(len(uint_crd), 2, Rel.EQ, "uint_crd_dim", cls_name) + validator.check_int(len(charge), 1, Rel.EQ, "charge_dim", cls_name) + validator.check_int(len(nl_numbers), 1, Rel.EQ, "nl_numbers_dim", cls_name) + validator.check_int(len(nl_serial), 2, Rel.LE, "nl_serial_dim", cls_name) + validator.check_int(len(excluded_list_start), 1, Rel.EQ, "excluded_list_start_dim", cls_name) + validator.check_int(len(excluded_atom_numbers), 1, Rel.EQ, "excluded_atom_numbers_dim", cls_name) + validator.check_int(len(excluded_list), 1, Rel.GE, "excluded_list", cls_name) + + validator.check_int(uint_crd[0], N, Rel.EQ, "uint_crd_shape[0]", cls_name) + validator.check_int(uint_crd[1], 3, Rel.EQ, "uint_crd_shape[1]", cls_name) + validator.check_int(charge[0], N, Rel.EQ, "charge_shape", cls_name) + validator.check_int(nl_numbers[0], N, Rel.EQ, "nl_numbers_shape[0]", cls_name) + validator.check_int(nl_serial[0], N, Rel.LE, "nl_serial_shape[0]", cls_name) + validator.check_int(nl_serial[1], 800, Rel.LE, "nl_serial_shape[1]", cls_name) + validator.check_int(excluded_list_start[0], N, Rel.EQ, "excluded_list_start_shape", cls_name) + validator.check_int(excluded_atom_numbers[0], N, Rel.EQ, "excluded_atom_numbers_shape", cls_name) + validator.check_int(excluded_list[0], 0, Rel.GE, "excluded_list_shape", cls_name) + return (1,), (1,), (1,), (1,) + + def infer_dtype(self, uint_crd, charge, nl_numbers, nl_serial, scaler, excluded_list_start, + excluded_list, excluded_atom_numbers): + validator.check_tensor_dtype_valid('uint_crd', uint_crd, [mstype.uint32], self.name) + validator.check_tensor_dtype_valid('charge', charge, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('nl_numbers', nl_numbers, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('nl_serial', nl_serial, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('scaler', scaler, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('excluded_list_start', excluded_list_start, [mstype.int32], + self.name) + validator.check_tensor_dtype_valid('excluded_list', excluded_list, [mstype.int32], + self.name) + validator.check_tensor_dtype_valid('excluded_atom_numbers', excluded_atom_numbers, [mstype.int32], + self.name) + return charge, charge, charge, charge + + +class LJEnergy(PrimitiveWithInfer): + """ + Calculate the Van der Waals interaction energy described by Lennard-Jones + potential for each atom. Assume the number of atoms is N, and the number + of Lennard-Jones types for all atoms is P, which means there will be + Q = P*(P+1)/2 types of possible Lennard-Jones interactions for all kinds + of atom pairs. + + + .. math:: + + dr = (x_a-x_b, y_a-y_b, z_a-z_b) + + .. math:: + E = A/|dr|^{12} - B/|dr|^{6} + + Agrs: + atom_numbers(int32): the number of atoms, N. + cutoff_square(float32): the square value of cutoff. + + Inputs: + - **uint_crd** (Tensor, uint32) - [N, 3], the unsigned int coordinate value of each atom. + - **LJtype** (Tensor, int32) - [N,], the Lennard-Jones type of each atom. + - **charge** (Tensor, float32) - [N,], the charge carried by each atom. + - **scaler** (Tensor, float32) - [3,], the scale factor between real + space coordinate and its unsigned int value. + - **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. + - **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. + Q is the number of atom pair. + + Outputs: + - **d_LJ_energy_atom** (Tensor, float32) - [N,], the Lennard-Jones potential energy of each atom. + - **d_LJ_energy_sum** (float32), the sum of Lennard-Jones potential energy of each atom. + + Supported Platforms: + ``GPU`` + """ + + @prim_attr_register + def __init__(self, atom_numbers, cutoff_square): + validator.check_value_type('atom_numbers', atom_numbers, (int), self.name) + validator.check_value_type('cutoff_square', cutoff_square, (float), self.name) + self.atom_numbers = atom_numbers + self.cutoff_square = cutoff_square + self.init_prim_io_names( + inputs=['uint_crd', 'LJtype', 'charge', 'scaler', 'nl_numbers', 'nl_serial', 'd_LJ_A', 'd_LJ_B'], + outputs=['d_LJ_energy_atom']) + self.add_prim_attr('atom_numbers', self.atom_numbers) + 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): + cls_name = self.name + N = self.atom_numbers + Q = d_LJ_A[0] + validator.check_int(len(uint_crd), 2, Rel.EQ, "uint_crd_dim", cls_name) + validator.check_int(len(LJtype), 1, Rel.EQ, "LJtype_dim", cls_name) + validator.check_int(len(charge), 1, Rel.EQ, "charge_dim", cls_name) + validator.check_int(len(scaler), 1, Rel.EQ, "scaler_dim", cls_name) + validator.check_int(len(nl_numbers), 1, Rel.EQ, "nl_numbers_dim", cls_name) + validator.check_int(len(nl_serial), 2, Rel.EQ, "nl_serial_dim", cls_name) + validator.check_int(len(scaler), 1, Rel.EQ, "scaler_dim", cls_name) + validator.check_int(len(d_LJ_B), 1, Rel.EQ, "d_LJ_B_dim", cls_name) + + validator.check_int(uint_crd[0], N, Rel.EQ, "uint_crd_shape[0]", cls_name) + validator.check_int(uint_crd[1], 3, Rel.EQ, "uint_crd_shape[1]", cls_name) + validator.check_int(LJtype[0], N, Rel.EQ, "LJtype_shape", cls_name) + validator.check_int(charge[0], N, Rel.EQ, "charge_shape", cls_name) + validator.check_int(scaler[0], 3, Rel.EQ, "scaler_shape", cls_name) + validator.check_int(nl_numbers[0], N, Rel.EQ, "nl_numbers_shape", cls_name) + validator.check_int(nl_serial[0], N, Rel.EQ, "nl_serial_shape[0]", cls_name) + validator.check_int(nl_serial[1], 800, Rel.LE, "nl_serial_shape[1]", cls_name) + validator.check_int(scaler[0], 3, Rel.EQ, "scaler_shape", cls_name) + validator.check_int(d_LJ_B[0], Q, Rel.EQ, "d_LJ_B_shape[0]", cls_name) + return charge + + def infer_dtype(self, uint_crd, LJtype, charge, scaler, nl_numbers, nl_serial, d_LJ_A, d_LJ_B): + validator.check_tensor_dtype_valid('uint_crd', uint_crd, [mstype.uint32], self.name) + validator.check_tensor_dtype_valid('LJtype', LJtype, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('charge', charge, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('scaler', scaler, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('nl_numbers', nl_numbers, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('nl_serial', nl_serial, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('d_LJ_A', d_LJ_A, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('d_LJ_B', d_LJ_B, [mstype.float32], self.name) + return charge + + +class LJForce(PrimitiveWithInfer): + """ + Calculate the Van der Waals interaction force described by Lennard-Jones + potential energy for each atom. + + .. math:: + + dr = (x_a-x_b, y_a-y_b, z_a-z_b) + + .. math:: + + F = (-12*A/|dr|^{14} + 6*B/|dr|^{8}) * dr + + Agrs: + atom_numbers(int32): the number of atoms, N. + cutoff_square(float32): the square value of cutoff. + + Inputs: + - **uint_crd** (Tensor, uint32) - [N, 3], the unsigned int coordinate value of each atom. + - **LJtype** (Tensor, int32) - [N,], the Lennard-Jones type of each atom. + - **charge** (Tensor, float32) - [N,], the charge carried by each atom. + - **scaler** (Tensor, float32) - [3,], the scale factor between real + space coordinate and its unsigned int value. + - **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. + - **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. + Q is the number of atom pair. + + outputs: + - **frc** (Tensor, float32) - [N, 3], the force felt by each atom. + + Supported Platforms: + ``GPU`` + """ + + @prim_attr_register + def __init__(self, atom_numbers, cutoff_square): + validator.check_value_type('atom_numbers', atom_numbers, (int), self.name) + validator.check_value_type('cutoff_square', cutoff_square, (float), self.name) + self.atom_numbers = atom_numbers + self.cutoff_square = cutoff_square + self.init_prim_io_names( + inputs=['uint_crd', 'LJtype', 'charge', 'scaler', 'nl_numbers', 'nl_serial', 'd_LJ_A', 'd_LJ_B'], + outputs=['frc']) + self.add_prim_attr('atom_numbers', self.atom_numbers) + 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): + cls_name = self.name + N = self.atom_numbers + Q = d_LJ_A[0] + validator.check_int(len(uint_crd), 2, Rel.EQ, "uint_crd_dim", cls_name) + validator.check_int(len(LJtype), 1, Rel.EQ, "LJtype_dim", cls_name) + validator.check_int(len(charge), 1, Rel.EQ, "charge_dim", cls_name) + validator.check_int(len(scaler), 1, Rel.EQ, "scaler_dim", cls_name) + validator.check_int(len(nl_numbers), 1, Rel.EQ, "nl_numbers_dim", cls_name) + validator.check_int(len(nl_serial), 2, Rel.EQ, "nl_serial_dim", cls_name) + validator.check_int(len(scaler), 1, Rel.EQ, "scaler_dim", cls_name) + validator.check_int(len(d_LJ_B), 1, Rel.EQ, "d_LJ_B_dim", cls_name) + + validator.check_int(uint_crd[0], N, Rel.EQ, "uint_crd_shape[0]", cls_name) + validator.check_int(uint_crd[1], 3, Rel.EQ, "uint_crd_shape[1]", cls_name) + validator.check_int(LJtype[0], N, Rel.EQ, "LJtype_shape", cls_name) + validator.check_int(charge[0], N, Rel.EQ, "charge_shape", cls_name) + validator.check_int(scaler[0], 3, Rel.EQ, "scaler_shape", cls_name) + validator.check_int(nl_numbers[0], N, Rel.EQ, "nl_numbers_shape", cls_name) + validator.check_int(nl_serial[0], N, Rel.EQ, "nl_serial_shape[0]", cls_name) + validator.check_int(nl_serial[1], 800, Rel.LE, "nl_serial_shape[1]", cls_name) + validator.check_int(scaler[0], 3, Rel.EQ, "scaler_shape", cls_name) + validator.check_int(d_LJ_B[0], Q, Rel.EQ, "d_LJ_B_shape[0]", cls_name) + return uint_crd + + def infer_dtype(self, uint_crd, LJtype, charge, scaler, nl_numbers, nl_serial, d_LJ_A, d_LJ_B): + validator.check_tensor_dtype_valid('uint_crd', uint_crd, [mstype.uint32], self.name) + validator.check_tensor_dtype_valid('LJtype', LJtype, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('charge', charge, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('scaler', scaler, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('nl_numbers', nl_numbers, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('nl_serial', nl_serial, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('d_LJ_A', d_LJ_A, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('d_LJ_B', d_LJ_B, [mstype.float32], self.name) + return charge + + +class LJForceWithPMEDirectForce(PrimitiveWithInfer): + """ + Calculate the Lennard-Jones force and PME direct force together. + + The calculation formula of Lennard-Jones part is the same as operator + LJForce(), and the PME direct part is within PME method. + + Agrs: + atom_numbers(int32): the number of atoms, N. + cutoff_square(float32): the square value of cutoff. + pme_beta(float32): PME beta parameter, same as operator PMEReciprocalForce(). + + Inputs: + - **uint_crd** (Tensor, uint32) - [N, 3], the unsigned int coordinate value of each atom. + - **LJtype** (Tensor, int32) - [N,], the Lennard-Jones type of each atom. + - **charge** (Tensor, float32) - [N,], the charge carried by each atom. + - **scaler** (Tensor, float32) - [3,], the scale factor between real + space coordinate and its unsigned int value. + - **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. + - **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. + Q is the number of atom pair. + + Outputs: + - **frc** (Tensor, float32), [N, 3], the force felt by each atom. + + Supported Platforms: + ``GPU`` + """ + + @prim_attr_register + def __init__(self, atom_numbers, cutoff, pme_beta): + validator.check_value_type('atom_numbers', atom_numbers, (int), self.name) + validator.check_value_type('cutoff', cutoff, (float), self.name) + validator.check_value_type('pme_beta', pme_beta, (float), self.name) + self.atom_numbers = atom_numbers + self.cutoff = cutoff + self.pme_beta = pme_beta + self.init_prim_io_names( + inputs=['uint_crd', 'LJtype', 'charge', 'scaler', 'nl_numbers', 'nl_serial', 'd_LJ_A', 'd_LJ_B'], + outputs=['frc']) + self.add_prim_attr('atom_numbers', self.atom_numbers) + self.add_prim_attr('cutoff', self.cutoff) + 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): + cls_name = self.name + N = self.atom_numbers + Q = d_LJ_A[0] + validator.check_int(len(uint_crd), 2, Rel.EQ, "uint_crd_dim", cls_name) + validator.check_int(len(LJtype), 1, Rel.EQ, "LJtype_dim", cls_name) + validator.check_int(len(charge), 1, Rel.EQ, "charge_dim", cls_name) + validator.check_int(len(scaler), 1, Rel.EQ, "scaler_dim", cls_name) + validator.check_int(len(nl_numbers), 1, Rel.EQ, "nl_numbers_dim", cls_name) + validator.check_int(len(nl_serial), 2, Rel.EQ, "nl_serial_dim", cls_name) + validator.check_int(len(scaler), 1, Rel.EQ, "scaler_dim", cls_name) + validator.check_int(len(d_LJ_B), 1, Rel.EQ, "d_LJ_B_dim", cls_name) + + validator.check_int(uint_crd[0], N, Rel.EQ, "uint_crd_shape[0]", cls_name) + validator.check_int(uint_crd[1], 3, Rel.EQ, "uint_crd_shape[1]", cls_name) + validator.check_int(LJtype[0], N, Rel.EQ, "LJtype_shape", cls_name) + validator.check_int(charge[0], N, Rel.EQ, "charge_shape", cls_name) + validator.check_int(scaler[0], 3, Rel.EQ, "scaler_shape", cls_name) + validator.check_int(nl_numbers[0], N, Rel.EQ, "nl_numbers_shape", cls_name) + validator.check_int(nl_serial[0], N, Rel.EQ, "nl_serial_shape[0]", cls_name) + validator.check_int(nl_serial[1], 800, Rel.LE, "nl_serial_shape[1]", cls_name) + validator.check_int(scaler[0], 3, Rel.EQ, "scaler_shape", cls_name) + validator.check_int(d_LJ_B[0], Q, Rel.EQ, "d_LJ_B_shape[0]", cls_name) + return uint_crd + + def infer_dtype(self, uint_crd, LJtype, charge, scaler, nl_numbers, nl_serial, d_LJ_A, d_LJ_B): + validator.check_tensor_dtype_valid('uint_crd', uint_crd, [mstype.uint32], self.name) + validator.check_tensor_dtype_valid('LJtype', LJtype, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('charge', charge, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('scaler', scaler, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('nl_numbers', nl_numbers, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('nl_serial', nl_serial, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('d_LJ_A', d_LJ_A, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('d_LJ_B', d_LJ_B, [mstype.float32], self.name) + return charge + + +class GetCenterOfGeometry(PrimitiveWithInfer): + """ + Get Center Of Geometry. + + Supported Platforms: + ``GPU`` + """ + + @prim_attr_register + 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_inverse = center_numbers_inverse + self.add_prim_attr('center_numbers', self.center_numbers) + self.add_prim_attr('center_numbers_inverse', self.center_numbers_inverse) + self.init_prim_io_names( + inputs=['center_atoms', 'crd_f'], + outputs=['center_of_geometry_f']) + + def infer_shape(self, center_atoms_shape, crd_f_shape): + cls_name = self.name + N = self.center_numbers + 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[0]", cls_name) + validator.check_int(crd_f_shape[1], 3, Rel.EQ, "crd_f[1]", cls_name) + return [3,] + + def infer_dtype(self, center_atoms_dtype, crd_f_dtype): + validator.check_tensor_dtype_valid('center_atoms', center_atoms_dtype, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('crd_f', crd_f_dtype, [mstype.float32], self.name) + + return crd_f_dtype + + +class MDTemperature(PrimitiveWithInfer): + """ + Calculate the MD Temperature. + + Calculate the temperature. + + Supported Platforms: + ``GPU`` + """ + + @prim_attr_register + def __init__(self, residue_numbers, atom_numbers): + assert isinstance(residue_numbers, int) + assert isinstance(atom_numbers, int) + self.residue_numbers = residue_numbers + self.atom_numbers = atom_numbers + self.add_prim_attr('residue_numbers', self.residue_numbers) + self.add_prim_attr('atom_numbers', self.atom_numbers) + self.init_prim_io_names( + inputs=['start', 'end', 'atom_vel_f', 'atom_mass'], + outputs=['ek']) + + def infer_shape(self, start_shape, end_shape, atom_vel_f_shape, atom_mass_shape): + cls_name = self.name + N = self.residue_numbers + M = self.atom_numbers + validator.check_int(len(start_shape), 1, Rel.EQ, "start", cls_name) + validator.check_int(start_shape[0], N, Rel.EQ, "end", cls_name) + validator.check_int(len(end_shape), 1, Rel.EQ, "start", cls_name) + validator.check_int(end_shape[0], N, Rel.EQ, "end", cls_name) + validator.check_int(atom_vel_f_shape[0], M, Rel.EQ, "atom_vel_f", cls_name) + validator.check_int(atom_vel_f_shape[1], 3, Rel.EQ, "atom_vel_f", cls_name) + validator.check_int(len(atom_mass_shape), 1, Rel.EQ, "atom_mass", cls_name) + validator.check_int(atom_mass_shape[0], M, Rel.EQ, "atom_mass", cls_name) + return [N,] + + def infer_dtype(self, start_dtype, end_dtype, atom_vel_f_dtype, atom_mass_dtype): + validator.check_tensor_dtype_valid('start', start_dtype, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('end', end_dtype, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('atom_vel_f', atom_vel_f_dtype, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('atom_mass', atom_mass_dtype, [mstype.float32], self.name) + return atom_mass_dtype + + +class NeighborListUpdate(PrimitiveWithInfer): + """ + Update (or construct if first time) the Verlet neighbor list for the + calculation of short-ranged force. Assume the number of atoms is N, + the number of grids divided is G, the maximum number of atoms in one + grid is M, the maximum number of atoms in single atom's neighbor list + is L, and the number of total atom in excluded list is E. + + Args: + grid_numbers(int32): the total number of grids divided. + not_first_time(int32): whether to construct the neighbor + list first time or not. + Nxy(int32): the total number of grids divided in xy plane. + excluded_atom_numbers(int32): the total atom numbers in the excluded list. + cutoff(float32): the cutoff distance for short-range force calculation. + skin(float32): the overflow value of cutoff to maintain a neighbor list. + cutoff_square(float32): the suqare value of cutoff. + half_skin_square(float32): skin*skin/4, indicates the maximum + square value of the distance atom allowed to move between two updates. + cutoff_with_skin(float32): cutoff + skin, indicates the + radius of the neighbor list for each atom. + half_cutoff_with_skin(float32): cutoff_with_skin/2. + cutoff_with_skin_square(float32): the square value of cutoff_with_skin. + refresh_interval(int32): the number of iteration steps between two updates of neighbor list. + max_atom_in_grid_numbers(int32): the maximum number of atoms in one grid. + + Inputs: + - **atom_numbers_in_grid_bucket** (Tensor, int32) - [G,], the number of atoms in each grid bucket. + - **bucket** (Tensor, int32) - (Tensor,int32) - [G, M], the atom indices in each grid bucket. + - **crd** (Tensor, float32) - [N,], the coordinates of each atom. + - **box_length** (Tensor, float32) - [3,], the length of 3 dimensions of the simulation box. + - **grid_N** (Tensor, int32) - [3,], the number of grids divided of 3 dimensions of the simulation box. + - **grid_length_inverse** (float32) - the inverse value of grid length. + - **atom_in_grid_serial** (Tensor, int32) - [N,], the grid index for each atom. + - **old_crd** (Tensor, float32) - [N, 3], the coordinates before update of each atom. + - **crd_to_uint_crd_cof** (Tensor, float32) - [3,], the scale factor + between the unsigned int value and the real space coordinates. + - **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. + 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_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 + the real space coordinates and the unsigned int value. + - **excluded_list_start** (Tensor, int32) - [N,], the start excluded index in excluded list for each atom. + - **excluded_numbers** (Tensor, int32) - [N,], the number of atom excluded in excluded list for each atom. + - **excluded_list** (Tensor, int32) - [E,], the contiguous join of excluded list of each atom. + - **need_refresh_flag** (Tensor, int32) - [N,], whether the neighbor list of each atom need update or not. + - **refresh_count** (Tensor, int32) - [1,], count how many iteration steps have passed since last update. + + Outputs: + - **res** (float32) + + Supported Platforms: + ``GPU`` + """ + + @prim_attr_register + def __init__(self, grid_numbers, atom_numbers, not_first_time, Nxy, excluded_atom_numbers, + cutoff_square, half_skin_square, cutoff_with_skin, half_cutoff_with_skin, cutoff_with_skin_square, + refresh_interval=20, cutoff=10.0, skin=2.0, max_atom_in_grid_numbers=64, max_neighbor_numbers=800): + self.grid_numbers = grid_numbers + self.atom_numbers = atom_numbers + self.refresh_interval = refresh_interval + self.not_first_time = not_first_time + self.cutoff = cutoff + self.skin = skin + self.max_atom_in_grid_numbers = max_atom_in_grid_numbers + self.Nxy = Nxy + self.excluded_atom_numbers = excluded_atom_numbers + self.cutoff_square = cutoff_square + self.half_skin_square = half_skin_square + self.cutoff_with_skin = cutoff_with_skin + self.half_cutoff_with_skin = half_cutoff_with_skin + self.cutoff_with_skin_square = cutoff_with_skin_square + self.max_neighbor_numbers = max_neighbor_numbers + self.init_prim_io_names( + inputs=['atom_numbers_in_grid_bucket', 'bucket', 'crd', 'box_length', 'grid_N', 'grid_length_inverse', + 'atom_in_grid_serial', 'old_crd', 'crd_to_uint_crd_cof', 'uint_crd', 'gpointer', 'nl_atom_numbers', + 'nl_atom_serial', 'uint_dr_to_dr_cof', 'excluded_list_start', 'excluded_list', 'excluded_numbers', + 'need_refresh_flag', 'refresh_count'], outputs=['res']) + + self.add_prim_attr('grid_numbers', self.grid_numbers) + self.add_prim_attr('atom_numbers', self.atom_numbers) + self.add_prim_attr('refresh_interval', self.refresh_interval) + self.add_prim_attr('not_first_time', self.not_first_time) + self.add_prim_attr('cutoff', self.cutoff) + self.add_prim_attr('skin', self.skin) + self.add_prim_attr('max_atom_in_grid_numbers', self.max_atom_in_grid_numbers) + self.add_prim_attr('Nxy', self.Nxy) + self.add_prim_attr('excluded_atom_numbers', self.excluded_atom_numbers) + self.add_prim_attr('cutoff_square', self.cutoff_square) + self.add_prim_attr('half_skin_square', self.half_skin_square) + self.add_prim_attr('cutoff_with_skin', self.cutoff_with_skin) + self.add_prim_attr('half_cutoff_with_skin', self.half_cutoff_with_skin) + self.add_prim_attr('cutoff_with_skin_square', self.cutoff_with_skin_square) + + def infer_shape(self, atom_numbers_in_grid_bucket_shape, bucket_shape, crd_shape, box_length_shape, grid_N_shape, + grid_length_inverse_shape, atom_in_grid_serial_shape, old_crd_shape, crd_to_uint_crd_cof_shape, + uint_crd_shape, gpointer_shape, nl_atom_numbers_shape, nl_atom_serial_shape, + uint_dr_to_dr_cof_shape, excluded_list_start_shape, excluded_list_shape, excluded_numbers_shape, + need_refresh_flag_shape, refresh_count_shape): + assert len(atom_numbers_in_grid_bucket_shape) == 1 + assert len(bucket_shape) == 2 + assert len(crd_shape) == 2 + assert len(box_length_shape) == 1 + assert len(grid_N_shape) == 1 + assert len(grid_length_inverse_shape) == 1 + assert len(atom_in_grid_serial_shape) == 1 + assert len(old_crd_shape) == 2 + assert len(crd_to_uint_crd_cof_shape) == 1 + assert len(uint_crd_shape) == 2 + assert len(gpointer_shape) == 2 + assert len(nl_atom_numbers_shape) == 1 + assert len(nl_atom_serial_shape) == 2 + assert len(uint_dr_to_dr_cof_shape) == 1 + assert len(excluded_list_start_shape) == 1 + assert len(excluded_list_shape) == 1 + assert len(excluded_numbers_shape) == 1 + assert len(need_refresh_flag_shape) == 1 + + validator.check_int(atom_numbers_in_grid_bucket_shape[0], self.grid_numbers, Rel.EQ, + "atom_numbers_in_grid_bucket", self.name) + validator.check_int(bucket_shape[0], self.grid_numbers, Rel.EQ, "bucket", self.name) + validator.check_int(bucket_shape[1], self.max_atom_in_grid_numbers, Rel.EQ, "bucket", self.name) + validator.check_int(crd_shape[0], self.atom_numbers, Rel.EQ, "crd", self.name) + validator.check_int(crd_shape[1], 3, Rel.EQ, "crd", self.name) + validator.check_int(box_length_shape[0], 3, Rel.EQ, "box_length", self.name) + validator.check_int(grid_N_shape[0], 3, Rel.EQ, "grid_N", self.name) + validator.check_int(grid_length_inverse_shape[0], 3, Rel.EQ, "grid_length_inverse", self.name) + validator.check_int(atom_in_grid_serial_shape[0], self.atom_numbers, Rel.EQ, "atom_in_grid_serial", + self.name) + validator.check_int(old_crd_shape[0], self.atom_numbers, Rel.EQ, "old_crd", self.name) + validator.check_int(old_crd_shape[1], 3, Rel.EQ, "old_crd", self.name) + validator.check_int(crd_to_uint_crd_cof_shape[0], 3, Rel.EQ, "crd_to_uint_crd_cof", self.name) + validator.check_int(uint_crd_shape[0], self.atom_numbers, Rel.EQ, "uint_crd", self.name) + validator.check_int(uint_crd_shape[1], 3, Rel.EQ, "uint_crd", self.name) + validator.check_int(gpointer_shape[0], self.grid_numbers, Rel.EQ, "gpointer", self.name) + validator.check_int(gpointer_shape[1], 125, Rel.EQ, "gpointer", self.name) + validator.check_int(nl_atom_numbers_shape[0], self.atom_numbers, Rel.EQ, "nl_atom_numbers", self.name) + validator.check_int(nl_atom_serial_shape[0], self.atom_numbers, Rel.EQ, "nl_atom_serial", self.name) + validator.check_int(nl_atom_serial_shape[1], self.max_neighbor_numbers, Rel.EQ, "nl_atom_serial", + self.name) + validator.check_int(uint_dr_to_dr_cof_shape[0], 3, Rel.EQ, "uint_dr_to_dr_cof", self.name) + validator.check_int(excluded_list_start_shape[0], self.atom_numbers, Rel.EQ, "excluded_list_start", + self.name) + validator.check_int(excluded_list_shape[0], self.excluded_atom_numbers, Rel.EQ, "excluded_list", + self.name) + validator.check_int(excluded_numbers_shape[0], self.atom_numbers, Rel.EQ, "excluded_numbers", self.name) + validator.check_int(need_refresh_flag_shape[0], 1, Rel.EQ, "need_refresh_flag", self.name) + + return [1,] + + def infer_dtype(self, atom_numbers_in_grid_bucket_dtype, bucket_dtype, crd_dtype, box_length_dtype, grid_N_dtype, + grid_length_inverse_dtype, atom_in_grid_serial_dtype, old_crd_dtype, crd_to_uint_crd_cof_dtype, + uint_crd_dtype, gpointer_dtype, nl_atom_numbers_dtype, nl_atom_serial_dtype, + uint_dr_to_dr_cof_dtype, excluded_list_start_dtype, excluded_list_dtype, excluded_numbers_dtype, + need_refresh_flag_dtype, refresh_count_dtype): + validator.check_tensor_dtype_valid('atom_numbers_in_grid_bucket', atom_numbers_in_grid_bucket_dtype, + [mstype.int32], self.name) + validator.check_tensor_dtype_valid('bucket', bucket_dtype, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('crd', crd_dtype, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('box_length', box_length_dtype, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('grid_N', grid_N_dtype, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('grid_length_inverse', grid_length_inverse_dtype, [mstype.float32], + self.name) + validator.check_tensor_dtype_valid('atom_in_grid_serial', atom_in_grid_serial_dtype, [mstype.int32], + self.name) + validator.check_tensor_dtype_valid('old_crd', old_crd_dtype, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('crd_to_uint_crd_cof', crd_to_uint_crd_cof_dtype, [mstype.float32], + self.name) + validator.check_tensor_dtype_valid('uint_crd', uint_crd_dtype, [mstype.uint32], self.name) + validator.check_tensor_dtype_valid('gpointer', gpointer_dtype, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('nl_atom_numbers', nl_atom_numbers_dtype, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('nl_atom_serial', nl_atom_serial_dtype, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('uint_dr_to_dr_cof', uint_dr_to_dr_cof_dtype, [mstype.float32], + self.name) + validator.check_tensor_dtype_valid('excluded_list_start', excluded_list_start_dtype, [mstype.int32], + self.name) + validator.check_tensor_dtype_valid('excluded_list', excluded_list_dtype, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('excluded_numbers', excluded_numbers_dtype, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('need_refresh_flag', need_refresh_flag_dtype, [mstype.int32], + self.name) + + return mstype.float32 + + +class MDIterationLeapFrogWithRF(PrimitiveWithInfer): + """ + One step of classical leap frog algorithm to solve the finite difference + Hamiltonian equations of motion for certain system, using Langevin dynamics + with Liu's thermostat scheme. Assume the number of atoms is N and the target + control temperature is T. + + Detailed iteration formula can be found in this paper: A unified thermostat + scheme for efficient configurational sampling for classical/quantum canonical + ensembles via molecular dynamics. DOI: 10.1063/1.4991621. + + Inputs: + - **float4_numbers** (int32) - total length to store random numbers. + - **atom_numbers** (int32) - the number of atoms N. + - **dt** (float32) - time step for finite difference. + - **half_dt** (float32) - half of time step for finite difference. + - **exp_gamma** (float32) - parameter in Liu's dynamic, equals + exp(-gamma_ln * dt), where gamma_ln is the firction factor in Langvin + dynamics. + - **max_velocity** (float32) - the upper limit of velocity, when the + veclocity overflows, scale it to the upper limit. + - **is_max_velocity** (int32) - whether the max velocity control is + open or not. + + - **mass_inverse** (Tensor, float32) - [N,], the inverse value of + mass of each atom. + - **sqrt_mass** (Tensor, float32) - [N,], the inverse square root value + of effect mass in Liu's dynamics of each atom. + - **vel** (Tensor, float32) - [N, 3], the velocity of each atom. + - **crd** (Tensor, float32) - [N, 3], the coordinate of each atom. + - **frc** (Tensor, float32) - [N, 3], the force felt by each atom. + - **acc** (Tensor, float32) - [N, 3], the acceleration of each atom. + - **random force** (Tensor, float32) - [N, 3], the random forces. + + Outputs: + - **res** (float32) + + Supported Platforms: + ``GPU`` + Examples: + """ + + @prim_attr_register + def __init__(self, float4_numbers, atom_numbers, half_dt, dt, exp_gamma, is_max_velocity, max_velocity): + validator.check_value_type('float4_numbers', float4_numbers, (int), self.name) + validator.check_value_type('atom_numbers', atom_numbers, (int), self.name) + validator.check_value_type('half_dt', half_dt, (float), self.name) + validator.check_value_type('dt', dt, (float), self.name) + validator.check_value_type('exp_gamma', exp_gamma, (float), self.name) + validator.check_value_type('is_max_velocity', is_max_velocity, (int), self.name) + validator.check_value_type('max_velocity', max_velocity, (float), self.name) + self.float4_numbers = float4_numbers + self.atom_numbers = atom_numbers + self.half_dt = half_dt + self.dt = dt + self.exp_gamma = exp_gamma + self.is_max_velocity = is_max_velocity + self.max_velocity = max_velocity + + self.init_prim_io_names( + inputs=['mass_inverse', 'sqrt_mass', 'vel_in', 'crd_in', 'frc_in', 'acc_in', 'random_force'], + outputs=['res']) + self.add_prim_attr('float4_numbers', self.float4_numbers) + self.add_prim_attr('atom_numbers', self.atom_numbers) + self.add_prim_attr('half_dt', self.half_dt) + self.add_prim_attr('dt', self.dt) + self.add_prim_attr('exp_gamma', self.exp_gamma) + self.add_prim_attr('is_max_velocity', self.is_max_velocity) + self.add_prim_attr('max_velocity', self.max_velocity) + + def infer_shape(self, mass_inverse_shape, sqrt_mass_shape, vel_in_shape, crd_in_shape, frc_in_shape, acc_in_shape, + random_force_shape): + N = self.atom_numbers + validator.check_int(len(mass_inverse_shape), 1, Rel.EQ, "mass_inverse", self.name) + validator.check_int(len(sqrt_mass_shape), 1, Rel.EQ, "mass_inverse", self.name) + validator.check_int(mass_inverse_shape[0], N, Rel.EQ, "mass_inverse", self.name) + validator.check_int(sqrt_mass_shape[0], N, Rel.EQ, "mass_inverse", self.name) + validator.check_int(vel_in_shape[0], N, Rel.EQ, "vel_in", self.name) + validator.check_int(vel_in_shape[1], 3, Rel.EQ, "vel_in", self.name) + validator.check_int(crd_in_shape[0], N, Rel.EQ, "crd_in", self.name) + validator.check_int(crd_in_shape[1], 3, Rel.EQ, "crd_in", self.name) + validator.check_int(frc_in_shape[0], N, Rel.EQ, "frc_in", self.name) + validator.check_int(frc_in_shape[1], 3, Rel.EQ, "frc_in", self.name) + validator.check_int(acc_in_shape[0], N, Rel.EQ, "acc_in", self.name) + validator.check_int(acc_in_shape[1], 3, Rel.EQ, "acc_in", self.name) + validator.check_int(random_force_shape[0], N, Rel.EQ, "random_force", self.name) + validator.check_int(random_force_shape[1], 3, Rel.EQ, "random_force", self.name) + + return [1,] + + def infer_dtype(self, mass_inverse_dtype, sqrt_mass_dtype, vel_in_dtype, crd_in_dtype, frc_in_dtype, acc_in_dtype, + rf_dtype): + validator.check_tensor_dtype_valid('mass_inverse', mass_inverse_dtype, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('sqrt_mass', sqrt_mass_dtype, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('vel_in', vel_in_dtype, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('crd_in', crd_in_dtype, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('frc_in', frc_in_dtype, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('acc_in', acc_in_dtype, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('rf', rf_dtype, [mstype.float32], self.name) + return mstype.float32 + + +class MDIterationLeapFrogLiujian(PrimitiveWithInfer): + """ + One step of classical leap frog algorithm to solve the finite difference + Hamiltonian equations of motion for certain system, using Langevin dynamics + with Liu's thermostat scheme. Assume the number of atoms is N and the target + control temperature is T. + + Detailed iteration formula can be found in this paper: A unified thermostat + scheme for efficient configurational sampling for classical/quantum canonical + ensembles via molecular dynamics. DOI: 10.1063/1.4991621. + + Inputs: + - **atom_numbers** (int32) - the number of atoms N. + - **dt** (float32) - time step for finite difference. + - **half_dt** (float32) - half of time step for finite difference. + - **exp_gamma** (float32) - parameter in Liu's dynamic, equals + exp(-gamma_ln * dt), where gamma_ln is the firction factor in Langvin + dynamics. + + - **inverse_mass** (Tensor, float32) - [N,], the inverse value of + mass of each atom. + - **sqrt_mass_inverse** (Tensor, float32) - [N,], the inverse square root value + of effect mass in Liu's dynamics of each atom. + - **vel** (Tensor, float32) - [N, 3], the velocity of each atom. + - **crd** (Tensor, float32) - [N, 3], the coordinate of each atom. + - **frc** (Tensor, float32) - [N, 3], the force felt by each atom. + - **acc** (Tensor, float32) - [N, 3], the acceleration of each atom. + - **rand_state** (Tensor, float32) - [math.ceil(atom_numbers * 3.0 / 4.0) * 16, ], random state to generate + random force. + - **rand_frc** (Tensor, float32) - [N, 3], the random forces. + + Outputs: + - **output** (float32) + + Supported Platforms: + ``GPU`` + Examples: + """ + + @prim_attr_register + def __init__(self, atom_numbers, half_dt, dt, exp_gamma): + self.atom_numbers = atom_numbers + self.half_dt = half_dt + self.dt = dt + self.exp_gamma = exp_gamma + + self.add_prim_attr('atom_numbers', self.atom_numbers) + self.add_prim_attr('half_dt', self.half_dt) + self.add_prim_attr('dt', self.dt) + self.add_prim_attr('exp_gamma', self.exp_gamma) + self.init_prim_io_names( + inputs=['inverse_mass', 'sqrt_mass_inverse', 'vel', 'crd', 'frc', 'acc', 'rand_state', 'rand_frc'], + outputs=['output']) + + def infer_shape(self, inverse_mass, sqrt_mass_inverse, vel, crd, frc, acc, rand_state, rand_frc): + N = self.atom_numbers + validator.check_int(len(inverse_mass), 1, Rel.EQ, "inverse_mass", self.name) + validator.check_int(len(sqrt_mass_inverse), 1, Rel.EQ, "sqrt_mass_inverse", self.name) + validator.check_int(inverse_mass[0], N, Rel.EQ, "inverse_mass", self.name) + validator.check_int(sqrt_mass_inverse[0], N, Rel.EQ, "sqrt_mass_inverse", self.name) + return [self.atom_numbers, 3] + + def infer_dtype(self, inverse_mass, sqrt_mass_inverse, vel, crd, frc, acc, rand_state, rand_frc): + validator.check_tensor_dtype_valid('inverse_mass', inverse_mass, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('sqrt_mass_inverse', sqrt_mass_inverse, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('vel', vel, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('crd', crd, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('frc', frc, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('acc', acc, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('rand_frc', rand_frc, [mstype.float32], self.name) + return mstype.float32 + + +class CrdToUintCrd(PrimitiveWithInfer): + """ + Convert FP32 coordinate to Uint32 coordinate. + + Inputs: + - **atom_numbers** (int32) - the number of atoms N. + + - **crd_to_uint_crd_cof** (Tensor, float32) - [3,], the . + - **crd** (Tensor, float32) - [N, 3], the coordinate of each atom. + + Outputs: + - **output** (uint32) + + Supported Platforms: + ``GPU`` + Examples: + """ + + @prim_attr_register + def __init__(self, atom_numbers): + self.atom_numbers = atom_numbers + self.add_prim_attr('atom_numbers', self.atom_numbers) + self.init_prim_io_names( + inputs=['crd_to_uint_crd_cof', 'crd'], + outputs=['output']) + + def infer_shape(self, crd_to_uint_crd_cof, crd): + validator.check_int(crd_to_uint_crd_cof[0], 3, Rel.EQ, "crd_to_uint_crd_cof", self.name) + validator.check_int(crd[0], self.atom_numbers, Rel.EQ, "crd[0]", self.name) + validator.check_int(crd[1], 3, Rel.EQ, "crd[1]", self.name) + return crd + + def infer_dtype(self, crd_to_uint_crd_cof, crd): + validator.check_tensor_dtype_valid('crd_to_uint_crd_cof', crd_to_uint_crd_cof, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('crd', crd, [mstype.float32], self.name) + return mstype.uint32 + + +class MDIterationSetupRandState(PrimitiveWithInfer): + """ + Convert FP32 coordinate to Uint32 coordinate. + + Inputs: + - **atom_numbers** (int32) - the number of atoms N. + - **seed** (int32) - random seed. + + Outputs: + - **output** (uint32) random state. + + Supported Platforms: + ``GPU`` + Examples: + """ + + @prim_attr_register + def __init__(self, atom_numbers, seed): + self.atom_numbers = atom_numbers + self.seed = seed + self.add_prim_attr('atom_numbers', self.atom_numbers) + self.add_prim_attr('seed', self.seed) + self.init_prim_io_names( + inputs=[], + outputs=['output']) + + def infer_shape(self): + float4_numbers = math.ceil(self.atom_numbers * 3 / 4.0) + curandStatePhilox4_32_10_t_size = 64 / 4 + + return [float4_numbers * int(curandStatePhilox4_32_10_t_size),] + + def infer_dtype(self): + return mstype.float32 diff --git a/model_zoo/research/hpc/sponge/main.py b/model_zoo/research/hpc/sponge/main.py index ae9532c57d..2c1f9ca2e3 100644 --- a/model_zoo/research/hpc/sponge/main.py +++ b/model_zoo/research/hpc/sponge/main.py @@ -12,44 +12,78 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================ -"""main""" - -import time +'''main''' import argparse -from mindspore import context -from src.simulation_initial import Simulation +import time + +from src.simulation import Simulation + +import mindspore.context as context +from mindspore import Tensor parser = argparse.ArgumentParser(description='Sponge Controller') parser.add_argument('--i', type=str, default=None, help='input file') -parser.add_argument('--amber_parm', type=str, default=None, - help='paramter file in AMBER type') -parser.add_argument('--c', type=str, default=None, - help='initial coordinates file') +parser.add_argument('--amber_parm', type=str, default=None, help='paramter file in AMBER type') +parser.add_argument('--c', type=str, default=None, help='initial coordinates file') parser.add_argument('--r', type=str, default="restrt", help='') parser.add_argument('--x', type=str, default="mdcrd", help='') parser.add_argument('--o', type=str, default="mdout", help="") parser.add_argument('--box', type=str, default="mdbox", help='') +parser.add_argument('--device_id', type=int, default=0, help='') args_opt = parser.parse_args() -context.set_context(mode=context.PYNATIVE_MODE, - device_target="GPU", device_id=0, save_graphs=True) +context.set_context(mode=context.GRAPH_MODE, device_target="GPU", device_id=args_opt.device_id, save_graphs=False) if __name__ == "__main__": - start = time.time() simulation = Simulation(args_opt) - simulation.Main_Initial() - res = simulation.Initial_Neighbor_List_Update(not_first_time=0) - md_info = simulation.md_info - md_info.step_limit = 1 - for i in range(1, md_info.step_limit + 1): - print("steps: ", i) - md_info.steps = i - simulation.Main_Before_Calculate_Force() - simulation.Main_Calculate_Force() - simulation.Main_Calculate_Energy() - simulation.Main_After_Calculate_Energy() - temperature = simulation.Main_Print() - simulation.Main_Iteration_2() + start = time.time() + compiler_time = 0 + save_path = args_opt.o + file = open(save_path, 'w') + for steps in range(simulation.md_info.step_limit): + print_step = steps % simulation.ntwx + if steps == simulation.md_info.step_limit - 1: + print_step = 0 + temperature, total_potential_energy, sigma_of_bond_ene, sigma_of_angle_ene, sigma_of_dihedral_ene, \ + nb14_lj_energy_sum, nb14_cf_energy_sum, LJ_energy_sum, ee_ene, _ = simulation(Tensor(steps), Tensor(print_step)) + if steps == 0: + compiler_time = time.time() + if steps % simulation.ntwx == 0 or steps == simulation.md_info.step_limit - 1: + if steps == 0: + print("_steps_ _TEMP_ _TOT_POT_ENE_ _BOND_ENE_ " + "_ANGLE_ENE_ _DIHEDRAL_ENE_ _14LJ_ENE_ _14CF_ENE_ _LJ_ENE_ _CF_PME_ENE_") + file.write("_steps_ _TEMP_ _TOT_POT_ENE_ _BOND_ENE_ " + "_ANGLE_ENE_ _DIHEDRAL_ENE_ _14LJ_ENE_ _14CF_ENE_ _LJ_ENE_ _CF_PME_ENE_\n") + + temperature = temperature.asnumpy() + total_potential_energy = total_potential_energy.asnumpy() + print("{:>7.0f} {:>7.3f} {:>11.3f}".format(steps, float(temperature), float(total_potential_energy)), + end=" ") + if simulation.bond.bond_numbers > 0: + sigma_of_bond_ene = sigma_of_bond_ene.asnumpy() + print("{:>10.3f}".format(float(sigma_of_bond_ene)), end=" ") + if simulation.angle.angle_numbers > 0: + sigma_of_angle_ene = sigma_of_angle_ene.asnumpy() + print("{:>11.3f}".format(float(sigma_of_angle_ene)), end=" ") + if simulation.dihedral.dihedral_numbers > 0: + sigma_of_dihedral_ene = sigma_of_dihedral_ene.asnumpy() + print("{:>14.3f}".format(float(sigma_of_dihedral_ene)), end=" ") + if simulation.nb14.nb14_numbers > 0: + nb14_lj_energy_sum = nb14_lj_energy_sum.asnumpy() + nb14_cf_energy_sum = nb14_cf_energy_sum.asnumpy() + print("{:>10.3f} {:>10.3f}".format(float(nb14_lj_energy_sum), float(nb14_cf_energy_sum)), end=" ") + LJ_energy_sum = LJ_energy_sum.asnumpy() + ee_ene = ee_ene.asnumpy() + print("{:>7.3f}".format(float(LJ_energy_sum)), end=" ") + print("{:>12.3f}".format(float(ee_ene))) + if file is not None: + file.write("{:>7.0f} {:>7.3f} {:>11.3f} {:>10.3f} {:>11.3f} {:>14.3f} {:>10.3f} {:>10.3f} {:>7.3f}" + " {:>12.3f}\n".format(steps, float(temperature), float(total_potential_energy), + float(sigma_of_bond_ene), float(sigma_of_angle_ene), + float(sigma_of_dihedral_ene), float(nb14_lj_energy_sum), + float(nb14_cf_energy_sum), float(LJ_energy_sum), float(ee_ene))) + end = time.time() + file.close() print("Main time(s):", end - start) - simulation.Main_Destroy() + print("Main time(s) without compiler:", end - compiler_time) diff --git a/model_zoo/research/hpc/sponge/src/__init__.py b/model_zoo/research/hpc/sponge/src/__init__.py deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/model_zoo/research/hpc/sponge/src/angle.py b/model_zoo/research/hpc/sponge/src/angle.py index f6d00edbdc..22d619e190 100644 --- a/model_zoo/research/hpc/sponge/src/angle.py +++ b/model_zoo/research/hpc/sponge/src/angle.py @@ -12,31 +12,19 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================ -"""angle class""" - -import numpy as np -import mindspore.common.dtype as mstype -from mindspore import Tensor, nn -from mindspore.ops import operations as P - - -class Angle(nn.Cell): - """Angle class""" - +'''Angle''' +class Angle: + '''Angle''' def __init__(self, controller): - super(Angle, self).__init__() if controller.amber_parm is not None: file_path = controller.amber_parm self.read_information_from_amberfile(file_path) - self.atom_a = Tensor(np.asarray(self.h_atom_a, np.int32), mstype.int32) - self.atom_b = Tensor(np.asarray(self.h_atom_b, np.int32), mstype.int32) - self.atom_c = Tensor(np.asarray(self.h_atom_c, np.int32), mstype.int32) - self.angle_k = Tensor(np.asarray(self.h_angle_k, np.float32), mstype.float32) - self.angle_theta0 = Tensor(np.asarray(self.h_angle_theta0, np.float32), mstype.float32) - - def read_process1(self, context): - """read_information_from_amberfile process1""" + def read_information_from_amberfile(self, file_path): + '''read amber file''' + file = open(file_path, 'r') + context = file.readlines() + file.close() for idx, val in enumerate(context): if idx < len(context) - 1: if "%FLAG POINTERS" in val + context[idx + 1] and "%FORMAT(10I8)" in val + context[idx + 1]: @@ -46,6 +34,7 @@ class Angle(nn.Cell): self.angle_with_H_numbers = value[4] self.angle_without_H_numbers = value[5] self.angle_numbers = self.angle_with_H_numbers + self.angle_without_H_numbers + # print(self.angle_numbers) information = [] information.extend(value) while count < 15: @@ -57,8 +46,10 @@ class Angle(nn.Cell): print("angle type numbers ", self.angle_type_numbers) break - def read_process2(self, context): - """read_information_from_amberfile process2""" + self.h_atom_a = [0] * self.angle_numbers + self.h_atom_b = [0] * self.angle_numbers + self.h_atom_c = [0] * self.angle_numbers + self.h_type = [0] * self.angle_numbers angle_count = 0 for idx, val in enumerate(context): if "%FLAG ANGLES_INC_HYDROGEN" in val: @@ -81,20 +72,6 @@ class Angle(nn.Cell): angle_count += 1 break - return angle_count - - def read_information_from_amberfile(self, file_path): - """read information from amberfile""" - file = open(file_path, 'r') - context = file.readlines() - file.close() - self.read_process1(context) - - self.h_atom_a = [0] * self.angle_numbers - self.h_atom_b = [0] * self.angle_numbers - self.h_atom_c = [0] * self.angle_numbers - self.h_type = [0] * self.angle_numbers - angle_count = self.read_process2(context) for idx, val in enumerate(context): if "%FLAG ANGLES_WITHOUT_HYDROGEN" in val: @@ -109,14 +86,17 @@ class Angle(nn.Cell): value = list(map(int, context[start_idx].strip().split())) information.extend(value) count += len(value) - for i in range(self.angle_without_H_numbers): + for _ in range(self.angle_without_H_numbers): self.h_atom_a[angle_count] = information[(angle_count - self.angle_with_H_numbers) * 4 + 0] / 3 self.h_atom_b[angle_count] = information[(angle_count - self.angle_with_H_numbers) * 4 + 1] / 3 self.h_atom_c[angle_count] = information[(angle_count - self.angle_with_H_numbers) * 4 + 2] / 3 self.h_type[angle_count] = information[(angle_count - self.angle_with_H_numbers) * 4 + 3] - 1 angle_count += 1 break + self.processor(context, angle_count) + def processor(self, context, angle_count): + ''' processor ''' self.type_k = [0] * self.angle_type_numbers for idx, val in enumerate(context): if "%FLAG ANGLE_FORCE_CONSTANT" in val: @@ -159,17 +139,3 @@ class Angle(nn.Cell): for i in range(self.angle_numbers): self.h_angle_k.append(self.type_k[self.h_type[i]]) self.h_angle_theta0.append(self.type_theta0[self.h_type[i]]) - - def Angle_Energy(self, uint_crd, uint_dr_to_dr_cof): - """compute angle energy""" - self.angle_energy = P.AngleEnergy(self.angle_numbers)(uint_crd, uint_dr_to_dr_cof, self.atom_a, self.atom_b, - self.atom_c, self.angle_k, self.angle_theta0) - self.sigma_of_angle_ene = P.ReduceSum()(self.angle_energy) - return self.sigma_of_angle_ene - - def Angle_Force_With_Atom_Energy(self, uint_crd, scaler): - """compute angle force with atom energy""" - print("angele angle numbers:", self.angle_numbers) - self.afae = P.AngleForceWithAtomEnergy(angle_numbers=self.angle_numbers) - frc, ene = self.afae(uint_crd, scaler, self.atom_a, self.atom_b, self.atom_c, self.angle_k, self.angle_theta0) - return frc, ene diff --git a/model_zoo/research/hpc/sponge/src/bond.py b/model_zoo/research/hpc/sponge/src/bond.py index 9fc526d6f0..4cc5b659bd 100644 --- a/model_zoo/research/hpc/sponge/src/bond.py +++ b/model_zoo/research/hpc/sponge/src/bond.py @@ -12,19 +12,10 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================ -"""bond class""" - -import numpy as np -import mindspore.common.dtype as mstype -from mindspore import Tensor, nn -from mindspore.ops import operations as P - - -class Bond(nn.Cell): - """bond class""" - +'''Bond''' +class Bond: + '''Bond''' def __init__(self, controller, md_info): - super(Bond, self).__init__() self.atom_numbers = md_info.atom_numbers @@ -32,13 +23,11 @@ class Bond(nn.Cell): file_path = controller.amber_parm self.read_information_from_amberfile(file_path) - self.atom_a = Tensor(np.asarray(self.h_atom_a, np.int32), mstype.int32) - self.atom_b = Tensor(np.asarray(self.h_atom_b, np.int32), mstype.int32) - self.bond_k = Tensor(np.asarray(self.h_k, np.float32), mstype.float32) - self.bond_r0 = Tensor(np.asarray(self.h_r0, np.float32), mstype.float32) - - def process1(self, context): - """process1: read information from amberfile""" + def read_information_from_amberfile(self, file_path): + '''read amber file''' + file = open(file_path, 'r') + context = file.readlines() + file.close() for idx, val in enumerate(context): if idx < len(context) - 1: if "%FLAG POINTERS" in val + context[idx + 1] and "%FORMAT(10I8)" in val + context[idx + 1]: @@ -48,7 +37,6 @@ class Bond(nn.Cell): self.bond_with_hydrogen = value[2] self.bond_numbers = value[3] self.bond_numbers += self.bond_with_hydrogen - print(self.bond_numbers) information = [] information.extend(value) while count < 16: @@ -76,13 +64,6 @@ class Bond(nn.Cell): self.bond_type_k = information[:self.bond_type_numbers] break - def read_information_from_amberfile(self, file_path): - """read information from amberfile""" - file = open(file_path, 'r') - context = file.readlines() - file.close() - self.process1(context) - for idx, val in enumerate(context): if "%FLAG BOND_EQUIL_VALUE" in val: count = 0 @@ -98,7 +79,10 @@ class Bond(nn.Cell): count += len(value) self.bond_type_r = information[:self.bond_type_numbers] break + self.processor(context) + def processor(self, context): + '''processor''' for idx, val in enumerate(context): if "%FLAG BONDS_INC_HYDROGEN" in val: self.h_atom_a = [0] * self.bond_numbers @@ -128,6 +112,7 @@ class Bond(nn.Cell): for idx, val in enumerate(context): if "%FLAG BONDS_WITHOUT_HYDROGEN" in val: + count = 0 start_idx = idx information = [] @@ -147,17 +132,3 @@ class Bond(nn.Cell): self.h_k[i] = self.bond_type_k[tmpi] self.h_r0[i] = self.bond_type_r[tmpi] break - - def Bond_Energy(self, uint_crd, uint_dr_to_dr_cof): - """compute bond energy""" - self.bond_energy = P.BondEnergy(self.bond_numbers, self.atom_numbers)(uint_crd, uint_dr_to_dr_cof, self.atom_a, - self.atom_b, self.bond_k, self.bond_r0) - self.sigma_of_bond_ene = P.ReduceSum()(self.bond_energy) - return self.sigma_of_bond_ene - - def Bond_Force_With_Atom_Energy(self, uint_crd, scaler): - """compute bond force with atom energy""" - self.bfatomenergy = P.BondForceWithAtomEnergy(bond_numbers=self.bond_numbers, - atom_numbers=self.atom_numbers) - frc, atom_energy = self.bfatomenergy(uint_crd, scaler, self.atom_a, self.atom_b, self.bond_k, self.bond_r0) - return frc, atom_energy diff --git a/model_zoo/research/hpc/sponge/src/dihedral.py b/model_zoo/research/hpc/sponge/src/dihedral.py index 5529a5ad29..2d06c0e3b1 100644 --- a/model_zoo/research/hpc/sponge/src/dihedral.py +++ b/model_zoo/research/hpc/sponge/src/dihedral.py @@ -12,38 +12,23 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================ -"""dihedral class""" - +'''Dihedral''' import math -import numpy as np - -import mindspore.common.dtype as mstype -from mindspore import Tensor, nn -from mindspore.ops import operations as P -class Dihedral(nn.Cell): - """dihedral class""" - +class Dihedral: + '''Dihedral''' def __init__(self, controller): - super(Dihedral, self).__init__() self.CONSTANT_Pi = 3.1415926535897932 if controller.amber_parm is not None: file_path = controller.amber_parm self.read_information_from_amberfile(file_path) - self.atom_a = Tensor(np.asarray(self.h_atom_a, np.int32), mstype.int32) - self.atom_b = Tensor(np.asarray(self.h_atom_b, np.int32), mstype.int32) - self.atom_c = Tensor(np.asarray(self.h_atom_c, np.int32), mstype.int32) - self.atom_d = Tensor(np.asarray(self.h_atom_d, np.int32), mstype.int32) - self.pk = Tensor(np.asarray(self.pk, np.float32), mstype.float32) - self.gamc = Tensor(np.asarray(self.gamc, np.float32), mstype.float32) - self.gams = Tensor(np.asarray(self.gams, np.float32), mstype.float32) - self.pn = Tensor(np.asarray(self.pn, np.float32), mstype.float32) - self.ipn = Tensor(np.asarray(self.ipn, np.int32), mstype.int32) - - def process1(self, context): - """process1: read information from amberfile""" + def read_information_from_amberfile(self, file_path): + '''read amber file''' + file = open(file_path, 'r') + context = file.readlines() + file.close() for idx, val in enumerate(context): if idx < len(context) - 1: if "%FLAG POINTERS" in val + context[idx + 1] and "%FORMAT(10I8)" in val + context[idx + 1]: @@ -115,15 +100,10 @@ class Dihedral(nn.Cell): count += len(value) self.pn_type = information[:self.dihedral_type_numbers] break + self.processor(context) - def read_information_from_amberfile(self, file_path): - """read information from amberfile""" - file = open(file_path, 'r') - context = file.readlines() - file.close() - - self.process1(context) - + def processor(self, context): + '''processor''' self.h_atom_a = [0] * self.dihedral_numbers self.h_atom_b = [0] * self.dihedral_numbers self.h_atom_c = [0] * self.dihedral_numbers @@ -204,18 +184,3 @@ class Dihedral(nn.Cell): for i in range(self.dihedral_numbers): if self.h_atom_c[i] < 0: self.h_atom_c[i] *= -1 - - def Dihedral_Engergy(self, uint_crd, uint_dr_to_dr_cof): - """compute dihedral energy""" - self.dihedral_energy = P.DihedralEnergy(self.dihedral_numbers)(uint_crd, uint_dr_to_dr_cof, self.atom_a, - self.atom_b, self.atom_c, self.atom_d, self.ipn, - self.pk, self.gamc, self.gams, self.pn) - self.sigma_of_dihedral_ene = P.ReduceSum()(self.dihedral_energy) - return self.sigma_of_dihedral_ene - - def Dihedral_Force_With_Atom_Energy(self, uint_crd, scaler): - """compute dihedral force and atom energy""" - self.dfae = P.DihedralForceWithAtomEnergy(dihedral_numbers=self.dihedral_numbers) - self.frc, self.ene = self.dfae(uint_crd, scaler, self.atom_a, self.atom_b, self.atom_c, self.atom_d, - self.ipn, self.pk, self.gamc, self.gams, self.pn) - return self.frc, self.ene diff --git a/model_zoo/research/hpc/sponge/src/Langevin_Liujian_md.py b/model_zoo/research/hpc/sponge/src/langevin_liujian_md.py similarity index 59% rename from model_zoo/research/hpc/sponge/src/Langevin_Liujian_md.py rename to model_zoo/research/hpc/sponge/src/langevin_liujian_md.py index 222b6db13c..0f25929f9d 100644 --- a/model_zoo/research/hpc/sponge/src/Langevin_Liujian_md.py +++ b/model_zoo/research/hpc/sponge/src/langevin_liujian_md.py @@ -12,18 +12,14 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================ -"""Langevin Liujian MD class""" - +'''LagevinLiuJian''' import math import numpy as np -from mindspore.common.tensor import Tensor -import mindspore.common.dtype as mstype class Langevin_Liujian: - """Langevin_Liujian class""" - + '''LagevinLiuJian''' def __init__(self, controller, atom_numbers): self.atom_numbers = atom_numbers if controller.amber_parm is not None: @@ -37,29 +33,27 @@ class Langevin_Liujian: controller.Command_Set["target_temperature"]) self.gamma_ln = 1.0 if "langevin_gamma" not in controller.Command_Set else float( controller.Command_Set["langevin_gamma"]) - self.rand_seed = 0 if "langevin_seed" not in controller.Command_Set else float( - controller.Command_Set["langevin_seed"]) # jiahong0315 + self.rand_seed = 1 if "langevin_seed" not in controller.Command_Set else float( + controller.Command_Set["langevin_seed"]) self.max_velocity = 10000.0 if "velocity_max" not in controller.Command_Set else float( controller.Command_Set["velocity_max"]) assert self.max_velocity > 0 - self.is_max_velocity = 0 if "velocity_max" not in controller.Command_Set else 1 print("target temperature is ", self.target_temperature) print("friction coefficient is ", self.gamma_ln, "ps^-1") print("random seed is ", self.rand_seed) self.dt = float(controller.Command_Set["dt"]) self.dt *= self.CONSTANT_TIME_CONVERTION self.half_dt = 0.5 * self.dt - self.float4_numbers = math.ceil(3.0 * self.atom_numbers / 4.0) + self.rand_state = np.float32(np.zeros([math.ceil(3 * self.atom_numbers / 4.0) * 16,])) self.gamma_ln = self.gamma_ln / self.CONSTANT_TIME_CONVERTION self.exp_gamma = math.exp(-1 * self.gamma_ln * self.dt) self.sqrt_gamma = math.sqrt((1. - self.exp_gamma * self.exp_gamma) * self.target_temperature * self.CONSTANT_kB) self.h_sqrt_mass = [0] * self.atom_numbers for i in range(self.atom_numbers): self.h_sqrt_mass[i] = self.sqrt_gamma * math.sqrt(1. / self.h_mass[i]) - self.d_sqrt_mass = Tensor(self.h_sqrt_mass, mstype.float32) def read_information_from_amberfile(self, file_path): - """read information from amberfile""" + '''read amber file''' file = open(file_path, 'r') context = file.readlines() file.close() @@ -81,29 +75,3 @@ class Langevin_Liujian: for i in range(self.atom_numbers): self.h_mass[i] = information[i] break - - def MDIterationLeapFrog_Liujian(self, atom_numbers, half_dt, dt, exp_gamma, inverse_mass, sqrt_mass_inverse, vel, - crd, frc, random_frc): - """compute MDIterationLeapFrog Liujian""" - inverse_mass = inverse_mass.reshape((-1, 1)) - sqrt_mass_inverse = sqrt_mass_inverse.reshape((-1, 1)) - acc = inverse_mass * frc - vel = vel + dt * acc - crd = crd + half_dt * vel - vel = exp_gamma * vel + sqrt_mass_inverse * random_frc - crd = crd + half_dt * vel - frc = Tensor(np.zeros((atom_numbers, 3)), mstype.float32) - return vel, crd, frc, acc - - def MD_Iteration_Leap_Frog(self, d_mass_inverse, vel_in, crd_in, frc_in): - """MD_Iteration_Leap_Frog""" - np.random.seed(int(self.rand_seed)) - self.rand_force = Tensor(np.zeros((self.atom_numbers, 3)), mstype.float32) - # self.rand_force = Tensor(np.random.randn(self.atom_numbers, 3), mstype.float32) - vel, crd, frc, acc = self.MDIterationLeapFrog_Liujian(atom_numbers=self.atom_numbers, half_dt=self.half_dt, - dt=self.dt, exp_gamma=self.exp_gamma, - inverse_mass=d_mass_inverse, - sqrt_mass_inverse=self.d_sqrt_mass, - vel=vel_in, crd=crd_in, - frc=frc_in, random_frc=self.rand_force) - return vel, crd, frc, acc diff --git a/model_zoo/research/hpc/sponge/src/lennard_jones.py b/model_zoo/research/hpc/sponge/src/lennard_jones.py index 4be20eeccd..b7617c11d6 100644 --- a/model_zoo/research/hpc/sponge/src/lennard_jones.py +++ b/model_zoo/research/hpc/sponge/src/lennard_jones.py @@ -12,30 +12,16 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================ -"""lennard jones""" - -import numpy as np -import mindspore.common.dtype as mstype -from mindspore import Tensor, nn -from mindspore.ops import operations as P - - -class Lennard_Jones_Information(nn.Cell): - """class Lennard Jones Information""" - +'''Lennard Jones''' +class Lennard_Jones_Information: + '''Lennard Jones''' def __init__(self, controller): - super(Lennard_Jones_Information, self).__init__() if controller.amber_parm is not None: file_path = controller.amber_parm self.read_information_from_amberfile(file_path) - self.atom_LJ_type = Tensor(np.asarray(self.atom_LJ_type, dtype=np.int32), mstype.int32) - self.LJ_A = Tensor(np.asarray(self.LJ_A, dtype=np.float32), mstype.float32) - self.LJ_B = Tensor(np.asarray(self.LJ_B, dtype=np.float32), mstype.float32) - self.LJ_energy_sum = 0 - self.LJ_energy = 0 def read_information_from_amberfile(self, file_path): - """read information from amberfile""" + '''read amber file''' file = open(file_path, 'r') context = file.readlines() file.close() @@ -48,8 +34,8 @@ class Lennard_Jones_Information(nn.Cell): value = list(map(int, context[start_idx].strip().split())) self.atom_numbers = value[0] self.atom_type_numbers = value[1] - self.pair_type_numbers = int(self.atom_type_numbers * (self.atom_type_numbers + 1) / 2) - print(self.pair_type_numbers) + self.pair_type_numbers = int( + self.atom_type_numbers * (self.atom_type_numbers + 1) / 2) # TODO 这个地方有问题啊 break self.atom_LJ_type = [0] * self.atom_numbers for idx, val in enumerate(context): @@ -102,21 +88,3 @@ class Lennard_Jones_Information(nn.Cell): for i in range(self.pair_type_numbers): self.LJ_B[i] = 6.0 * information[i] break - - def LJ_Energy(self, uint_crd_with_LJ, uint_dr_to_dr_cof, nl_atom_numbers, nl_atom_serial, cutoff_square): - """compute LJ energy""" - uint_crd, LJtype, charge = uint_crd_with_LJ - self.LJ_energy = P.LJEnergy(self.atom_numbers, cutoff_square) \ - (uint_crd, LJtype, charge, uint_dr_to_dr_cof, nl_atom_numbers, nl_atom_serial, self.LJ_A, self.LJ_B) - self.LJ_energy_sum = P.ReduceSum()(self.LJ_energy) - return self.LJ_energy_sum - - def LJ_Force_With_PME_Direct_Force(self, atom_numbers, uint_crd_with_LJ, uint_dr_to_dr_cof, nl_number, nl_serial, - cutoff, beta): - """compute LJ force with PME direct force""" - assert atom_numbers == self.atom_numbers - assert isinstance(uint_crd_with_LJ, tuple) - uint_crd_f, LJtype, charge = uint_crd_with_LJ - self.ljfd = P.LJForceWithPMEDirectForce(atom_numbers, cutoff, beta) - frc = self.ljfd(uint_crd_f, LJtype, charge, uint_dr_to_dr_cof, nl_number, nl_serial, self.LJ_A, self.LJ_B) - return frc diff --git a/model_zoo/research/hpc/sponge/src/md_information.py b/model_zoo/research/hpc/sponge/src/md_information.py index d9fccbac60..81051e540a 100644 --- a/model_zoo/research/hpc/sponge/src/md_information.py +++ b/model_zoo/research/hpc/sponge/src/md_information.py @@ -12,36 +12,29 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================ -"""md information""" - +'''MD Information''' import numpy as np -import mindspore.common.dtype as mstype -from mindspore import Tensor, nn -from mindspore.ops import operations as P -class md_information(nn.Cell): - """class md information""" - +class md_information: + '''MD Information''' def __init__(self, controller): - super(md_information, self).__init__() CONSTANT_TIME_CONVERTION = 20.455 CONSTANT_UINT_MAX_FLOAT = 4294967296.0 self.md_task = controller.md_task self.mode = 0 if "mode" not in controller.Command_Set else int(controller.Command_Set["mode"]) - self.dt = 0.001 * CONSTANT_TIME_CONVERTION if "dt" not in controller.Command_Set \ - else float(controller.Command_Set["dt"]) * CONSTANT_TIME_CONVERTION - self.skin = 2.0 if "skin" not in controller.Command_Set \ - else float(controller.Command_Set["skin"]) + self.dt = 0.001 * CONSTANT_TIME_CONVERTION if "dt" not in controller.Command_Set else float( + controller.Command_Set["dt"]) * CONSTANT_TIME_CONVERTION + self.skin = 2.0 if "skin" not in controller.Command_Set else float(controller.Command_Set["skin"]) self.trans_vec = [self.skin, self.skin, self.skin] self.trans_vec_minus = -1 * self.trans_vec self.step_limit = 1000 if "step_limit" not in controller.Command_Set else int( controller.Command_Set["step_limit"]) self.netfrc = 0 if "net_force" not in controller.Command_Set else int(controller.Command_Set["net_force"]) - self.ntwx = 1000 if "write_information_interval" not in controller.Command_Set else \ - int(controller.Command_Set["write_information_interval"]) - self.ntce = self.step_limit + 1 if "calculate_energy_interval" not in controller.Command_Set else \ - int(controller.Command_Set["calculate_energy_interval"]) + self.ntwx = 1000 if "write_information_interval" not in controller.Command_Set else int( + controller.Command_Set["write_information_interval"]) + self.ntce = self.step_limit + 1 if "calculate_energy_interval" not in controller.Command_Set else int( + controller.Command_Set["calculate_energy_interval"]) self.atom_numbers = 0 self.residue_numbers = 0 self.density = 0.0 @@ -51,7 +44,6 @@ class md_information(nn.Cell): self.h_mass = [] self.h_mass_inverse = [] self.h_charge = [] - self.steps = 0 if controller.amber_parm is not None: self.read_basic_system_information_from_amber_file(controller.amber_parm) @@ -67,23 +59,13 @@ class md_information(nn.Cell): self.uint_dr_to_dr_cof = [1.0 / self.crd_to_uint_crd_cof[0], 1.0 / self.crd_to_uint_crd_cof[1], 1.0 / self.crd_to_uint_crd_cof[2]] self.density *= 1e24 / 6.023e23 / (self.box_length[0] * self.box_length[1] * self.box_length[2]) - self.frc = Tensor(np.zeros((self.atom_numbers, 3)), mstype.float32) - self.crd = Tensor(np.array(self.coordinate, dtype=np.float32).reshape((self.atom_numbers, 3)), mstype.float32) - self.crd_n = np.array(self.coordinate).reshape([self.atom_numbers, 3]) - self.crd_old = Tensor(np.zeros([self.atom_numbers, 3], dtype=np.float32), mstype.float32) - self.uint_crd = Tensor(np.zeros([self.atom_numbers, 3], dtype=np.uint32), mstype.uint32) - self.charge = Tensor(self.h_charge, mstype.float32) - self.crd_to_uint_crd_cof_n = np.array(self.crd_to_uint_crd_cof) - self.crd_to_uint_crd_cof = Tensor(self.crd_to_uint_crd_cof, mstype.float32) - self.uint_dr_to_dr_cof = Tensor(self.uint_dr_to_dr_cof, mstype.float32) - self.uint_crd_with_LJ = None - self.d_mass_inverse = Tensor(self.h_mass_inverse, mstype.float32) - self.d_res_start = Tensor(self.h_res_start, mstype.int32) - self.d_res_end = Tensor(self.h_res_end, mstype.int32) - self.d_mass = Tensor(self.h_mass, mstype.float32) - def process1(self, context): - """process1: read basic system information from amber file""" + self.velocity = np.reshape(np.asarray(self.velocity, np.float32), [self.atom_numbers, 3]) + + def read_basic_system_information_from_amber_file(self, path): + '''read amber file''' + file = open(path, 'r') + context = file.readlines() for idx, val in enumerate(context): if idx < len(context) - 1: if "%FLAG POINTERS" in val + context[idx + 1] and "%FORMAT(10I8)" in val + context[idx + 1]: @@ -95,16 +77,10 @@ class md_information(nn.Cell): start_idx += 1 value = list(map(int, context[start_idx].strip().split())) count += len(value) - self.residue_numbers = list(map(int, context[start_idx].strip().split()))[10 - (count - 10)] + self.residue_numbers = list(map(int, context[start_idx].strip().split()))[ + 10 - (count - 10)] # may exist bug break - def read_basic_system_information_from_amber_file(self, path): - """read basic system information from amber file""" - file = open(path, 'r') - context = file.readlines() - file.close() - self.process1(context) - if self.residue_numbers != 0 and self.atom_numbers != 0: for idx, val in enumerate(context): if "%FLAG RESIDUE_POINTER" in val: @@ -124,42 +100,45 @@ class md_information(nn.Cell): self.h_res_start.append(self.lin_serial[-1] - 1) self.h_res_end.append(self.atom_numbers + 1 - 1) break + self.processor(context) - for idx, val in enumerate(context): - if "%FLAG MASS" in val: - count = 0 - start_idx = idx - while count != self.atom_numbers: - start_idx += 1 - if "%FORMAT" in context[start_idx]: - continue - else: - value = list(map(float, context[start_idx].strip().split())) - self.h_mass.extend(value) - count += len(value) - for i in range(self.atom_numbers): - if self.h_mass[i] == 0: - self.h_mass_inverse.append(0.0) - else: - self.h_mass_inverse.append(1.0 / self.h_mass[i]) - self.density += self.h_mass[i] - break - for idx, val in enumerate(context): - if "%FLAG CHARGE" in val: - count = 0 - start_idx = idx - while count != self.atom_numbers: - start_idx += 1 - if "%FORMAT" in context[start_idx]: - continue - else: - value = list(map(float, context[start_idx].strip().split())) - self.h_charge.extend(value) - count += len(value) - break + def processor(self, context): + '''processor''' + for idx, val in enumerate(context): + if "%FLAG MASS" in val: + count = 0 + start_idx = idx + while count != self.atom_numbers: + start_idx += 1 + if "%FORMAT" in context[start_idx]: + continue + else: + value = list(map(float, context[start_idx].strip().split())) + self.h_mass.extend(value) + count += len(value) + for i in range(self.atom_numbers): + if self.h_mass[i] == 0: + self.h_mass_inverse.append(0.0) + else: + self.h_mass_inverse.append(1.0 / self.h_mass[i]) + self.density += self.h_mass[i] + break + for idx, val in enumerate(context): + if "%FLAG CHARGE" in val: + count = 0 + start_idx = idx + while count != self.atom_numbers: + start_idx += 1 + if "%FORMAT" in context[start_idx]: + continue + else: + value = list(map(float, context[start_idx].strip().split())) + self.h_charge.extend(value) + count += len(value) + break def read_basic_system_information_from_rst7(self, path, irest): - """read basic system information from rst7""" + '''read rst7 file''' file = open(path, 'r') context = file.readlines() file.close() @@ -191,22 +170,4 @@ class md_information(nn.Cell): self.coordinate = information[: 3 * self.atom_numbers] self.velocity = [0.0] * (3 * self.atom_numbers) self.box_length = information[3 * self.atom_numbers:3 * self.atom_numbers + 3] - self.vel = Tensor(self.velocity, mstype.float32).reshape((self.atom_numbers, 3)) - self.acc = Tensor(np.zeros((self.atom_numbers, 3), dtype=np.float32), mstype.float32) - - def MD_Information_Crd_To_Uint_Crd(self): - """transform the crd to uint crd""" - uint_crd = self.crd.asnumpy() * (0.5 * self.crd_to_uint_crd_cof.asnumpy()) * 2 - self.uint_crd = Tensor(uint_crd, mstype.uint32) - - return self.uint_crd - - def Centerize(self): - return - - def MD_Information_Temperature(self): - """compute temperature""" - self.mdtemp = P.MDTemperature(self.residue_numbers, self.atom_numbers) - self.res_ek_energy = self.mdtemp(self.d_res_start, self.d_res_end, self.vel, self.d_mass) - self.d_temperature = P.ReduceSum()(self.res_ek_energy) - return self.d_temperature + print("system size is ", self.box_length[0], self.box_length[1], self.box_length[2]) diff --git a/model_zoo/research/hpc/sponge/src/nb14.py b/model_zoo/research/hpc/sponge/src/nb14.py index fce4edd10a..9c37ec79e0 100644 --- a/model_zoo/research/hpc/sponge/src/nb14.py +++ b/model_zoo/research/hpc/sponge/src/nb14.py @@ -12,20 +12,10 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================ -"""nb14""" - -import numpy as np -import mindspore.common.dtype as mstype -from mindspore import Tensor, nn -from mindspore.ops import operations as P - - -class NON_BOND_14(nn.Cell): - """class Non bond 14""" - +'''NON BOND''' +class NON_BOND_14: + '''NON BOND''' def __init__(self, controller, dihedral, atom_numbers): - super(NON_BOND_14, self).__init__() - self.dihedral_with_hydrogen = dihedral.dihedral_with_hydrogen self.dihedral_numbers = dihedral.dihedral_numbers self.dihedral_type_numbers = dihedral.dihedral_type_numbers @@ -34,14 +24,20 @@ class NON_BOND_14(nn.Cell): if controller.amber_parm is not None: file_path = controller.amber_parm self.read_information_from_amberfile(file_path) + self.h_atom_a = self.h_atom_a[:self.nb14_numbers] + self.h_atom_b = self.h_atom_b[:self.nb14_numbers] + self.h_lj_scale_factor = self.h_lj_scale_factor[:self.nb14_numbers] + self.h_cf_scale_factor = self.h_cf_scale_factor[:self.nb14_numbers] - self.atom_a = Tensor(np.asarray(self.h_atom_a, np.int32), mstype.int32) - self.atom_b = Tensor(np.asarray(self.h_atom_b, np.int32), mstype.int32) - self.lj_scale_factor = Tensor(np.asarray(self.h_lj_scale_factor, np.float32), mstype.float32) - self.cf_scale_factor = Tensor(np.asarray(self.h_cf_scale_factor, np.float32), mstype.float32) + def read_information_from_amberfile(self, file_path): + '''read amber file''' + file = open(file_path, 'r') + context = file.readlines() + file.close() + + self.cf_scale_type = [0] * self.dihedral_type_numbers + self.lj_scale_type = [0] * self.dihedral_type_numbers - def process1(self, context): - """process1: read information from amberfile""" for idx, val in enumerate(context): if "%FLAG SCEE_SCALE_FACTOR" in val: count = 0 @@ -73,16 +69,10 @@ class NON_BOND_14(nn.Cell): count += len(value) self.lj_scale_type = information[:self.dihedral_type_numbers] break + self.processor(context) - def read_information_from_amberfile(self, file_path): - """read information from amberfile""" - file = open(file_path, 'r') - context = file.readlines() - file.close() - - self.cf_scale_type = [0] * self.dihedral_type_numbers - self.lj_scale_type = [0] * self.dihedral_type_numbers - self.process1(context) + def processor(self, context): + '''processor''' self.h_atom_a = [0] * self.dihedral_numbers self.h_atom_b = [0] * self.dihedral_numbers self.h_lj_scale_factor = [0] * self.dihedral_numbers @@ -154,42 +144,3 @@ class NON_BOND_14(nn.Cell): break self.nb14_numbers = nb14_numbers - - def Non_Bond_14_LJ_Energy(self, uint_crd_with_LJ, uint_dr_to_dr_cof, LJ_A, LJ_B): - """compute Non bond 14 LJ energy""" - assert isinstance(uint_crd_with_LJ, tuple) - uint_crd, LJtype, charge = uint_crd_with_LJ - self.LJ_energy = P.Dihedral14LJEnergy(self.nb14_numbers, self.atom_numbers)(uint_crd, LJtype, charge, - uint_dr_to_dr_cof, self.atom_a, - self.atom_b, self.lj_scale_factor, - LJ_A, LJ_B) - self.nb14_lj_energy_sum = P.ReduceSum()(self.LJ_energy) - - return self.nb14_lj_energy_sum - - def Non_Bond_14_CF_Energy(self, uint_crd_with_LJ, uint_dr_to_dr_cof): - """compute Non bond 14 CF energy""" - assert isinstance(uint_crd_with_LJ, tuple) - uint_crd, LJtype, charge = uint_crd_with_LJ - self.CF_energy = P.Dihedral14CFEnergy(self.nb14_numbers, self.atom_numbers)(uint_crd, LJtype, charge, - uint_dr_to_dr_cof, self.atom_a, - self.atom_b, self.cf_scale_factor) - self.nb14_cf_energy_sum = P.ReduceSum()(self.CF_energy) - return self.nb14_cf_energy_sum - - def Non_Bond_14_LJ_CF_Energy(self, uint_crd_with_LJ, uint_dr_to_dr_cof, LJ_A, LJ_B): - """compute Non bond 14 LJ and CF energy""" - assert isinstance(uint_crd_with_LJ, tuple) - self.nb14_lj_energy_sum = self.Non_Bond_14_LJ_Energy(uint_crd_with_LJ, uint_dr_to_dr_cof, LJ_A, LJ_B) - self.nb14_cf_energy_sum = self.Non_Bond_14_CF_Energy(uint_crd_with_LJ, uint_dr_to_dr_cof) - - return self.nb14_lj_energy_sum, self.nb14_cf_energy_sum - - def Non_Bond_14_LJ_CF_Force_With_Atom_Energy(self, uint_crd_with_LJ, boxlength, LJ_A, LJ_B): - """compute Non bond 14 LJ CF force and atom energy""" - self.d14lj = P.Dihedral14LJCFForceWithAtomEnergy(nb14_numbers=self.nb14_numbers, atom_numbers=self.atom_numbers) - assert isinstance(uint_crd_with_LJ, tuple) - uint_crd_f, LJtype, charge = uint_crd_with_LJ - self.frc, self.atom_ene = self.d14lj(uint_crd_f, LJtype, charge, boxlength, self.atom_a, self.atom_b, - self.lj_scale_factor, self.cf_scale_factor, LJ_A, LJ_B) - return self.frc, self.atom_ene diff --git a/model_zoo/research/hpc/sponge/src/neighbor_list.py b/model_zoo/research/hpc/sponge/src/neighbor_list.py index d9991fa849..607f6d258c 100644 --- a/model_zoo/research/hpc/sponge/src/neighbor_list.py +++ b/model_zoo/research/hpc/sponge/src/neighbor_list.py @@ -12,25 +12,16 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================ -"""neighbour list""" - -import numpy as np -import mindspore.common.dtype as mstype -from mindspore import Tensor, nn -from mindspore.ops import operations as P - - -class nb_infomation(nn.Cell): - """neighbour list""" - +'''Neighbor List''' +class neighbor_list: + '''Neighbor List''' def __init__(self, controller, atom_numbers, box_length): - super(nb_infomation, self).__init__() - self.refresh_interval = 20 if "neighbor_list_refresh_interval" not in controller.Command_Set else \ - int(controller.Command_Set["neighbor_list_refresh_interval"]) - self.max_atom_in_grid_numbers = 64 if "max_atom_in_grid_numbers" not in controller.Command_Set else \ - int(controller.Command_Set["max_atom_in_grid_numbers"]) - self.max_neighbor_numbers = 800 if "max_neighbor_numbers" not in controller.Command_Set else \ - int(controller.Command_Set["max_neighbor_numbers"]) + self.refresh_interval = 20 if "neighbor_list_refresh_interval" not in controller.Command_Set else int( + controller.Command_Set["neighbor_list_refresh_interval"]) + self.max_atom_in_grid_numbers = 64 if "max_atom_in_grid_numbers" not in controller.Command_Set else int( + controller.Command_Set["max_atom_in_grid_numbers"]) + self.max_neighbor_numbers = 800 if "max_neighbor_numbers" not in controller.Command_Set else int( + controller.Command_Set["max_neighbor_numbers"]) self.skin = 2.0 if "skin" not in controller.Command_Set else float(controller.Command_Set["skin"]) self.cutoff = 10.0 if "cut" not in controller.Command_Set else float(controller.Command_Set["cut"]) self.cutoff_square = self.cutoff * self.cutoff @@ -47,24 +38,10 @@ class nb_infomation(nn.Cell): self.Initial_Neighbor_Grid() self.not_first_time = 0 - self.refresh_count = 0 - - self.atom_numbers_in_grid_bucket = Tensor(np.asarray(self.atom_numbers_in_grid_bucket, np.int32), mstype.int32) - self.bucket = Tensor( - np.asarray(self.bucket, np.int32).reshape([self.grid_numbers, self.max_atom_in_grid_numbers]), mstype.int32) - self.grid_N = Tensor(np.asarray(self.grid_N, np.int32), mstype.int32) - self.grid_length_inverse = Tensor(np.asarray(self.grid_length_inverse, np.float32), mstype.float32) - self.atom_in_grid_serial = Tensor(np.zeros(self.atom_numbers, np.int32), mstype.int32) - self.pointer = Tensor(np.asarray(self.pointer, np.int32).reshape([self.grid_numbers, 125]), mstype.int32) - self.nl_atom_numbers = Tensor(np.zeros(self.atom_numbers, np.int32), mstype.int32) - self.nl_atom_serial = Tensor(np.zeros([self.atom_numbers, self.max_neighbor_numbers], np.int32), mstype.int32) - self.excluded_list_start = Tensor(np.asarray(self.excluded_list_start, np.int32), mstype.int32) - self.excluded_list = Tensor(np.asarray(self.excluded_list, np.int32), mstype.int32) - self.excluded_numbers = Tensor(np.asarray(self.excluded_numbers, np.int32), mstype.int32) - self.need_refresh_flag = Tensor(np.asarray([0], np.int32), mstype.int32) + self.refresh_count = [0] def read_information_from_amberfile(self, file_path): - """read information from amberfile""" + '''read amber file''' file = open(file_path, 'r') context = file.readlines() file.close() @@ -85,6 +62,7 @@ class nb_infomation(nn.Cell): information.extend(value) count += len(value) self.excluded_atom_numbers = information[10] + print("excluded atom numbers ", self.excluded_atom_numbers) break for idx, val in enumerate(context): if "%FLAG NUMBER_EXCLUDED_ATOMS" in val: @@ -125,37 +103,22 @@ class nb_infomation(nn.Cell): count = 0 for i in range(self.atom_numbers): tmp_list = [] - for _ in range(self.excluded_numbers[i]): + if self.excluded_numbers[i] == 1: tmp_list.append(information[count] - 1) + if information[count] == 0: + self.excluded_numbers[i] = 0 count += 1 - tmp_list = sorted(tmp_list) + else: + for _ in range(self.excluded_numbers[i]): + tmp_list.append(information[count] - 1) + + count += 1 + tmp_list = sorted(tmp_list) self.excluded_list.extend(tmp_list) break - def fun(self, Nx, Ny, Nz, l, m, temp_grid_serial, count): - """fun to replace the for""" - for n in range(-2, 3): - xx = Nx + l - if xx < 0: - xx = xx + self.Nx - elif xx >= self.Nx: - xx = xx - self.Nx - yy = Ny + m - if yy < 0: - yy = yy + self.Ny - elif yy >= self.Ny: - yy = yy - self.Ny - zz = Nz + n - if zz < 0: - zz = zz + self.Nz - elif zz >= self.Nz: - zz = zz - self.Nz - temp_grid_serial[count] = zz * self.Nxy + yy * self.Nx + xx - count += 1 - return temp_grid_serial, count - def Initial_Neighbor_Grid(self): - """initial neighbour grid""" + '''init neighbor grid''' half_cutoff = self.half_cutoff_with_skin self.Nx = int(self.box_length[0] / half_cutoff) self.Ny = int(self.box_length[1] / half_cutoff) @@ -177,31 +140,23 @@ class nb_infomation(nn.Cell): count = 0 for l in range(-2, 3): for m in range(-2, 3): - temp_grid_serial, count = self.fun(Nx, Ny, Nz, l, m, temp_grid_serial, count) + for n in range(-2, 3): + xx = Nx + l + if xx < 0: + xx = xx + self.Nx + elif xx >= self.Nx: + xx = xx - self.Nx + yy = Ny + m + if yy < 0: + yy = yy + self.Ny + elif yy >= self.Ny: + yy = yy - self.Ny + zz = Nz + n + if zz < 0: + zz = zz + self.Nz + elif zz >= self.Nz: + zz = zz - self.Nz + temp_grid_serial[count] = zz * self.Nxy + yy * self.Nx + xx + count += 1 temp_grid_serial = sorted(temp_grid_serial) self.pointer.extend(temp_grid_serial) - - def NeighborListUpdate(self, crd, old_crd, uint_crd, crd_to_uint_crd_cof, uint_dr_to_dr_cof, box_length, - not_first_time=0): - """NeighborList Update""" - self.not_first_time = not_first_time - self.neighbor_list_update = P.NeighborListUpdate(grid_numbers=self.grid_numbers, atom_numbers=self.atom_numbers, - refresh_count=self.refresh_count, - not_first_time=self.not_first_time, - Nxy=self.Nxy, excluded_atom_numbers=self.excluded_atom_numbers, - cutoff_square=self.cutoff_square, - half_skin_square=self.half_skin_square, - cutoff_with_skin=self.cutoff_with_skin, - half_cutoff_with_skin=self.half_cutoff_with_skin, - cutoff_with_skin_square=self.cutoff_with_skin_square, - refresh_interval=self.refresh_interval, cutoff=self.cutoff, - skin=self.skin, - max_atom_in_grid_numbers=self.max_atom_in_grid_numbers, - max_neighbor_numbers=self.max_neighbor_numbers) - - res = self.neighbor_list_update(self.atom_numbers_in_grid_bucket, self.bucket, crd, box_length, self.grid_N, - self.grid_length_inverse, self.atom_in_grid_serial, old_crd, - crd_to_uint_crd_cof, uint_crd, self.pointer, self.nl_atom_numbers, - self.nl_atom_serial, uint_dr_to_dr_cof, self.excluded_list_start, - self.excluded_list, self.excluded_numbers, self.need_refresh_flag) - return res diff --git a/model_zoo/research/hpc/sponge/src/particle_mesh_ewald.py b/model_zoo/research/hpc/sponge/src/particle_mesh_ewald.py index 833a680d6e..e1e45b4e35 100644 --- a/model_zoo/research/hpc/sponge/src/particle_mesh_ewald.py +++ b/model_zoo/research/hpc/sponge/src/particle_mesh_ewald.py @@ -12,20 +12,13 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================ -"""PME""" - +'''PME''' import math -import numpy as np -import mindspore.common.dtype as mstype -from mindspore import Tensor, nn -from mindspore.ops import operations as P -class Particle_Mesh_Ewald(nn.Cell): - """class Particle_Mesh_Ewald""" - +class Particle_Mesh_Ewald(): + '''PME''' def __init__(self, controller, md_info): - super(Particle_Mesh_Ewald, self).__init__() self.cutoff = 10.0 if "cut" not in controller.Command_Set else float(controller.Command_Set["cut"]) self.tolerance = 0.00001 if "PME_Direct_Tolerance" not in controller.Command_Set else float( controller.Command_Set["PME_Direct_Tolerance"]) @@ -43,12 +36,9 @@ class Particle_Mesh_Ewald(nn.Cell): self.fftz = self.Get_Fft_Patameter(self.box_length[2]) self.beta = self.Get_Beta(self.cutoff, self.tolerance) - self.box_length = Tensor(np.asarray(self.box_length, np.float32), mstype.float32) - - print("========== ", self.fftx, self.ffty, self.fftz, self.tolerance, self.beta) def Get_Beta(self, cutoff, tolerance): - """get beta""" + '''GET BETA''' high = 1.0 ihigh = 1 while 1: @@ -69,7 +59,7 @@ class Particle_Mesh_Ewald(nn.Cell): return beta def Check_2357_Factor(self, number): - """check 2357 factor""" + '''CHECK FACTOR''' while number > 0: if number == 1: return 1 @@ -101,7 +91,7 @@ class Particle_Mesh_Ewald(nn.Cell): return 0 def Get_Fft_Patameter(self, length): - """get fft parameter""" + '''GET FFT PARAMETER''' tempi = math.ceil(length + 3) >> 2 << 2 if 60 <= tempi <= 68: tempi = 64 @@ -117,31 +107,3 @@ class Particle_Mesh_Ewald(nn.Cell): if self.Check_2357_Factor(tempi): return tempi tempi += 4 - - def PME_Energy(self, uint_crd, charge, nl_atom_numbers, nl_atom_serial, uint_dr_to_dr_cof, excluded_list_start, - excluded_list, excluded_numbers, excluded_atom_numbers): - """PME_Energy""" - self.pmee = P.PMEEnergy(self.atom_numbers, excluded_atom_numbers, self.beta, self.fftx, self.ffty, self.fftz) - self.reciprocal_energy, self.self_energy, self.direct_energy, self.correction_energy = \ - self.pmee(self.box_length, uint_crd, charge, nl_atom_numbers, nl_atom_serial, uint_dr_to_dr_cof, - excluded_list_start, excluded_list, excluded_numbers) - return self.reciprocal_energy, self.self_energy, self.direct_energy, self.correction_energy - - def PME_Excluded_Force(self, uint_crd, scaler, charge, excluded_list_start, excluded_list, - excluded_numbers, excluded_atom_numbers): - """PME Excluded Force""" - self.pmeef = P.PMEExcludedForce(atom_numbers=self.atom_numbers, excluded_numbers=excluded_atom_numbers, - beta=self.beta) - self.frc = self.pmeef(uint_crd, scaler, charge, excluded_list_start, excluded_list, excluded_numbers) - return self.frc - - def PME_Reciprocal_Force(self, uint_crd, charge): - """PME reciprocal force""" - self.pmerf = P.PMEReciprocalForce(self.atom_numbers, self.beta, self.fftx, self.ffty, self.fftz) - self.frc = self.pmerf(self.box_length, uint_crd, charge) - return self.frc - - def Energy_Device_To_Host(self): - """Energy_Device_To_Host""" - self.ee_ene = self.reciprocal_energy + self.self_energy + self.direct_energy + self.correction_energy - return self.ee_ene diff --git a/model_zoo/research/hpc/sponge/src/simulation.py b/model_zoo/research/hpc/sponge/src/simulation.py new file mode 100644 index 0000000000..507eb67ee0 --- /dev/null +++ b/model_zoo/research/hpc/sponge/src/simulation.py @@ -0,0 +1,439 @@ +# Copyright 2021 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================ +'''Simulation''' +import numpy as np + +import mindspore.common.dtype as mstype +from mindspore import Tensor +from mindspore import nn +from mindspore.common.parameter import Parameter +from mindspore.ops import functional as F +from mindspore.ops import operations as P +from src.angle import Angle +from src.bond import Bond +from src.dihedral import Dihedral +from src.langevin_liujian_md import Langevin_Liujian +from src.lennard_jones import Lennard_Jones_Information +from src.md_information import md_information +from src.nb14 import NON_BOND_14 +from src.neighbor_list import neighbor_list +from src.particle_mesh_ewald import Particle_Mesh_Ewald + + +class controller: + '''controller''' + def __init__(self, args_opt): + self.input_file = args_opt.i + self.initial_coordinates_file = args_opt.c + self.amber_parm = args_opt.amber_parm + self.restrt = args_opt.r + self.mdcrd = args_opt.x + self.mdout = args_opt.o + self.mdbox = args_opt.box + + self.Command_Set = {} + self.md_task = None + self.commands_from_in_file() + + def commands_from_in_file(self): + '''command from in file''' + file = open(self.input_file, 'r') + context = file.readlines() + file.close() + self.md_task = context[0].strip() + for val in context: + if "=" in val: + assert len(val.strip().split("=")) == 2 + flag, value = val.strip().split("=") + value = value.replace(",", '') + flag = flag.replace(" ", "") + if flag not in self.Command_Set: + self.Command_Set[flag] = value + else: + print("ERROR COMMAND FILE") + + +class Simulation(nn.Cell): + '''simulation''' + def __init__(self, args_opt): + super(Simulation, self).__init__() + self.control = controller(args_opt) + self.md_info = md_information(self.control) + self.bond = Bond(self.control, self.md_info) + self.angle = Angle(self.control) + self.dihedral = Dihedral(self.control) + self.nb14 = NON_BOND_14(self.control, self.dihedral, self.md_info.atom_numbers) + self.nb_info = neighbor_list(self.control, self.md_info.atom_numbers, self.md_info.box_length) + self.LJ_info = Lennard_Jones_Information(self.control) + self.liujian_info = Langevin_Liujian(self.control, self.md_info.atom_numbers) + self.pme_method = Particle_Mesh_Ewald(self.control, self.md_info) + self.bond_energy_sum = Tensor(0, mstype.int32) + self.angle_energy_sum = Tensor(0, mstype.int32) + self.dihedral_energy_sum = Tensor(0, mstype.int32) + self.nb14_lj_energy_sum = Tensor(0, mstype.int32) + self.nb14_cf_energy_sum = Tensor(0, mstype.int32) + self.lj_energy_sum = Tensor(0, mstype.int32) + self.ee_ene = Tensor(0, mstype.int32) + self.total_energy = Tensor(0, mstype.int32) + # Init scalar + self.ntwx = self.md_info.ntwx + self.atom_numbers = self.md_info.atom_numbers + self.residue_numbers = self.md_info.residue_numbers + self.bond_numbers = self.bond.bond_numbers + self.angle_numbers = self.angle.angle_numbers + self.dihedral_numbers = self.dihedral.dihedral_numbers + self.nb14_numbers = self.nb14.nb14_numbers + self.Nxy = self.nb_info.Nxy + self.grid_numbers = self.nb_info.grid_numbers + self.max_atom_in_grid_numbers = self.nb_info.max_atom_in_grid_numbers + self.max_neighbor_numbers = self.nb_info.max_neighbor_numbers + self.excluded_atom_numbers = self.nb_info.excluded_atom_numbers + self.refresh_count = Parameter(Tensor(self.nb_info.refresh_count, mstype.int32), requires_grad=False) + self.refresh_interval = self.nb_info.refresh_interval + self.skin = self.nb_info.skin + self.cutoff = self.nb_info.cutoff + self.cutoff_square = self.nb_info.cutoff_square + self.cutoff_with_skin = self.nb_info.cutoff_with_skin + self.half_cutoff_with_skin = self.nb_info.half_cutoff_with_skin + self.cutoff_with_skin_square = self.nb_info.cutoff_with_skin_square + self.half_skin_square = self.nb_info.half_skin_square + self.beta = self.pme_method.beta + self.fftx = self.pme_method.fftx + self.ffty = self.pme_method.ffty + self.fftz = self.pme_method.fftz + self.random_seed = self.liujian_info.rand_seed + self.dt = self.liujian_info.dt + self.half_dt = self.liujian_info.half_dt + self.exp_gamma = self.liujian_info.exp_gamma + self.init_Tensor() + self.op_define() + + def init_Tensor(self): + '''init tensor''' + self.crd = Parameter( + Tensor(np.float32(np.asarray(self.md_info.coordinate).reshape([self.atom_numbers, 3])), mstype.float32), + requires_grad=False) + self.crd_to_uint_crd_cof = Tensor(np.asarray(self.md_info.crd_to_uint_crd_cof, np.float32), mstype.float32) + self.uint_dr_to_dr_cof = Parameter( + Tensor(np.asarray(self.md_info.uint_dr_to_dr_cof, np.float32), mstype.float32), requires_grad=False) + self.box_length = Tensor(self.md_info.box_length, mstype.float32) + self.charge = Tensor(np.asarray(self.md_info.h_charge, dtype=np.float32), mstype.float32) + self.old_crd = Parameter(Tensor(np.zeros([self.atom_numbers, 3], dtype=np.float32), mstype.float32), + requires_grad=False) + self.uint_crd = Parameter(Tensor(np.zeros([self.atom_numbers, 3], dtype=np.uint32), mstype.uint32), + requires_grad=False) + self.mass_inverse = Tensor(self.md_info.h_mass_inverse, mstype.float32) + self.res_start = Tensor(self.md_info.h_res_start, mstype.int32) + self.res_end = Tensor(self.md_info.h_res_end, mstype.int32) + self.mass = Tensor(self.md_info.h_mass, mstype.float32) + self.velocity = Parameter(Tensor(self.md_info.velocity, mstype.float32), requires_grad=False) + self.acc = Parameter(Tensor(np.zeros([self.atom_numbers, 3], np.float32), mstype.float32), requires_grad=False) + self.bond_atom_a = Tensor(np.asarray(self.bond.h_atom_a, np.int32), mstype.int32) + self.bond_atom_b = Tensor(np.asarray(self.bond.h_atom_b, np.int32), mstype.int32) + self.bond_k = Tensor(np.asarray(self.bond.h_k, np.float32), mstype.float32) + self.bond_r0 = Tensor(np.asarray(self.bond.h_r0, np.float32), mstype.float32) + self.angle_atom_a = Tensor(np.asarray(self.angle.h_atom_a, np.int32), mstype.int32) + self.angle_atom_b = Tensor(np.asarray(self.angle.h_atom_b, np.int32), mstype.int32) + self.angle_atom_c = Tensor(np.asarray(self.angle.h_atom_c, np.int32), mstype.int32) + self.angle_k = Tensor(np.asarray(self.angle.h_angle_k, np.float32), mstype.float32) + self.angle_theta0 = Tensor(np.asarray(self.angle.h_angle_theta0, np.float32), mstype.float32) + self.dihedral_atom_a = Tensor(np.asarray(self.dihedral.h_atom_a, np.int32), mstype.int32) + self.dihedral_atom_b = Tensor(np.asarray(self.dihedral.h_atom_b, np.int32), mstype.int32) + self.dihedral_atom_c = Tensor(np.asarray(self.dihedral.h_atom_c, np.int32), mstype.int32) + self.dihedral_atom_d = Tensor(np.asarray(self.dihedral.h_atom_d, np.int32), mstype.int32) + self.pk = Tensor(np.asarray(self.dihedral.pk, np.float32), mstype.float32) + self.gamc = Tensor(np.asarray(self.dihedral.gamc, np.float32), mstype.float32) + self.gams = Tensor(np.asarray(self.dihedral.gams, np.float32), mstype.float32) + self.pn = Tensor(np.asarray(self.dihedral.pn, np.float32), mstype.float32) + self.ipn = Tensor(np.asarray(self.dihedral.ipn, np.int32), mstype.int32) + self.nb14_atom_a = Tensor(np.asarray(self.nb14.h_atom_a, np.int32), mstype.int32) + self.nb14_atom_b = Tensor(np.asarray(self.nb14.h_atom_b, np.int32), mstype.int32) + self.lj_scale_factor = Tensor(np.asarray(self.nb14.h_lj_scale_factor, np.float32), mstype.float32) + self.cf_scale_factor = Tensor(np.asarray(self.nb14.h_cf_scale_factor, np.float32), mstype.float32) + self.grid_N = Tensor(self.nb_info.grid_N, mstype.int32) + self.grid_length_inverse = Tensor(self.nb_info.grid_length_inverse, mstype.float32) + self.bucket = Parameter(Tensor( + np.asarray(self.nb_info.bucket, np.int32).reshape([self.grid_numbers, self.max_atom_in_grid_numbers]), + mstype.int32), requires_grad=False) + self.atom_numbers_in_grid_bucket = Parameter(Tensor(self.nb_info.atom_numbers_in_grid_bucket, mstype.int32), + requires_grad=False) + self.atom_in_grid_serial = Parameter(Tensor(np.zeros([self.nb_info.atom_numbers,], np.int32), mstype.int32), + requires_grad=False) + self.pointer = Parameter( + Tensor(np.asarray(self.nb_info.pointer, np.int32).reshape([self.grid_numbers, 125]), mstype.int32), + requires_grad=False) + self.nl_atom_numbers = Parameter(Tensor(np.zeros([self.atom_numbers,], np.int32), mstype.int32), + requires_grad=False) + self.nl_atom_serial = Parameter( + Tensor(np.zeros([self.atom_numbers, self.max_neighbor_numbers], np.int32), mstype.int32), + requires_grad=False) + self.excluded_list_start = Tensor(np.asarray(self.nb_info.excluded_list_start, np.int32), mstype.int32) + self.excluded_list = Tensor(np.asarray(self.nb_info.excluded_list, np.int32), mstype.int32) + self.excluded_numbers = Tensor(np.asarray(self.nb_info.excluded_numbers, np.int32), mstype.int32) + self.need_refresh_flag = Tensor(np.asarray([0], np.int32), mstype.int32) + self.atom_LJ_type = Tensor(np.asarray(self.LJ_info.atom_LJ_type, dtype=np.int32), mstype.int32) + self.LJ_A = Tensor(np.asarray(self.LJ_info.LJ_A, dtype=np.float32), mstype.float32) + self.LJ_B = Tensor(np.asarray(self.LJ_info.LJ_B, dtype=np.float32), mstype.float32) + self.sqrt_mass = Tensor(self.liujian_info.h_sqrt_mass, mstype.float32) + self.rand_state = Parameter(Tensor(self.liujian_info.rand_state, mstype.float32)) + self.zero_fp_tensor = Tensor(np.asarray([0,], np.float32)) + + def op_define(self): + '''op define''' + self.crd_to_uint_crd = P.CrdToUintCrd(self.atom_numbers) + self.mdtemp = P.MDTemperature(self.residue_numbers, self.atom_numbers) + self.setup_random_state = P.MDIterationSetupRandState(self.atom_numbers, self.random_seed) + self.bond_force_with_atom_energy = P.BondForceWithAtomEnergy(bond_numbers=self.bond_numbers, + atom_numbers=self.atom_numbers) + self.angle_force_with_atom_energy = P.AngleForceWithAtomEnergy(angle_numbers=self.angle_numbers) + self.dihedral_force_with_atom_energy = P.DihedralForceWithAtomEnergy(dihedral_numbers=self.dihedral_numbers) + self.nb14_force_with_atom_energy = P.Dihedral14LJCFForceWithAtomEnergy(nb14_numbers=self.nb14_numbers, + atom_numbers=self.atom_numbers) + self.lj_force_pme_direct_force = P.LJForceWithPMEDirectForce(self.atom_numbers, self.cutoff, self.beta) + self.pme_excluded_force = P.PMEExcludedForce(atom_numbers=self.atom_numbers, + excluded_numbers=self.excluded_atom_numbers, beta=self.beta) + self.pme_reciprocal_force = P.PMEReciprocalForce(self.atom_numbers, self.beta, self.fftx, self.ffty, self.fftz, + self.md_info.box_length[0], self.md_info.box_length[1], + self.md_info.box_length[2]) + + self.bond_energy = P.BondEnergy(self.bond_numbers, self.atom_numbers) + self.angle_energy = P.AngleEnergy(self.angle_numbers) + self.dihedral_energy = P.DihedralEnergy(self.dihedral_numbers) + self.nb14_lj_energy = P.Dihedral14LJEnergy(self.nb14_numbers, self.atom_numbers) + self.nb14_cf_energy = P.Dihedral14CFEnergy(self.nb14_numbers, self.atom_numbers) + self.lj_energy = P.LJEnergy(self.atom_numbers, self.cutoff_square) + self.pme_energy = P.PMEEnergy(self.atom_numbers, self.excluded_atom_numbers, self.beta, self.fftx, self.ffty, + self.fftz, self.md_info.box_length[0], self.md_info.box_length[1], + self.md_info.box_length[2]) + + self.md_iteration_leap_frog_liujian = P.MDIterationLeapFrogLiujian(self.atom_numbers, self.half_dt, self.dt, + self.exp_gamma) + + self.neighbor_list_update_init = P.NeighborListUpdate(grid_numbers=self.grid_numbers, + atom_numbers=self.atom_numbers, not_first_time=0, + Nxy=self.Nxy, + excluded_atom_numbers=self.excluded_atom_numbers, + cutoff_square=self.cutoff_square, + half_skin_square=self.half_skin_square, + cutoff_with_skin=self.cutoff_with_skin, + half_cutoff_with_skin=self.half_cutoff_with_skin, + cutoff_with_skin_square=self.cutoff_with_skin_square, + refresh_interval=self.refresh_interval, + cutoff=self.cutoff, skin=self.skin, + max_atom_in_grid_numbers=self.max_atom_in_grid_numbers, + max_neighbor_numbers=self.max_neighbor_numbers) + + self.neighbor_list_update = P.NeighborListUpdate(grid_numbers=self.grid_numbers, atom_numbers=self.atom_numbers, + not_first_time=1, Nxy=self.Nxy, + excluded_atom_numbers=self.excluded_atom_numbers, + cutoff_square=self.cutoff_square, + half_skin_square=self.half_skin_square, + cutoff_with_skin=self.cutoff_with_skin, + half_cutoff_with_skin=self.half_cutoff_with_skin, + cutoff_with_skin_square=self.cutoff_with_skin_square, + refresh_interval=self.refresh_interval, cutoff=self.cutoff, + skin=self.skin, + max_atom_in_grid_numbers=self.max_atom_in_grid_numbers, + max_neighbor_numbers=self.max_neighbor_numbers) + self.random_force = Tensor(np.zeros([self.atom_numbers, 3], np.float32), mstype.float32) + + def Simulation_Beforce_Caculate_Force(self): + '''simulation before calculate force''' + crd_to_uint_crd_cof = 0.5 * self.crd_to_uint_crd_cof + uint_crd = self.crd_to_uint_crd(crd_to_uint_crd_cof, self.crd) + return uint_crd + + def Simulation_Caculate_Force(self, uint_crd, scaler, nl_atom_numbers, nl_atom_serial): + '''simulation calculate force''' + bond_force, _ = self.bond_force_with_atom_energy(uint_crd, scaler, self.bond_atom_a, + self.bond_atom_b, self.bond_k, self.bond_r0) + + angle_force, _ = self.angle_force_with_atom_energy(uint_crd, scaler, self.angle_atom_a, + self.angle_atom_b, self.angle_atom_c, + self.angle_k, self.angle_theta0) + + dihedral_force, _ = self.dihedral_force_with_atom_energy(uint_crd, scaler, + self.dihedral_atom_a, + self.dihedral_atom_b, + self.dihedral_atom_c, + self.dihedral_atom_d, self.ipn, + self.pk, self.gamc, self.gams, + self.pn) + + nb14_force, _ = self.nb14_force_with_atom_energy(uint_crd, self.atom_LJ_type, self.charge, + scaler, self.nb14_atom_a, self.nb14_atom_b, + self.lj_scale_factor, self.cf_scale_factor, + self.LJ_A, self.LJ_B) + + lj_force = self.lj_force_pme_direct_force(uint_crd, self.atom_LJ_type, self.charge, scaler, nl_atom_numbers, + nl_atom_serial, self.LJ_A, self.LJ_B) + pme_excluded_force = self.pme_excluded_force(uint_crd, scaler, self.charge, self.excluded_list_start, + self.excluded_list, self.excluded_numbers) + pme_reciprocal_force = self.pme_reciprocal_force(uint_crd, self.charge) + force = P.AddN()( + [bond_force, angle_force, dihedral_force, nb14_force, lj_force, pme_excluded_force, pme_reciprocal_force]) + return force + + def Simulation_Caculate_Energy(self, uint_crd, uint_dr_to_dr_cof): + '''simulation calculate energy''' + bond_energy = self.bond_energy(uint_crd, uint_dr_to_dr_cof, self.bond_atom_a, self.bond_atom_b, self.bond_k, + self.bond_r0) + bond_energy_sum = P.ReduceSum(True)(bond_energy) + + angle_energy = self.angle_energy(uint_crd, uint_dr_to_dr_cof, self.angle_atom_a, self.angle_atom_b, + self.angle_atom_c, self.angle_k, self.angle_theta0) + angle_energy_sum = P.ReduceSum(True)(angle_energy) + + dihedral_energy = self.dihedral_energy(uint_crd, uint_dr_to_dr_cof, self.dihedral_atom_a, self.dihedral_atom_b, + self.dihedral_atom_c, self.dihedral_atom_d, self.ipn, self.pk, self.gamc, + self.gams, self.pn) + dihedral_energy_sum = P.ReduceSum(True)(dihedral_energy) + + nb14_lj_energy = self.nb14_lj_energy(uint_crd, self.atom_LJ_type, self.charge, uint_dr_to_dr_cof, + self.nb14_atom_a, self.nb14_atom_b, self.lj_scale_factor, self.LJ_A, + self.LJ_B) + nb14_cf_energy = self.nb14_cf_energy(uint_crd, self.atom_LJ_type, self.charge, uint_dr_to_dr_cof, + self.nb14_atom_a, self.nb14_atom_b, self.cf_scale_factor) + nb14_lj_energy_sum = P.ReduceSum(True)(nb14_lj_energy) + nb14_cf_energy_sum = P.ReduceSum(True)(nb14_cf_energy) + + lj_energy = self.lj_energy(uint_crd, self.atom_LJ_type, self.charge, uint_dr_to_dr_cof, self.nl_atom_numbers, + self.nl_atom_serial, self.LJ_A, self.LJ_B) + lj_energy_sum = P.ReduceSum(True)(lj_energy) + + reciprocal_energy, self_energy, direct_energy, correction_energy = self.pme_energy(uint_crd, self.charge, + self.nl_atom_numbers, + self.nl_atom_serial, + uint_dr_to_dr_cof, + self.excluded_list_start, + self.excluded_list, + self.excluded_numbers) + ee_ene = reciprocal_energy + self_energy + direct_energy + correction_energy + total_energy = P.AddN()( + [bond_energy_sum, angle_energy_sum, dihedral_energy_sum, nb14_lj_energy_sum, nb14_cf_energy_sum, + lj_energy_sum, ee_ene]) + return bond_energy_sum, angle_energy_sum, dihedral_energy_sum, nb14_lj_energy_sum, nb14_cf_energy_sum, \ + lj_energy_sum, ee_ene, total_energy + + def Simulation_Temperature(self): + '''caculate temperature''' + res_ek_energy = self.mdtemp(self.res_start, self.res_end, self.velocity, self.mass) + temperature = P.ReduceSum()(res_ek_energy) + return temperature + + def Simulation_MDIterationLeapFrog_Liujian(self, inverse_mass, sqrt_mass_inverse, crd, frc, rand_state, random_frc): + '''simulation leap frog iteration liujian''' + crd = self.md_iteration_leap_frog_liujian(inverse_mass, sqrt_mass_inverse, self.velocity, crd, frc, self.acc, + rand_state, random_frc) + vel = F.depend(self.velocity, crd) + acc = F.depend(self.acc, crd) + return vel, crd, acc + + def construct(self, step, print_step): + '''construct''' + if step == 0: + res = self.neighbor_list_update_init(self.atom_numbers_in_grid_bucket, self.bucket, self.crd, + self.box_length, self.grid_N, self.grid_length_inverse, + self.atom_in_grid_serial, self.old_crd, self.crd_to_uint_crd_cof, + self.uint_crd, self.pointer, self.nl_atom_numbers, self.nl_atom_serial, + self.uint_dr_to_dr_cof, self.excluded_list_start, self.excluded_list, + self.excluded_numbers, self.need_refresh_flag, self.refresh_count) + self.nl_atom_numbers = F.depend(self.nl_atom_numbers, res) + self.nl_atom_serial = F.depend(self.nl_atom_serial, res) + self.uint_dr_to_dr_cof = F.depend(self.uint_dr_to_dr_cof, res) + self.old_crd = F.depend(self.old_crd, res) + self.atom_numbers_in_grid_bucket = F.depend(self.atom_numbers_in_grid_bucket, res) + self.bucket = F.depend(self.bucket, res) + self.atom_in_grid_serial = F.depend(self.atom_in_grid_serial, res) + self.pointer = F.depend(self.pointer, res) + uint_crd = F.depend(self.uint_crd, res) + + force = self.Simulation_Caculate_Force(uint_crd, self.uint_dr_to_dr_cof, self.nl_atom_numbers, + self.nl_atom_serial) + bond_energy_sum, angle_energy_sum, dihedral_energy_sum, nb14_lj_energy_sum, nb14_cf_energy_sum, \ + lj_energy_sum, ee_ene, total_energy = self.Simulation_Caculate_Energy(uint_crd, self.uint_dr_to_dr_cof) + temperature = self.Simulation_Temperature() + self.rand_state = self.setup_random_state() + self.velocity, self.crd, _ = self.Simulation_MDIterationLeapFrog_Liujian(self.mass_inverse, + self.sqrt_mass, self.crd, force, + self.rand_state, + self.random_force) + + res = self.neighbor_list_update(self.atom_numbers_in_grid_bucket, + self.bucket, + self.crd, + self.box_length, + self.grid_N, + self.grid_length_inverse, + self.atom_in_grid_serial, + self.old_crd, + self.crd_to_uint_crd_cof, + self.uint_crd, + self.pointer, + self.nl_atom_numbers, + self.nl_atom_serial, + self.uint_dr_to_dr_cof, + self.excluded_list_start, + self.excluded_list, + self.excluded_numbers, + self.need_refresh_flag, + self.refresh_count) + self.nl_atom_numbers = F.depend(self.nl_atom_numbers, res) + self.nl_atom_serial = F.depend(self.nl_atom_serial, res) + else: + uint_crd = self.Simulation_Beforce_Caculate_Force() + force = self.Simulation_Caculate_Force(uint_crd, self.uint_dr_to_dr_cof, self.nl_atom_numbers, + self.nl_atom_serial) + if print_step == 0: + bond_energy_sum, angle_energy_sum, dihedral_energy_sum, nb14_lj_energy_sum, nb14_cf_energy_sum, \ + lj_energy_sum, ee_ene, total_energy = self.Simulation_Caculate_Energy( + uint_crd, self.uint_dr_to_dr_cof) + else: + bond_energy_sum = self.zero_fp_tensor + angle_energy_sum = self.zero_fp_tensor + dihedral_energy_sum = self.zero_fp_tensor + nb14_lj_energy_sum = self.zero_fp_tensor + nb14_cf_energy_sum = self.zero_fp_tensor + lj_energy_sum = self.zero_fp_tensor + ee_ene = self.zero_fp_tensor + total_energy = self.zero_fp_tensor + temperature = self.Simulation_Temperature() + self.velocity, self.crd, _ = self.Simulation_MDIterationLeapFrog_Liujian(self.mass_inverse, + self.sqrt_mass, self.crd, force, + self.rand_state, + self.random_force) + res = self.neighbor_list_update(self.atom_numbers_in_grid_bucket, + self.bucket, + self.crd, + self.box_length, + self.grid_N, + self.grid_length_inverse, + self.atom_in_grid_serial, + self.old_crd, + self.crd_to_uint_crd_cof, + self.uint_crd, + self.pointer, + self.nl_atom_numbers, + self.nl_atom_serial, + self.uint_dr_to_dr_cof, + self.excluded_list_start, + self.excluded_list, + self.excluded_numbers, + self.need_refresh_flag, + self.refresh_count) + self.nl_atom_numbers = F.depend(self.nl_atom_numbers, res) + self.nl_atom_serial = F.depend(self.nl_atom_serial, res) + return temperature, total_energy, bond_energy_sum, angle_energy_sum, dihedral_energy_sum, nb14_lj_energy_sum, \ + nb14_cf_energy_sum, lj_energy_sum, ee_ene, res diff --git a/model_zoo/research/hpc/sponge/src/simulation_initial.py b/model_zoo/research/hpc/sponge/src/simulation_initial.py deleted file mode 100644 index 03771c8035..0000000000 --- a/model_zoo/research/hpc/sponge/src/simulation_initial.py +++ /dev/null @@ -1,245 +0,0 @@ -# Copyright 2021 Huawei Technologies Co., Ltd -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ============================================================================ -"""simulation""" - -import numpy as np -import mindspore.common.dtype as mstype -from mindspore import Tensor, nn - -from .Langevin_Liujian_md import Langevin_Liujian -from .angle import Angle -from .bond import Bond -from .dihedral import Dihedral -from .lennard_jones import Lennard_Jones_Information -from .md_information import md_information -from .nb14 import NON_BOND_14 -from .neighbor_list import nb_infomation -from .particle_mesh_ewald import Particle_Mesh_Ewald - - -class controller: - """class controller""" - def __init__(self, args_opt): - self.input_file = args_opt.i - self.initial_coordinates_file = args_opt.c - self.amber_parm = args_opt.amber_parm - self.restrt = args_opt.r - self.mdcrd = args_opt.x - self.mdout = args_opt.o - self.mdbox = args_opt.box - - self.Command_Set = {} - self.md_task = None - self.commands_from_in_file() - - def commands_from_in_file(self): - """commands from in file""" - file = open(self.input_file, 'r') - context = file.readlines() - file.close() - self.md_task = context[0].strip() - for val in context: - if "=" in val: - assert len(val.strip().split("=")) == 2 - flag, value = val.strip().split("=") - value = value.replace(",", '') - flag = flag.replace(" ", "") - if flag not in self.Command_Set: - self.Command_Set[flag] = value - else: - print("ERROR COMMAND FILE") - - -class Simulation(nn.Cell): - """class simulation""" - - def __init__(self, args_opt): - super(Simulation, self).__init__() - self.control = controller(args_opt) - self.md_info = md_information(self.control) - self.bond = Bond(self.control, self.md_info) - self.angle = Angle(self.control) - self.dihedral = Dihedral(self.control) - self.nb14 = NON_BOND_14(self.control, self.dihedral, self.md_info.atom_numbers) - self.nb_info = nb_infomation(self.control, self.md_info.atom_numbers, self.md_info.box_length) - self.LJ_info = Lennard_Jones_Information(self.control) - self.liujian_info = Langevin_Liujian(self.control, self.md_info.atom_numbers) - self.pme_method = Particle_Mesh_Ewald(self.control, self.md_info) - self.box_length = Tensor(np.asarray(self.md_info.box_length, np.float32), mstype.float32) - self.file = None - - def Main_Before_Calculate_Force(self): - """Main Before Calculate Force""" - _ = self.md_info.MD_Information_Crd_To_Uint_Crd() - self.md_info.uint_crd_with_LJ = (self.md_info.uint_crd, self.LJ_info.atom_LJ_type, self.md_info.charge) - return self.md_info.uint_crd, self.md_info.uint_crd_with_LJ - - def Initial_Neighbor_List_Update(self, not_first_time): - """Initial Neighbor List Update""" - res = self.nb_info.NeighborListUpdate(self.md_info.crd, self.md_info.crd_old, self.md_info.uint_crd, - self.md_info.crd_to_uint_crd_cof, self.md_info.uint_dr_to_dr_cof, - self.box_length, not_first_time) - - return res - - def Main_Calculate_Force(self): - """main calculate force""" - self.bond.atom_numbers = self.md_info.atom_numbers - md_info = self.md_info - LJ_info = self.LJ_info - nb_info = self.nb_info - 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) - 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) - frc_t += angle_frc.asnumpy() - - dihedral_frc, _ = self.dihedral.Dihedral_Force_With_Atom_Energy(md_info.uint_crd, md_info.uint_dr_to_dr_cof) - frc_t += dihedral_frc.asnumpy() - - nb14_frc, _ = self.nb14.Non_Bond_14_LJ_CF_Force_With_Atom_Energy(md_info.uint_crd_with_LJ, - md_info.uint_dr_to_dr_cof, LJ_info.LJ_A, - LJ_info.LJ_B) - frc_t += nb14_frc.asnumpy() - - lj_frc = LJ_info.LJ_Force_With_PME_Direct_Force( - md_info.atom_numbers, md_info.uint_crd_with_LJ, md_info.uint_dr_to_dr_cof, nb_info.nl_atom_numbers, - nb_info.nl_atom_serial, nb_info.cutoff, pme_method.beta) - frc_t += lj_frc.asnumpy() - - pme_excluded_frc = pme_method.PME_Excluded_Force( - md_info.uint_crd, md_info.uint_dr_to_dr_cof, md_info.charge, - nb_info.excluded_list_start, nb_info.excluded_list, - nb_info.excluded_numbers, nb_info.excluded_atom_numbers) - frc_t += pme_excluded_frc.asnumpy() - - pme_reciprocal_frc = pme_method.PME_Reciprocal_Force(md_info.uint_crd, md_info.charge) - frc_t += pme_reciprocal_frc.asnumpy() - - self.md_info.frc = Tensor(frc_t, mstype.float32) - - return self.md_info.frc - - def Main_Calculate_Energy(self): - """main calculate energy""" - _ = self.bond.Bond_Energy(self.md_info.uint_crd, self.md_info.uint_dr_to_dr_cof) - _ = self.angle.Angle_Energy(self.md_info.uint_crd, self.md_info.uint_dr_to_dr_cof) - _ = self.dihedral.Dihedral_Engergy(self.md_info.uint_crd, self.md_info.uint_dr_to_dr_cof) - _ = self.nb14.Non_Bond_14_LJ_CF_Energy(self.md_info.uint_crd_with_LJ, self.md_info.uint_dr_to_dr_cof, - self.LJ_info.LJ_A, - self.LJ_info.LJ_B) - - _ = self.LJ_info.LJ_Energy(self.md_info.uint_crd_with_LJ, self.md_info.uint_dr_to_dr_cof, - self.nb_info.nl_atom_numbers, self.nb_info.nl_atom_serial, - self.nb_info.cutoff_square) - _ = self.pme_method.PME_Energy( - self.md_info.uint_crd, self.md_info.charge, self.nb_info.nl_atom_numbers, self.nb_info.nl_atom_serial, - self.md_info.uint_dr_to_dr_cof, self.nb_info.excluded_list_start, self.nb_info.excluded_list, - self.nb_info.excluded_numbers, self.nb_info.excluded_atom_numbers) - _ = self.pme_method.Energy_Device_To_Host() - - def Main_After_Calculate_Energy(self): - """main after calculate energy""" - md_info = self.md_info - LJ_info = self.LJ_info - bond = self.bond - angle = self.angle - dihedral = self.dihedral - nb14 = self.nb14 - pme_method = self.pme_method - - md_info.total_potential_energy = 0 - md_info.total_potential_energy += bond.sigma_of_bond_ene - md_info.total_potential_energy += angle.sigma_of_angle_ene - md_info.total_potential_energy += dihedral.sigma_of_dihedral_ene - md_info.total_potential_energy += nb14.nb14_lj_energy_sum + nb14.nb14_cf_energy_sum - md_info.total_potential_energy += LJ_info.LJ_energy_sum - pme_method.Energy_Device_To_Host() - md_info.total_potential_energy += pme_method.ee_ene - print("md_info.total_potential_energy", md_info.total_potential_energy) - - def Main_Iteration_2(self): - """main iteration2""" - md_info = self.md_info - control = self.control - liujian_info = self.liujian_info - - 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.d_mass_inverse, md_info.vel, md_info.crd, md_info.frc) - self.Main_After_Iteration() - - def Main_After_Iteration(self): - """main after iteration""" - md_info = self.md_info - nb_info = self.nb_info - md_info.Centerize() - _ = nb_info.NeighborListUpdate(md_info.crd, md_info.crd_old, md_info.uint_crd, - md_info.crd_to_uint_crd_cof, - md_info.uint_dr_to_dr_cof, self.box_length, not_first_time=1) - - def Main_Print(self): - """compute the temperature""" - md_info = self.md_info - temperature = md_info.MD_Information_Temperature() - md_info.h_temperature = temperature - steps = md_info.steps - temperature = temperature.asnumpy() - total_potential_energy = md_info.total_potential_energy.asnumpy() - sigma_of_bond_ene = self.bond.sigma_of_bond_ene.asnumpy() - sigma_of_angle_ene = self.angle.sigma_of_angle_ene.asnumpy() - sigma_of_dihedral_ene = self.dihedral.sigma_of_dihedral_ene.asnumpy() - nb14_lj_energy_sum = self.nb14.nb14_lj_energy_sum.asnumpy() - nb14_cf_energy_sum = self.nb14.nb14_cf_energy_sum.asnumpy() - LJ_energy_sum = self.LJ_info.LJ_energy_sum.asnumpy() - ee_ene = self.pme_method.ee_ene.asnumpy() - print("_steps_ _TEMP_ _TOT_POT_ENE_ _BOND_ENE_ " - "_ANGLE_ENE_ _DIHEDRAL_ENE_ _14LJ_ENE_ _14CF_ENE_ _LJ_ENE_ _CF_PME_ENE_") - print("{:>7.0f} {:>7.3f} {:>11.3f}".format(steps, float(temperature), float(total_potential_energy)), end=" ") - if self.bond.bond_numbers > 0: - print("{:>10.3f}".format(float(sigma_of_bond_ene)), end=" ") - if self.angle.angle_numbers > 0: - print("{:>11.3f}".format(float(sigma_of_angle_ene)), end=" ") - if self.dihedral.dihedral_numbers > 0: - print("{:>14.3f}".format(float(sigma_of_dihedral_ene)), end=" ") - if self.nb14.nb14_numbers > 0: - print("{:>10.3f} {:>10.3f}".format(float(nb14_lj_energy_sum), float(nb14_cf_energy_sum)), end=" ") - - print("{:>7.3f}".format(float(LJ_energy_sum)), end=" ") - print("{:>12.3f}".format(float(ee_ene))) - - if self.file is not None: - self.file.write("{:>7.0f} {:>7.3f} {:>11.3f} {:>10.3f} {:>11.3f} {:>14.3f} {:>10.3f} {:>10.3f} {:>7.3f}" - " {:>12.3f}\n".format(steps, float(temperature), float(total_potential_energy), - float(sigma_of_bond_ene), float(sigma_of_angle_ene), - float(sigma_of_dihedral_ene), float(nb14_lj_energy_sum), - float(nb14_cf_energy_sum), float(LJ_energy_sum), float(ee_ene))) - - return temperature - - def Main_Initial(self): - """main initial""" - if self.control.mdout: - self.file = open(self.control.mdout, 'w') - self.file.write("_steps_ _TEMP_ _TOT_POT_ENE_ _BOND_ENE_ " - "_ANGLE_ENE_ _DIHEDRAL_ENE_ _14LJ_ENE_ _14CF_ENE_ _LJ_ENE_ _CF_PME_ENE_\n") - - def Main_Destroy(self): - """main destroy""" - if self.file is not None: - self.file.close() - print("Save successfully!")