def test_batchnorm_fold_freeze():
c = 64
x = np.random.uniform(1, 10, size=[3, c, 32, 32]).astype('float32')
mean = np.random.uniform(1, 10, size=[c]).astype('float32')
variance = np.random.uniform(1, 10, size=[c]).astype('float32')
current_step = np.array([10]).astype('int32')
ms_mean_t = Tensor(mean)
ms_var_t = Tensor(variance)
ms_mean = Parameter(initializer(ms_mean_t, ms_mean_t.shape), name='mean')
ms_var = Parameter(initializer(ms_var_t, ms_var_t.shape), name='var')
net = Net(ms_mean, ms_var)
batch_mean, batch_var, delay_mean, delay_std = net(Tensor(x), Tensor(current_step))
_, _, _, _, expect5, expect6 = np_result(x, mean, variance, 0.9, 1e-12)
assert np.allclose(batch_mean.asnumpy(), np.zeros_like(mean), rtol=1.e-7, atol=1.e-5)
assert np.allclose(batch_var.asnumpy(), np.ones_like(mean), rtol=1.e-7, atol=1.e-5)
assert np.allclose(ms_mean.asnumpy(), mean, rtol=1.e-7, atol=1.e-5)
assert np.allclose(ms_var.asnumpy(), variance, rtol=1.e-7, atol=1.e-5)
assert np.allclose(delay_mean.asnumpy(), expect5, rtol=1.e-7, atol=1.e-5)
assert np.allclose(delay_std.asnumpy(), expect6, rtol=1.e-7, atol=1.e-5)
def test_rmsprop():
learning_rate, decay, momentum, epsilon, centered = [
0.5, 0.8, 0.9, 1e-3, True
]
variable_np = np.array([1.0, 2.0], dtype=np.float32)
gradients_np = np.array([0.1, 0.2], dtype=np.float32)
mean_gradients_np = np.array([0.0, 0.0], dtype=np.float32)
mean_square_np = np.array([epsilon, epsilon], dtype=np.float32)
moment_np = np.array([0.0, 0.0], dtype=np.float32)
variable = Tensor(variable_np)
gradients = Tensor(gradients_np)
mean_gradients = Tensor(mean_gradients_np)
mean_square = Tensor(mean_square_np)
moment = Tensor(moment_np)
variable_ms = Parameter(initializer(variable, variable.shape), name='var')
gradients_ms = Parameter(initializer(gradients, gradients.shape),
name='grad')
mean_gradients_ms = Parameter(initializer(mean_gradients,
mean_gradients.shape),
name='mg')
mean_square_ms = Parameter(initializer(mean_square, mean_square.shape),
name='msr')
moment_ms = Parameter(initializer(moment, moment.shape), name='mom')
if centered:
variable_np, gradients_np, mean_gradients_np, mean_square_np, moment_np = \
rmspropcented_numpy(variable_np, gradients_np, mean_gradients_np, mean_square_np, moment_np,
learning_rate, decay, momentum, epsilon)
net = NetCenteredRMSProp(learning_rate, decay, momentum, epsilon,
variable_ms, gradients_ms, mean_gradients_ms,
mean_square_ms, moment_ms)
_ = net()
else:
variable_np, gradients_np, mean_square_np, moment_np = \
rmsprop_numpy(variable_np, gradients_np, mean_square_np, moment_np,
learning_rate, decay, momentum, epsilon)
net = NetRMSProp(learning_rate, decay, momentum, epsilon, variable_ms,
gradients_ms, mean_gradients_ms, mean_square_ms,
moment_ms)
_ = net()
error = np.ones(shape=variable_np.shape) * 10e-6
diff = variable_ms.asnumpy() - variable_np
assert np.all(diff < error)
error = np.ones(shape=gradients_np.shape) * 10e-6
diff = gradients_ms.asnumpy() - gradients_np
assert np.all(diff < error)
error = np.ones(shape=mean_gradients_np.shape) * 10e-6
diff = mean_gradients_ms.asnumpy() - mean_gradients_np
assert np.all(diff < error)
error = np.ones(shape=mean_square_np.shape) * 10e-6
diff = mean_square_ms.asnumpy() - mean_square_np
assert np.all(diff < error)
error = np.ones(shape=moment_np.shape) * 10e-6
diff = moment_ms.asnumpy() - moment_np
assert np.all(diff < error)
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()
self.update = False
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 = Parameter(Tensor(
np.asarray(self.md_info.h_charge, dtype=np.float32),
mstype.float32),
requires_grad=False)
self.old_crd = Parameter(Tensor(
np.zeros([self.atom_numbers, 3], dtype=np.float32),
mstype.float32),
requires_grad=False)
self.last_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 Main_Print(self, *args):
"""compute the temperature"""
steps, 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 = list(args)
if steps == 0:
print(
"_steps_ _TEMP_ _TOT_POT_ENE_ _BOND_ENE_ "
"_ANGLE_ENE_ _DIHEDRAL_ENE_ _14LJ_ENE_ _14CF_ENE_ _LJ_ENE_ _CF_PME_ENE_"
)
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 self.bond.bond_numbers > 0:
sigma_of_bond_ene = sigma_of_bond_ene.asnumpy()
print("{:>10.3f}".format(float(sigma_of_bond_ene)), end=" ")
if self.angle.angle_numbers > 0:
sigma_of_angle_ene = sigma_of_angle_ene.asnumpy()
print("{:>11.3f}".format(float(sigma_of_angle_ene)), end=" ")
if self.dihedral.dihedral_numbers > 0:
sigma_of_dihedral_ene = sigma_of_dihedral_ene.asnumpy()
print("{:>14.3f}".format(float(sigma_of_dihedral_ene)), end=" ")
if self.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 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)))
if self.datfile is not None:
self.datfile.write(self.crd.asnumpy())
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"
)
if self.control.mdcrd:
self.datfile = open(self.control.mdcrd, 'wb')
def Main_Destroy(self):
"""main destroy"""
if self.file is not None:
self.file.close()
print("Save .out file successfully!")
if self.datfile is not None:
self.datfile.close()
print("Save .dat file successfully!")
def construct(self, step, print_step):
'''construct'''
self.last_crd = self.crd
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
