1. In opt_orb_pytorch_dpsi, split info to info_stru, info_element, info_opt

Author: PeizeLin

tools/opt_orb_pytorch_dpsi/IO/change_info.py

@@ -0,0 +1,31 @@
+import addict
+
+def change_info(info_old, weight_old):
+	info_stru = [None] * info_old.Nst
+	for ist in range(len(info_stru)):
+		info_stru[ist] = addict.Dict()
+	for ist,Na in enumerate(info_old.Na):
+		info_stru[ist].Na = Na
+	for ist,weight in enumerate(weight_old):
+		info_stru[ist].weight = weight
+		info_stru[ist].Nb = weight.shape[0]
+
+	info_element = addict.Dict()
+	for it,Nu in info_old.Nu.items():
+		info_element[it].Nu = Nu
+		info_element[it].Nl = len(Nu)
+	for it,Rcut in info_old.Rcut.items():
+		info_element[it].Rcut = Rcut
+	for it,dr in info_old.dr.items():
+		info_element[it].dr = dr
+	for it,Ecut in info_old.Ecut.items():
+		info_element[it].Ecut = Ecut
+	for it,Ne in info_old.Ne.items():
+		info_element[it].Ne = Ne
+
+	info_opt = addict.Dict()
+	info_opt.lr = info_old.lr
+	info_opt.cal_T = info_old.cal_T
+	info_opt.cal_smooth = info_old.cal_smooth
+
+	return info_stru, info_element, info_opt

tools/opt_orb_pytorch_dpsi/IO/func_C.py

@@ -2,20 +2,20 @@
 import torch
 import numpy as np
 
-def random_C_init(info):
+def random_C_init(info_element):
 	""" C[it][il][ie,iu]	<jY|\phi> """
 	C = dict()
-	for it in info.Nt_all:
-		C[it] = ND_list(info.Nl[it])
-		for il in range(info.Nl[it]):
-			C[it][il] = torch.tensor(np.random.uniform(-1,1, (info.Ne[it], info.Nu[it][il])), dtype=torch.float64, requires_grad=True)
+	for it in info_element.keys():
+		C[it] = ND_list(info_element[it].Nl)
+		for il in range(info_element[it].Nl):
+			C[it][il] = torch.tensor(np.random.uniform(-1,1, (info_element[it].Ne, info_element[it].Nu[il])), dtype=torch.float64, requires_grad=True)
 	return C
 
 
 
-def read_C_init(file_name,info):
+def read_C_init(file_name,info_element):
 	""" C[it][il][ie,iu]	<jY|\phi> """
-	C = random_C_init(info)
+	C = random_C_init(info_element)
 
 	with open(file_name,"r") as file:
 
@@ -29,11 +29,12 @@ def read_C_init(file_name,info):
 		while True:
 			line = file.readline().strip()
 			if line.startswith("Type"):
-				it,il,iu = list(map(int,file.readline().split()));	
-				it=info.Nt_all[it-1];	iu-=1
+				it,il,iu = file.readline().split();	
+				il = int(il)
+				iu = int(iu)-1
 				C_read_index.add((it,il,iu))
 				line = file.readline().split()
-				for ie in range(info.Ne[it]):
+				for ie in range(info_element[it].Ne):
 					if not line:	line = file.readline().split()
 					C[it][il].data[ie,iu] = float(line.pop(0))
 			elif line.startswith("</Coefficient>"):
@@ -44,17 +45,17 @@ def read_C_init(file_name,info):
 
 
 
-def copy_C(C,info):
+def copy_C(C,info_element):
 	C_copy = dict()
-	for it in info.Nt_all:
-		C_copy[it] = ND_list(info.Nl[it])
-		for il in range(info.Nl[it]):
+	for it in info_element.keys():
+		C_copy[it] = ND_list(info_element[it].Nl)
+		for il in range(info_element[it].Nl):
 			C_copy[it][il] = C[it][il].clone()
 	return C_copy
 
 
 
-def write_C(file_name,info,C,Spillage):
+def write_C(file_name,C,Spillage):
 	with open(file_name,"w") as file:
 		print("<Coefficient>", file=file)
 		#print("\tTotal number of radial orbitals.", file=file)
