Huawei_Technology
/
mindspore
Not watched
1
0
Fork 0
Code
Releases 13 Wiki evaluate Activity
Issues 0 Pull Requests 0 Datasets Model Cloudbrain HPC
You can not select more than 25 topics Topics must start with a chinese character,a letter or number, can include dashes ('-') and can be up to 35 characters long.
Tree: 97e3232b6d
Branches Tags
${ item.name }
Create branch ${ searchTerm }
from '97e3232b6d'
${ noResults }
mindspore/model_zoo/research/hpc/sponge/src/simulation.py

440 lines
29 kB
Raw Blame History 

  1. # Copyright 2021 Huawei Technologies Co., Ltd

  2. #

  3. # Licensed under the Apache License, Version 2.0 (the "License");

  4. # you may not use this file except in compliance with the License.

  5. # You may obtain a copy of the License at

  6. #

  7. # http://www.apache.org/licenses/LICENSE-2.0

  8. #

  9. # Unless required by applicable law or agreed to in writing, software

  10. # distributed under the License is distributed on an "AS IS" BASIS,

  11. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

  12. # See the License for the specific language governing permissions and

  13. # limitations under the License.

  14. # ============================================================================

  15. '''Simulation'''

  16. import numpy as np

  17. 

  18. import mindspore.common.dtype as mstype

  19. from mindspore import Tensor

  20. from mindspore import nn

  21. from mindspore.common.parameter import Parameter

  22. from mindspore.ops import functional as F

  23. from mindspore.ops import operations as P

  24. from src.angle import Angle

  25. from src.bond import Bond

  26. from src.dihedral import Dihedral

  27. from src.langevin_liujian_md import Langevin_Liujian

  28. from src.lennard_jones import Lennard_Jones_Information

  29. from src.md_information import md_information

  30. from src.nb14 import NON_BOND_14

  31. from src.neighbor_list import neighbor_list

  32. from src.particle_mesh_ewald import Particle_Mesh_Ewald

  33. 

  34. 

  35. class controller:

  36.     '''controller'''

  37.     def __init__(self, args_opt):

  38.         self.input_file = args_opt.i

  39.         self.initial_coordinates_file = args_opt.c

  40.         self.amber_parm = args_opt.amber_parm

  41.         self.restrt = args_opt.r

  42.         self.mdcrd = args_opt.x

  43.         self.mdout = args_opt.o

  44.         self.mdbox = args_opt.box

  45. 

  46.         self.Command_Set = {}

  47.         self.md_task = None

  48.         self.commands_from_in_file()

  49. 

  50.     def commands_from_in_file(self):

  51.         '''command from in file'''

  52.         file = open(self.input_file, 'r')

  53.         context = file.readlines()

  54.         file.close()

  55.         self.md_task = context[0].strip()

  56.         for val in context:

  57.             if "=" in val:

  58.                 assert len(val.strip().split("=")) == 2

  59.                 flag, value = val.strip().split("=")

  60.                 value = value.replace(",", '')

  61.                 flag = flag.replace(" ", "")

  62.                 if flag not in self.Command_Set:

  63.                     self.Command_Set[flag] = value

  64.                 else:

  65.                     print("ERROR COMMAND FILE")

  66. 

  67. 

  68. class Simulation(nn.Cell):

  69.     '''simulation'''

  70.     def __init__(self, args_opt):

  71.         super(Simulation, self).__init__()

  72.         self.control = controller(args_opt)

  73.         self.md_info = md_information(self.control)

  74.         self.bond = Bond(self.control, self.md_info)

  75.         self.angle = Angle(self.control)

  76.         self.dihedral = Dihedral(self.control)

  77.         self.nb14 = NON_BOND_14(self.control, self.dihedral, self.md_info.atom_numbers)

  78.         self.nb_info = neighbor_list(self.control, self.md_info.atom_numbers, self.md_info.box_length)

  79.         self.LJ_info = Lennard_Jones_Information(self.control)

  80.         self.liujian_info = Langevin_Liujian(self.control, self.md_info.atom_numbers)

  81.         self.pme_method = Particle_Mesh_Ewald(self.control, self.md_info)

  82.         self.bond_energy_sum = Tensor(0, mstype.int32)

  83.         self.angle_energy_sum = Tensor(0, mstype.int32)

  84.         self.dihedral_energy_sum = Tensor(0, mstype.int32)

  85.         self.nb14_lj_energy_sum = Tensor(0, mstype.int32)

  86.         self.nb14_cf_energy_sum = Tensor(0, mstype.int32)

  87.         self.lj_energy_sum = Tensor(0, mstype.int32)

  88.         self.ee_ene = Tensor(0, mstype.int32)

  89.         self.total_energy = Tensor(0, mstype.int32)

  90.         # Init scalar

  91.         self.ntwx = self.md_info.ntwx

  92.         self.atom_numbers = self.md_info.atom_numbers

  93.         self.residue_numbers = self.md_info.residue_numbers

  94.         self.bond_numbers = self.bond.bond_numbers

  95.         self.angle_numbers = self.angle.angle_numbers

  96.         self.dihedral_numbers = self.dihedral.dihedral_numbers

  97.         self.nb14_numbers = self.nb14.nb14_numbers

  98.         self.Nxy = self.nb_info.Nxy

  99.         self.grid_numbers = self.nb_info.grid_numbers

  100.         self.max_atom_in_grid_numbers = self.nb_info.max_atom_in_grid_numbers

  101.         self.max_neighbor_numbers = self.nb_info.max_neighbor_numbers

  102.         self.excluded_atom_numbers = self.nb_info.excluded_atom_numbers

  103.         self.refresh_count = Parameter(Tensor(self.nb_info.refresh_count, mstype.int32), requires_grad=False)

  104.         self.refresh_interval = self.nb_info.refresh_interval

  105.         self.skin = self.nb_info.skin

  106.         self.cutoff = self.nb_info.cutoff

  107.         self.cutoff_square = self.nb_info.cutoff_square

  108.         self.cutoff_with_skin = self.nb_info.cutoff_with_skin

  109.         self.half_cutoff_with_skin = self.nb_info.half_cutoff_with_skin

  110.         self.cutoff_with_skin_square = self.nb_info.cutoff_with_skin_square

  111.         self.half_skin_square = self.nb_info.half_skin_square

  112.         self.beta = self.pme_method.beta

  113.         self.fftx = self.pme_method.fftx

  114.         self.ffty = self.pme_method.ffty

  115.         self.fftz = self.pme_method.fftz

  116.         self.random_seed = self.liujian_info.rand_seed

  117.         self.dt = self.liujian_info.dt

  118.         self.half_dt = self.liujian_info.half_dt

  119.         self.exp_gamma = self.liujian_info.exp_gamma

  120.         self.init_Tensor()

  121.         self.op_define()

  122. 

  123.     def init_Tensor(self):

  124.         '''init tensor'''

  125.         self.crd = Parameter(

  126.             Tensor(np.float32(np.asarray(self.md_info.coordinate).reshape([self.atom_numbers, 3])), mstype.float32),

  127.             requires_grad=False)

  128.         self.crd_to_uint_crd_cof = Tensor(np.asarray(self.md_info.crd_to_uint_crd_cof, np.float32), mstype.float32)

  129.         self.uint_dr_to_dr_cof = Parameter(

  130.             Tensor(np.asarray(self.md_info.uint_dr_to_dr_cof, np.float32), mstype.float32), requires_grad=False)

  131.         self.box_length = Tensor(self.md_info.box_length, mstype.float32)

  132.         self.charge = Tensor(np.asarray(self.md_info.h_charge, dtype=np.float32), mstype.float32)

  133.         self.old_crd = Parameter(Tensor(np.zeros([self.atom_numbers, 3], dtype=np.float32), mstype.float32),

  134.                                  requires_grad=False)

  135.         self.uint_crd = Parameter(Tensor(np.zeros([self.atom_numbers, 3], dtype=np.uint32), mstype.uint32),

  136.                                   requires_grad=False)

  137.         self.mass_inverse = Tensor(self.md_info.h_mass_inverse, mstype.float32)

  138.         self.res_start = Tensor(self.md_info.h_res_start, mstype.int32)

  139.         self.res_end = Tensor(self.md_info.h_res_end, mstype.int32)

  140.         self.mass = Tensor(self.md_info.h_mass, mstype.float32)

  141.         self.velocity = Parameter(Tensor(self.md_info.velocity, mstype.float32), requires_grad=False)

  142.         self.acc = Parameter(Tensor(np.zeros([self.atom_numbers, 3], np.float32), mstype.float32), requires_grad=False)

  143.         self.bond_atom_a = Tensor(np.asarray(self.bond.h_atom_a, np.int32), mstype.int32)

  144.         self.bond_atom_b = Tensor(np.asarray(self.bond.h_atom_b, np.int32), mstype.int32)

  145.         self.bond_k = Tensor(np.asarray(self.bond.h_k, np.float32), mstype.float32)

  146.         self.bond_r0 = Tensor(np.asarray(self.bond.h_r0, np.float32), mstype.float32)

  147.         self.angle_atom_a = Tensor(np.asarray(self.angle.h_atom_a, np.int32), mstype.int32)

  148.         self.angle_atom_b = Tensor(np.asarray(self.angle.h_atom_b, np.int32), mstype.int32)

  149.         self.angle_atom_c = Tensor(np.asarray(self.angle.h_atom_c, np.int32), mstype.int32)

  150.         self.angle_k = Tensor(np.asarray(self.angle.h_angle_k, np.float32), mstype.float32)

  151.         self.angle_theta0 = Tensor(np.asarray(self.angle.h_angle_theta0, np.float32), mstype.float32)

  152.         self.dihedral_atom_a = Tensor(np.asarray(self.dihedral.h_atom_a, np.int32), mstype.int32)

  153.         self.dihedral_atom_b = Tensor(np.asarray(self.dihedral.h_atom_b, np.int32), mstype.int32)

  154.         self.dihedral_atom_c = Tensor(np.asarray(self.dihedral.h_atom_c, np.int32), mstype.int32)

  155.         self.dihedral_atom_d = Tensor(np.asarray(self.dihedral.h_atom_d, np.int32), mstype.int32)

  156.         self.pk = Tensor(np.asarray(self.dihedral.pk, np.float32), mstype.float32)

  157.         self.gamc = Tensor(np.asarray(self.dihedral.gamc, np.float32), mstype.float32)

  158.         self.gams = Tensor(np.asarray(self.dihedral.gams, np.float32), mstype.float32)

  159.         self.pn = Tensor(np.asarray(self.dihedral.pn, np.float32), mstype.float32)

  160.         self.ipn = Tensor(np.asarray(self.dihedral.ipn, np.int32), mstype.int32)

  161.         self.nb14_atom_a = Tensor(np.asarray(self.nb14.h_atom_a, np.int32), mstype.int32)

  162.         self.nb14_atom_b = Tensor(np.asarray(self.nb14.h_atom_b, np.int32), mstype.int32)

  163.         self.lj_scale_factor = Tensor(np.asarray(self.nb14.h_lj_scale_factor, np.float32), mstype.float32)

  164.         self.cf_scale_factor = Tensor(np.asarray(self.nb14.h_cf_scale_factor, np.float32), mstype.float32)

  165.         self.grid_N = Tensor(self.nb_info.grid_N, mstype.int32)

  166.         self.grid_length_inverse = Tensor(self.nb_info.grid_length_inverse, mstype.float32)

  167.         self.bucket = Parameter(Tensor(

  168.             np.asarray(self.nb_info.bucket, np.int32).reshape([self.grid_numbers, self.max_atom_in_grid_numbers]),

  169.             mstype.int32), requires_grad=False)

  170.         self.atom_numbers_in_grid_bucket = Parameter(Tensor(self.nb_info.atom_numbers_in_grid_bucket, mstype.int32),

  171.                                                      requires_grad=False)

  172.         self.atom_in_grid_serial = Parameter(Tensor(np.zeros([self.nb_info.atom_numbers,], np.int32), mstype.int32),

  173.                                              requires_grad=False)

  174.         self.pointer = Parameter(

  175.             Tensor(np.asarray(self.nb_info.pointer, np.int32).reshape([self.grid_numbers, 125]), mstype.int32),

  176.             requires_grad=False)

  177.         self.nl_atom_numbers = Parameter(Tensor(np.zeros([self.atom_numbers,], np.int32), mstype.int32),

  178.                                          requires_grad=False)

  179.         self.nl_atom_serial = Parameter(

  180.             Tensor(np.zeros([self.atom_numbers, self.max_neighbor_numbers], np.int32), mstype.int32),

  181.             requires_grad=False)

  182.         self.excluded_list_start = Tensor(np.asarray(self.nb_info.excluded_list_start, np.int32), mstype.int32)

  183.         self.excluded_list = Tensor(np.asarray(self.nb_info.excluded_list, np.int32), mstype.int32)

  184.         self.excluded_numbers = Tensor(np.asarray(self.nb_info.excluded_numbers, np.int32), mstype.int32)

  185.         self.need_refresh_flag = Tensor(np.asarray([0], np.int32), mstype.int32)

  186.         self.atom_LJ_type = Tensor(np.asarray(self.LJ_info.atom_LJ_type, dtype=np.int32), mstype.int32)

  187.         self.LJ_A = Tensor(np.asarray(self.LJ_info.LJ_A, dtype=np.float32), mstype.float32)

  188.         self.LJ_B = Tensor(np.asarray(self.LJ_info.LJ_B, dtype=np.float32), mstype.float32)

  189.         self.sqrt_mass = Tensor(self.liujian_info.h_sqrt_mass, mstype.float32)

  190.         self.rand_state = Parameter(Tensor(self.liujian_info.rand_state, mstype.float32))

  191.         self.zero_fp_tensor = Tensor(np.asarray([0,], np.float32))

  192. 

  193.     def op_define(self):

  194.         '''op define'''

  195.         self.crd_to_uint_crd = P.CrdToUintCrd(self.atom_numbers)

  196.         self.mdtemp = P.MDTemperature(self.residue_numbers, self.atom_numbers)

  197.         self.setup_random_state = P.MDIterationSetupRandState(self.atom_numbers, self.random_seed)

  198.         self.bond_force_with_atom_energy = P.BondForceWithAtomEnergy(bond_numbers=self.bond_numbers,

  199.                                                                      atom_numbers=self.atom_numbers)

  200.         self.angle_force_with_atom_energy = P.AngleForceWithAtomEnergy(angle_numbers=self.angle_numbers)

  201.         self.dihedral_force_with_atom_energy = P.DihedralForceWithAtomEnergy(dihedral_numbers=self.dihedral_numbers)

  202.         self.nb14_force_with_atom_energy = P.Dihedral14LJCFForceWithAtomEnergy(nb14_numbers=self.nb14_numbers,

  203.                                                                                atom_numbers=self.atom_numbers)

  204.         self.lj_force_pme_direct_force = P.LJForceWithPMEDirectForce(self.atom_numbers, self.cutoff, self.beta)

  205.         self.pme_excluded_force = P.PMEExcludedForce(atom_numbers=self.atom_numbers,

  206.                                                      excluded_numbers=self.excluded_atom_numbers, beta=self.beta)

  207.         self.pme_reciprocal_force = P.PMEReciprocalForce(self.atom_numbers, self.beta, self.fftx, self.ffty, self.fftz,

  208.                                                          self.md_info.box_length[0], self.md_info.box_length[1],

  209.                                                          self.md_info.box_length[2])

  210. 

  211.         self.bond_energy = P.BondEnergy(self.bond_numbers, self.atom_numbers)

  212.         self.angle_energy = P.AngleEnergy(self.angle_numbers)

  213.         self.dihedral_energy = P.DihedralEnergy(self.dihedral_numbers)

  214.         self.nb14_lj_energy = P.Dihedral14LJEnergy(self.nb14_numbers, self.atom_numbers)

  215.         self.nb14_cf_energy = P.Dihedral14CFEnergy(self.nb14_numbers, self.atom_numbers)

  216.         self.lj_energy = P.LJEnergy(self.atom_numbers, self.cutoff_square)

  217.         self.pme_energy = P.PMEEnergy(self.atom_numbers, self.excluded_atom_numbers, self.beta, self.fftx, self.ffty,

  218.                                       self.fftz, self.md_info.box_length[0], self.md_info.box_length[1],

  219.                                       self.md_info.box_length[2])

  220. 

  221.         self.md_iteration_leap_frog_liujian = P.MDIterationLeapFrogLiujian(self.atom_numbers, self.half_dt, self.dt,

  222.                                                                            self.exp_gamma)

  223. 

  224.         self.neighbor_list_update_init = P.NeighborListUpdate(grid_numbers=self.grid_numbers,

  225.                                                               atom_numbers=self.atom_numbers, not_first_time=0,

  226.                                                               Nxy=self.Nxy,

  227.                                                               excluded_atom_numbers=self.excluded_atom_numbers,

  228.                                                               cutoff_square=self.cutoff_square,

  229.                                                               half_skin_square=self.half_skin_square,

  230.                                                               cutoff_with_skin=self.cutoff_with_skin,

  231.                                                               half_cutoff_with_skin=self.half_cutoff_with_skin,

  232.                                                               cutoff_with_skin_square=self.cutoff_with_skin_square,

  233.                                                               refresh_interval=self.refresh_interval,

  234.                                                               cutoff=self.cutoff, skin=self.skin,

  235.                                                               max_atom_in_grid_numbers=self.max_atom_in_grid_numbers,

  236.                                                               max_neighbor_numbers=self.max_neighbor_numbers)

  237. 

  238.         self.neighbor_list_update = P.NeighborListUpdate(grid_numbers=self.grid_numbers, atom_numbers=self.atom_numbers,

  239.                                                          not_first_time=1, Nxy=self.Nxy,

  240.                                                          excluded_atom_numbers=self.excluded_atom_numbers,

  241.                                                          cutoff_square=self.cutoff_square,

  242.                                                          half_skin_square=self.half_skin_square,

  243.                                                          cutoff_with_skin=self.cutoff_with_skin,

  244.                                                          half_cutoff_with_skin=self.half_cutoff_with_skin,

  245.                                                          cutoff_with_skin_square=self.cutoff_with_skin_square,

  246.                                                          refresh_interval=self.refresh_interval, cutoff=self.cutoff,

  247.                                                          skin=self.skin,

  248.                                                          max_atom_in_grid_numbers=self.max_atom_in_grid_numbers,

  249.                                                          max_neighbor_numbers=self.max_neighbor_numbers)

  250.         self.random_force = Tensor(np.zeros([self.atom_numbers, 3], np.float32), mstype.float32)

  251. 

  252.     def Simulation_Beforce_Caculate_Force(self):

  253.         '''simulation before calculate force'''

  254.         crd_to_uint_crd_cof = 0.5 * self.crd_to_uint_crd_cof

  255.         uint_crd = self.crd_to_uint_crd(crd_to_uint_crd_cof, self.crd)

  256.         return uint_crd

  257. 

  258.     def Simulation_Caculate_Force(self, uint_crd, scaler, nl_atom_numbers, nl_atom_serial):

  259.         '''simulation calculate force'''

  260.         bond_force, _ = self.bond_force_with_atom_energy(uint_crd, scaler, self.bond_atom_a,

  261.                                                          self.bond_atom_b, self.bond_k, self.bond_r0)

  262. 

  263.         angle_force, _ = self.angle_force_with_atom_energy(uint_crd, scaler, self.angle_atom_a,

  264.                                                            self.angle_atom_b, self.angle_atom_c,

  265.                                                            self.angle_k, self.angle_theta0)

  266. 

  267.         dihedral_force, _ = self.dihedral_force_with_atom_energy(uint_crd, scaler,

  268.                                                                  self.dihedral_atom_a,

  269.                                                                  self.dihedral_atom_b,

  270.                                                                  self.dihedral_atom_c,

  271.                                                                  self.dihedral_atom_d, self.ipn,

  272.                                                                  self.pk, self.gamc, self.gams,

  273.                                                                  self.pn)

  274. 

  275.         nb14_force, _ = self.nb14_force_with_atom_energy(uint_crd, self.atom_LJ_type, self.charge,

  276.                                                          scaler, self.nb14_atom_a, self.nb14_atom_b,

  277.                                                          self.lj_scale_factor, self.cf_scale_factor,

  278.                                                          self.LJ_A, self.LJ_B)

  279. 

  280.         lj_force = self.lj_force_pme_direct_force(uint_crd, self.atom_LJ_type, self.charge, scaler, nl_atom_numbers,

  281.                                                   nl_atom_serial, self.LJ_A, self.LJ_B)

  282.         pme_excluded_force = self.pme_excluded_force(uint_crd, scaler, self.charge, self.excluded_list_start,

  283.                                                      self.excluded_list, self.excluded_numbers)

  284.         pme_reciprocal_force = self.pme_reciprocal_force(uint_crd, self.charge)

  285.         force = P.AddN()(

  286.             [bond_force, angle_force, dihedral_force, nb14_force, lj_force, pme_excluded_force, pme_reciprocal_force])

  287.         return force

  288. 

  289.     def Simulation_Caculate_Energy(self, uint_crd, uint_dr_to_dr_cof):

  290.         '''simulation calculate energy'''

  291.         bond_energy = self.bond_energy(uint_crd, uint_dr_to_dr_cof, self.bond_atom_a, self.bond_atom_b, self.bond_k,

  292.                                        self.bond_r0)

  293.         bond_energy_sum = P.ReduceSum(True)(bond_energy)

  294. 

  295.         angle_energy = self.angle_energy(uint_crd, uint_dr_to_dr_cof, self.angle_atom_a, self.angle_atom_b,

  296.                                          self.angle_atom_c, self.angle_k, self.angle_theta0)

  297.         angle_energy_sum = P.ReduceSum(True)(angle_energy)

  298. 

  299.         dihedral_energy = self.dihedral_energy(uint_crd, uint_dr_to_dr_cof, self.dihedral_atom_a, self.dihedral_atom_b,

  300.                                                self.dihedral_atom_c, self.dihedral_atom_d, self.ipn, self.pk, self.gamc,

  301.                                                self.gams, self.pn)

  302.         dihedral_energy_sum = P.ReduceSum(True)(dihedral_energy)

  303. 

  304.         nb14_lj_energy = self.nb14_lj_energy(uint_crd, self.atom_LJ_type, self.charge, uint_dr_to_dr_cof,

  305.                                              self.nb14_atom_a, self.nb14_atom_b, self.lj_scale_factor, self.LJ_A,

  306.                                              self.LJ_B)

  307.         nb14_cf_energy = self.nb14_cf_energy(uint_crd, self.atom_LJ_type, self.charge, uint_dr_to_dr_cof,

  308.                                              self.nb14_atom_a, self.nb14_atom_b, self.cf_scale_factor)

  309.         nb14_lj_energy_sum = P.ReduceSum(True)(nb14_lj_energy)

  310.         nb14_cf_energy_sum = P.ReduceSum(True)(nb14_cf_energy)

  311. 

  312.         lj_energy = self.lj_energy(uint_crd, self.atom_LJ_type, self.charge, uint_dr_to_dr_cof, self.nl_atom_numbers,

  313.                                    self.nl_atom_serial, self.LJ_A, self.LJ_B)

  314.         lj_energy_sum = P.ReduceSum(True)(lj_energy)

  315. 

  316.         reciprocal_energy, self_energy, direct_energy, correction_energy = self.pme_energy(uint_crd, self.charge,

  317.                                                                                            self.nl_atom_numbers,

  318.                                                                                            self.nl_atom_serial,

  319.                                                                                            uint_dr_to_dr_cof,

  320.                                                                                            self.excluded_list_start,

  321.                                                                                            self.excluded_list,

  322.                                                                                            self.excluded_numbers)

  323.         ee_ene = reciprocal_energy + self_energy + direct_energy + correction_energy

  324.         total_energy = P.AddN()(

  325.             [bond_energy_sum, angle_energy_sum, dihedral_energy_sum, nb14_lj_energy_sum, nb14_cf_energy_sum,

  326.              lj_energy_sum, ee_ene])

  327.         return bond_energy_sum, angle_energy_sum, dihedral_energy_sum, nb14_lj_energy_sum, nb14_cf_energy_sum, \

  328.                lj_energy_sum, ee_ene, total_energy

  329. 

  330.     def Simulation_Temperature(self):

  331.         '''caculate temperature'''

  332.         res_ek_energy = self.mdtemp(self.res_start, self.res_end, self.velocity, self.mass)

  333.         temperature = P.ReduceSum()(res_ek_energy)

  334.         return temperature

  335. 

  336.     def Simulation_MDIterationLeapFrog_Liujian(self, inverse_mass, sqrt_mass_inverse, crd, frc, rand_state, random_frc):

  337.         '''simulation leap frog iteration liujian'''

  338.         crd = self.md_iteration_leap_frog_liujian(inverse_mass, sqrt_mass_inverse, self.velocity, crd, frc, self.acc,

  339.                                                   rand_state, random_frc)

  340.         vel = F.depend(self.velocity, crd)

  341.         acc = F.depend(self.acc, crd)

  342.         return vel, crd, acc

  343. 

  344.     def construct(self, step, print_step):

  345.         '''construct'''

  346.         if step == 0:

  347.             res = self.neighbor_list_update_init(self.atom_numbers_in_grid_bucket, self.bucket, self.crd,

  348.                                                  self.box_length, self.grid_N, self.grid_length_inverse,

  349.                                                  self.atom_in_grid_serial, self.old_crd, self.crd_to_uint_crd_cof,

  350.                                                  self.uint_crd, self.pointer, self.nl_atom_numbers, self.nl_atom_serial,

  351.                                                  self.uint_dr_to_dr_cof, self.excluded_list_start, self.excluded_list,

  352.                                                  self.excluded_numbers, self.need_refresh_flag, self.refresh_count)

  353.             self.nl_atom_numbers = F.depend(self.nl_atom_numbers, res)

  354.             self.nl_atom_serial = F.depend(self.nl_atom_serial, res)

  355.             self.uint_dr_to_dr_cof = F.depend(self.uint_dr_to_dr_cof, res)

  356.             self.old_crd = F.depend(self.old_crd, res)

  357.             self.atom_numbers_in_grid_bucket = F.depend(self.atom_numbers_in_grid_bucket, res)

  358.             self.bucket = F.depend(self.bucket, res)

  359.             self.atom_in_grid_serial = F.depend(self.atom_in_grid_serial, res)

  360.             self.pointer = F.depend(self.pointer, res)

  361.             uint_crd = F.depend(self.uint_crd, res)

  362. 

  363.             force = self.Simulation_Caculate_Force(uint_crd, self.uint_dr_to_dr_cof, self.nl_atom_numbers,

  364.                                                    self.nl_atom_serial)

  365.             bond_energy_sum, angle_energy_sum, dihedral_energy_sum, nb14_lj_energy_sum, nb14_cf_energy_sum, \

  366.             lj_energy_sum, ee_ene, total_energy = self.Simulation_Caculate_Energy(uint_crd, self.uint_dr_to_dr_cof)

  367.             temperature = self.Simulation_Temperature()

  368.             self.rand_state = self.setup_random_state()

  369.             self.velocity, self.crd, _ = self.Simulation_MDIterationLeapFrog_Liujian(self.mass_inverse,

  370.                                                                                      self.sqrt_mass, self.crd, force,

  371.                                                                                      self.rand_state,

  372.                                                                                      self.random_force)

  373. 

  374.             res = self.neighbor_list_update(self.atom_numbers_in_grid_bucket,

  375.                                             self.bucket,

  376.                                             self.crd,

  377.                                             self.box_length,

  378.                                             self.grid_N,

  379.                                             self.grid_length_inverse,

  380.                                             self.atom_in_grid_serial,

  381.                                             self.old_crd,

  382.                                             self.crd_to_uint_crd_cof,

  383.                                             self.uint_crd,

  384.                                             self.pointer,

  385.                                             self.nl_atom_numbers,

  386.                                             self.nl_atom_serial,

  387.                                             self.uint_dr_to_dr_cof,

  388.                                             self.excluded_list_start,

  389.                                             self.excluded_list,

  390.                                             self.excluded_numbers,

  391.                                             self.need_refresh_flag,

  392.                                             self.refresh_count)

  393.             self.nl_atom_numbers = F.depend(self.nl_atom_numbers, res)

  394.             self.nl_atom_serial = F.depend(self.nl_atom_serial, res)

  395.         else:

  396.             uint_crd = self.Simulation_Beforce_Caculate_Force()

  397.             force = self.Simulation_Caculate_Force(uint_crd, self.uint_dr_to_dr_cof, self.nl_atom_numbers,

  398.                                                    self.nl_atom_serial)

  399.             if print_step == 0:

  400.                 bond_energy_sum, angle_energy_sum, dihedral_energy_sum, nb14_lj_energy_sum, nb14_cf_energy_sum, \

  401.                 lj_energy_sum, ee_ene, total_energy = self.Simulation_Caculate_Energy(

  402.                     uint_crd, self.uint_dr_to_dr_cof)

  403.             else:

  404.                 bond_energy_sum = self.zero_fp_tensor

  405.                 angle_energy_sum = self.zero_fp_tensor

  406.                 dihedral_energy_sum = self.zero_fp_tensor

  407.                 nb14_lj_energy_sum = self.zero_fp_tensor

  408.                 nb14_cf_energy_sum = self.zero_fp_tensor

  409.                 lj_energy_sum = self.zero_fp_tensor

  410.                 ee_ene = self.zero_fp_tensor

  411.                 total_energy = self.zero_fp_tensor

  412.             temperature = self.Simulation_Temperature()

  413.             self.velocity, self.crd, _ = self.Simulation_MDIterationLeapFrog_Liujian(self.mass_inverse,

  414.                                                                                      self.sqrt_mass, self.crd, force,

  415.                                                                                      self.rand_state,

  416.                                                                                      self.random_force)

  417.             res = self.neighbor_list_update(self.atom_numbers_in_grid_bucket,

  418.                                             self.bucket,

  419.                                             self.crd,

  420.                                             self.box_length,

  421.                                             self.grid_N,

  422.                                             self.grid_length_inverse,

  423.                                             self.atom_in_grid_serial,

  424.                                             self.old_crd,

  425.                                             self.crd_to_uint_crd_cof,

  426.                                             self.uint_crd,

  427.                                             self.pointer,

  428.                                             self.nl_atom_numbers,

  429.                                             self.nl_atom_serial,

  430.                                             self.uint_dr_to_dr_cof,

  431.                                             self.excluded_list_start,

  432.                                             self.excluded_list,

  433.                                             self.excluded_numbers,

  434.                                             self.need_refresh_flag,

  435.                                             self.refresh_count)

  436.             self.nl_atom_numbers = F.depend(self.nl_atom_numbers, res)

  437.             self.nl_atom_serial = F.depend(self.nl_atom_serial, res)

  438.         return temperature, total_energy, bond_energy_sum, angle_energy_sum, dihedral_energy_sum, nb14_lj_energy_sum, \

  439.                nb14_cf_energy_sum, lj_energy_sum, ee_ene, res