def __init__(self, train_dict, loss_dict):
self.train_dict = train_dict
mode = train_dict["nn_mode"]
if mode == 'ff':
output_dim = 1
else:
print('hf not implemented yet')
quit()
self.mbpw_obj, self.net_list = self.net_builder(train_dict)
self.mlvv = velocity_verlet3(self.mbpw_obj)
self.opt = optim.SGD(self.mbpw_obj.parameters(), lr=train_dict["lr"])
self.sch = optim.lr_scheduler.StepLR(self.opt, train_dict["sch_step"],
train_dict["sch_decay"])
#self.sch = DecayCosineAnnealingWarmRestarts(self.opt,train_dict["sch_step"],train_dict["sch_decay"])
self.opt_tau = optim.SGD(self.mbpw_obj.tau_parameters(),
train_dict["tau_lr"])
lj = lennard_jones2d()
self.loss_obj = loss(lj, loss_dict["eweight"],
loss_dict["polynomial_degree"])
def __init__(self, train_dict, loss_dict):
self.train_dict = train_dict
pwnet_input = 5 # for pw force function
mbnet_input = 25 # for mb force function
mode = train_dict["nn_mode"]
if mode == 'hf':
output_dim = 1
else:
output_dim = 2
pwnet1 = mydevice.load(
pw_net(pwnet_input, output_dim, train_dict["nnodes"]))
pwnet2 = mydevice.load(
pw_net(pwnet_input, output_dim, train_dict["nnodes"]))
mbnet1 = mydevice.load(
mb_net(mbnet_input, output_dim, train_dict["nnodes"]))
mbnet2 = mydevice.load(
mb_net(mbnet_input, output_dim, train_dict["nnodes"]))
self.net_list = [pwnet1, pwnet2, mbnet1, mbnet2]
print('pwnet1', pwnet1)
print('pwnet2', pwnet2)
print('mbnet1', mbnet1)
print('mbnet2', mbnet2)
if mode == 'hf':
pwforce = pw_hf(pwnet1, pwnet2, train_dict["force_clip"])
mbforce = mb_hf(mbnet1, mbnet2, train_dict["ngrids"],
train_dict["b"], train_dict["force_clip"])
else:
pwforce = pw_ff(pwnet1, pwnet2, train_dict["force_clip"])
mbforce = mb_ff(mbnet1, mbnet2, train_dict["ngrids"],
train_dict["b"], train_dict["force_clip"])
self.mbpwff = mbpw(mbforce, pwforce)
param = itertools.chain(pwnet1.parameters(), pwnet2.parameters(),
mbnet1.parameters(), mbnet2.parameters())
self.opt = optim.SGD(param, lr=train_dict["lr"])
#sch = optim.lr_scheduler.StepLR(self.opt,train_dict["sch_step"],train_dict["sch_decay"])
self.sch = DecayCosineAnnealingWarmRestarts(self.opt,
train_dict["sch_step"],
train_dict["sch_decay"])
self.opta = optim.SGD(self.mbpwff.parameters(), train_dict["alpha_lr"])
self.mlvv = velocity_verlet(self.mbpwff)
lj = lennard_jones2d()
self.loss_obj = loss(lj, loss_dict["eweight"],
loss_dict["polynomial_degree"])
self.ckpt = checkpoint(self.net_list, self.mbpwff, self.opt, self.opta,
self.sch)
def __init__(self, qp_list_init, qp_list_final, l_list):
self.q_list_init = qp_list_init[:, 0, :, :]
self.p_list_init = qp_list_init[:, 1, :, :]
self.q_list_final = qp_list_final[:, 0, :, :]
self.p_list_final = qp_list_final[:, 1, :, :]
self.l_list = l_list
self.nsamples, self.nparticles, _ = self.q_list_init.shape
self.potential_function = lennard_jones2d()