@@ -70,7 +71,7 @@ def write_C(file_name,info,C,Spillage):
 			for il,C_tl in enumerate(C_t):
 				for iu in range(C_tl.size()[1]):
 					print("\tType\tL\tZeta-Orbital", file=file)
-					print(f"\t  {info.Nt_all.index(it)+1} \t{il}\t    {iu+1}", file=file)
+					print(f"\t  {it} \t{il}\t    {iu+1}", file=file)
 					for ie in range(C_tl.size()[0]):
 						print("\t", '%18.14f'%C_tl[ie,iu].item(), file=file)
 		print("</Coefficient>", file=file)

tools/opt_orb_pytorch_dpsi/IO/print_orbital.py

@@ -23,12 +23,12 @@
                     ## 'Fl': 114, 'Uup': 115, 'Lv': 116, 'Uus': 117, 'Uuo': 118
                } 
 
-def print_orbital(orb,info):
+def print_orbital(orb,info_element):
 	""" orb[it][il][iu][r] """
 	for it,orb_t in orb.items():
 		#with open("orb_{0}.dat".format(it),"w") as file:
 		with open("ORBITAL_{0}U.dat".format( periodtable[it] ),"w") as file:
-			print_orbital_head(file,info,it)
+			print_orbital_head(file,info_element,it)
 			for il,orb_tl in enumerate(orb_t):
 				for iu,orb_tlu in enumerate(orb_tl):
 					print("""                Type                   L                   N""",file=file)
@@ -52,19 +52,19 @@ def plot_orbital(orb,Rcut,dr):
 				print(file=file)
 
 
-def print_orbital_head(file,info,it):
+def print_orbital_head(file,info_element,it):
 	print( "---------------------------------------------------------------------------", file=file )
 	print( "Element                     {0}".format(it), file=file )
-	print( "Energy Cutoff(Ry)           {0}".format(info.Ecut[it]), file=file )
-	print( "Radius Cutoff(a.u.)         {0}".format(info.Rcut[it]), file=file )
-	print( "Lmax                        {0}".format(info.Nl[it]-1), file=file )
+	print( "Energy Cutoff(Ry)           {0}".format(info_element[it].Ecut), file=file )
+	print( "Radius Cutoff(a.u.)         {0}".format(info_element[it].Rcut), file=file )
+	print( "Lmax                        {0}".format(info_element[it].Nl-1), file=file )
 	l_name = ["S","P","D"]+list(map(chr,range(ord('F'),ord('Z')+1)))
-	for il,iu in enumerate(info.Nu[it]):
+	for il,iu in enumerate(info_element[it].Nu):
 		print( "Number of {0}orbital-->       {1}".format(l_name[il],iu), file=file )
 	print( "---------------------------------------------------------------------------", file=file )
 	print( "SUMMARY  END", file=file )
 	print( file=file )
-	print( "Mesh                        {0}".format(int(info.Rcut[it]/info.dr[it])+1), file=file )
-	print( "dr                          {0}".format(info.dr[it]), file=file )
+	print( "Mesh                        {0}".format(int(info_element[it].Rcut/info_element[it].dr)+1), file=file )
+	print( "dr                          {0}".format(info_element[it].dr), file=file )
 
 

tools/opt_orb_pytorch_dpsi/IO/read_QSV.py

@@ -71,6 +71,7 @@ def read_file(info,file_list,V_info):
 			print("read VI:",ist_true,ik)
 			vi = read_VI(info_true,V_info,ist_true,data)
 			VI.append( vi )
+	print()
 
 	info_all = copy.deepcopy(info)
 	info_all.Nst = sum(info_true.Nk,0)

tools/opt_orb_pytorch_dpsi/main.py

@@ -13,6 +13,8 @@
 import torch_optimizer 
 import IO.cal_weight
 import util
+import IO.change_info
+import pprint
 
 def main():
 	seed = int(1000*time.time())%(2**32)
@@ -24,51 +26,59 @@ def main():
 
 	weight = IO.cal_weight.cal_weight(weight_info, V_info["same_band"], file_list["origin"])
 
-	QI,SI,VI_origin,info = IO.read_QSV.read_file(info_true,file_list["origin"],V_info)
-	print(info, flush=True)
+	QI,SI,VI_origin,info_kst = IO.read_QSV.read_file(info_true,file_list["origin"],V_info)
 	if "linear" in file_list.keys():
 		QI_linear, SI_linear, VI_linear, info_linear = list(zip(*( IO.read_QSV.read_file(info_true,file,V_info) for file in file_list["linear"] )))
 
+	info_stru, info_element, info_opt = IO.change_info.change_info(info_kst,weight)
+
+	print(pprint.pformat(info_stru), end="\n"*2, flush=True)
+	print(pprint.pformat(info_element,width=40), end="\n"*2, flush=True)
+	print(pprint.pformat(info_opt,width=40), end="\n"*2, flush=True)
+
 	if C_init_info["init_from_file"]:
-		C, C_read_index = IO.func_C.read_C_init( C_init_info["C_init_file"], info )
+		C, C_read_index = IO.func_C.read_C_init( C_init_info["C_init_file"], info_element )
 	else:
-		C = IO.func_C.random_C_init(info)
-	E = orbital.set_E(info,info.Rcut)
-	orbital.normalize( orbital.generate_orbital(info,C,E,info.Rcut,info.dr), info.dr,C,flag_norm_C=True)
+		C = IO.func_C.random_C_init(info_element)
+	E = orbital.set_E(info_element)
+	orbital.normalize(
+		orbital.generate_orbital(info_element,C,E),
+		{it:info_element[it].dr for it in info_element},
+		C, flag_norm_C=True)
 
 	opt_orb = opt_orbital.Opt_Orbital()
 
-	#opt = torch.optim.Adam(sum( ([c.real,c.imag] for c in sum(C,[])), []), lr=info.lr, eps=1e-8)
-	#opt = torch.optim.Adam( sum(C.values(),[]), lr=info.lr, eps=1e-20, weight_decay=info.weight_decay)
-	#opt = radam.RAdam( sum(C.values(),[]), lr=info.lr, eps=1e-20 )
-	opt = torch_optimizer.SWATS( sum(C.values(),[]), lr=info.lr, eps=1e-20 )
+	#opt = torch.optim.Adam(sum( ([c.real,c.imag] for c in sum(C,[])), []), lr=info_opt.lr, eps=1e-8)
+	#opt = torch.optim.Adam( sum(C.values(),[]), lr=info_opt.lr, eps=1e-20, weight_decay=info_opt.weight_decay)
+	#opt = radam.RAdam( sum(C.values(),[]), lr=info_opt.lr, eps=1e-20 )
+	opt = torch_optimizer.SWATS( sum(C.values(),[]), lr=info_opt.lr, eps=1e-20 )
 
 
 	with open("Spillage.dat","w") as S_file:
 
 		print( "\nSee \"Spillage.dat\" for detail status: " , flush=True )
-		if info.cal_T:
+		if info_opt.cal_T:
 			print( '%5s'%"istep", "%20s"%"Spillage", "%20s"%"T.item()", "%20s"%"Loss", flush=True )
 		else:
 			print( '%5s'%"istep", "%20s"%"Spillage", flush=True )
 
 		loss_old = np.inf
-		for istep in range(200):
+		for istep in range(3):
 
 			Spillage = 0
-			for ist in range(info.Nst):
+			for ist in range(len(info_stru)):
 
-				Q = opt_orb.change_index_Q(opt_orb.cal_Q(QI[ist],C,info,ist),info,ist)
-				S = opt_orb.change_index_S(opt_orb.cal_S(SI[ist],C,info,ist),info,ist)
+				Q = opt_orb.change_index_Q(opt_orb.cal_Q(QI[ist],C,info_stru[ist],info_element),info_stru[ist])
+				S = opt_orb.change_index_S(opt_orb.cal_S(SI[ist],C,info_stru[ist],info_element),info_stru[ist],info_element)
 				coef = opt_orb.cal_coef(Q,S)
 				V = opt_orb.cal_V(coef,Q)
 				V_origin = opt_orb.cal_V_origin(V,V_info)
 
 				if "linear" in file_list.keys():
 					V_linear = [None] * len(file_list["linear"])
 					for i in range(len(file_list["linear"])):
-						Q_linear = opt_orb.change_index_Q(opt_orb.cal_Q(QI_linear[i][ist],C,info,ist),info,ist)
-						S_linear = opt_orb.change_index_S(opt_orb.cal_S(SI_linear[i][ist],C,info,ist),info,ist)
+						Q_linear = opt_orb.change_index_Q(opt_orb.cal_Q(QI_linear[i][ist],C,info_stru[ist],info_element),info_stru[ist])
+						S_linear = opt_orb.change_index_S(opt_orb.cal_S(SI_linear[i][ist],C,info_stru[ist],info_element),info_stru[ist],info_element)
 						V_linear[i] = opt_orb.cal_V_linear(coef,Q_linear,S_linear,V,V_info)
 
 				def cal_Spillage(V_delta):
@@ -83,14 +93,14 @@ def cal_delta(VI, V):
 					for i in range(len(file_list["linear"])):
 						Spillage += cal_Spillage(cal_delta(VI_linear[i],V_linear[i]))
 
-			if info.cal_T:
+			if info_opt.cal_T:
 				T = opt_orb.cal_T(C,E)
 				if not "TSrate" in vars():	TSrate = torch.abs(0.002*Spillage/T).data[0]
 				Loss = Spillage + TSrate*T
 			else:
 				Loss = Spillage
 
-			if info.cal_T:
+			if info_opt.cal_T:
 				print_content = [istep, Spillage.item(), T.item(), Loss.item()]
 			else:
 				print_content = [istep, Spillage.item()]
@@ -100,7 +110,7 @@ def cal_delta(VI, V):
 
 			if Loss.item() < loss_old:
 				loss_old = Loss.item()
-				C_old = IO.func_C.copy_C(C,info)
+				C_old = IO.func_C.copy_C(C,info_element)
 				flag_finish = 0
 			else:
 				flag_finish += 1
@@ -113,16 +123,27 @@ def cal_delta(VI, V):
 				for it,il,iu in C_read_index:
 					C[it][il].grad[:,iu] = 0
 			opt.step()
-	#		orbital.normalize( orbital.generate_orbital(info,C,E,info.Rcut,info.dr), info.dr,C,flag_norm_C=True)
-
-		orb = orbital.generate_orbital(info,C_old,E,info.Rcut,info.dr)
-		if info.cal_smooth:
-			orbital.smooth_orbital(orb,info.Rcut,info.dr,0.1)
-		orbital.orth(orb,info.dr)
-		IO.print_orbital.print_orbital(orb,info)
-		IO.print_orbital.plot_orbital(orb,info.Rcut,info.dr)
-
-		IO.func_C.write_C("ORBITAL_RESULTS.txt",info,C_old,Spillage)
+	#		orbital.normalize(
+	# 			orbital.generate_orbital(info_element,C,E),
+	# 			{it:info_element[it].dr for it in info_element},
+	# 			C, flag_norm_C=True)
+
+		orb = orbital.generate_orbital(info_element,C_old,E)
+		if info_opt.cal_smooth:
+			orbital.smooth_orbital(
+				orb,
+				{it:info_element[it].Rcut for it in info_element}, {it:info_element[it].dr for it in info_element},
+				0.1)
+		orbital.orth(
+			orb,
+			{it:info_element[it].dr for it in info_element})
+		IO.print_orbital.print_orbital(orb,info_element)
+		IO.print_orbital.plot_orbital(
+			orb,
+			{it:info_element[it].Rcut for it in info_element},
+			{it:info_element[it].dr for it in info_element})
+
+		IO.func_C.write_C("ORBITAL_RESULTS.txt",C_old,Spillage)
 
 		print("Time (PyTorch):     %s\n"%(time.time()-time_start), flush=True )