Adjust the output way to ensure that log.train is well aligned

Author: xuan112358

deepks/model/evaluator.py

@@ -220,35 +220,34 @@ def __call__(self, model, sample):
         loss.append(tot_loss)
         return loss
 
-    def print_head(self,name,data_keys):
-        len=20
-        info=f"{name}_energy".rjust(len)
+    def print_head(self,name,data_keys,align_len=20):
+        info=f"{name}_energy".rjust(align_len)
         if self.g_penalty > 0 and "eg0" in data_keys:
-            info+=f"{name}_grad".rjust(len)
+            info+=f"{name}_grad".rjust(align_len)
         # optional force calculation
         if self.f_factor > 0 and "lb_f" in data_keys:
-            info+=f"{name}_force".rjust(len)
+            info+=f"{name}_force".rjust(align_len)
         # optional stress calculation
         if self.s_factor > 0 and "lb_s" in data_keys:
-            info+=f"{name}_stress".rjust(len)
+            info+=f"{name}_stress".rjust(align_len)
         # optional orbital(bandgap) calculation
         if self.o_factor > 0 and "lb_o" in data_keys:
-            info+=f"{name}_bandgap".rjust(len)
+            info+=f"{name}_bandgap".rjust(align_len)
         # optional v_delta calculation
         if self.vd_factor > 0 and "lb_vd" in data_keys:
-            info+=f"{name}_v_delta".rjust(len)
+            info+=f"{name}_v_delta".rjust(align_len)
         # optional phi calculation
         if self.phi_factor > 0 and "lb_phi" in data_keys:
-            info+=f"{name}_phi".rjust(len)
+            info+=f"{name}_phi".rjust(align_len)
         # optional band energy calculation
         if self.band_factor > 0 and "lb_band" in data_keys:
-            info+=f"{name}_band".rjust(len)
+            info+=f"{name}_band".rjust(align_len)
         # optional density matrix calculation
         if self.density_m_factor > 0 and "lb_phi" in data_keys:
-            info+=f"{name}_dm".rjust(len)             
+            info+=f"{name}_dm".rjust(align_len)             
         # density loss with fix head grad
         if self.d_factor > 0 and "gldv" in data_keys:
-            info+=f"{name}_density".rjust(len)
+            info+=f"{name}_density".rjust(align_len)
         print(info,end='')
 
 class NatomLossList:
@@ -275,11 +274,11 @@ def avg_atom_loss(self):
         # avg upon data
         return {natom:np.mean(losses,axis=0) for (natom,losses) in self.natom_loss_list.items()}
 
-    def print_avg_atom_loss(self):
+    def print_avg_atom_loss(self,align_len=20):
         avg_atom_loss = sorted(self.avg_atom_loss().items(), key=lambda x: x[0])
         for (atom,aal) in avg_atom_loss:
             for avg_atom_loss_term in aal[:-1]:
-                print(f"{avg_atom_loss_term:>18.4e}",end='')
+                print(f"{avg_atom_loss_term:>{align_len}.4e}",end='')
 
     def avg_loss(self):
         # avg upon data and natom

deepks/model/train.py

@@ -84,9 +84,10 @@ def train(model, g_reader, n_epoch=1000, test_reader=None, *,
     data_keys = g_reader.readers[0].sample_all().keys()
     # L_inv_in=1 if "L_inv" in data_keys else 0
     # print("if L_inv in sample:",L_inv_in)
-    evaluator.print_head("trn_loss",data_keys)
+    align_len=20
+    evaluator.print_head("trn_loss",data_keys,align_len)
     if display_detail_test:
-        test_eval.print_head("tst_loss",data_keys)
+        test_eval.print_head("tst_loss",data_keys,align_len)
     # print("")
 
     tic = time()
@@ -109,24 +110,24 @@ def train(model, g_reader, n_epoch=1000, test_reader=None, *,
     tst_time = time() - tic
     if display_natom_loss:
         for natom in trn_natom_loss_list.natoms():
-            evaluator.print_head(str(natom)+"_trn",data_keys)        
+            evaluator.print_head(str(natom)+"_trn",data_keys,align_len)        
         for natom in tst_natom_loss_list.natoms():
             if display_detail_test:
-                test_eval.print_head(str(natom)+"_tst",data_keys)
+                test_eval.print_head(str(natom)+"_tst",data_keys,align_len)
             else:
-                test_eval.print_head(str(natom)+"_tst",[])#just energy
+                test_eval.print_head(str(natom)+"_tst",[],align_len)#just energy
     print("")
 
     print(f"  {0:<8d}  {np.sqrt(np.abs(trn_loss[-1])):>.2e}  {np.sqrt(np.abs(tst_loss[-1])):>.2e}"
           f"  {start_lr:>.2e}  {0:>8.2f}  {tst_time:>8.2f}",end='')
     for loss_term in trn_loss[:-1]:
-        print(f"{loss_term:>18.4e}",end='')
+        print(f"{loss_term:>{align_len}.4e}",end='')
     if display_detail_test:
         for loss_term in tst_loss[:-1]:
-            print(f"{loss_term:>18.4e}",end='')
+            print(f"{loss_term:>{align_len}.4e}",end='')
     if display_natom_loss:
-        trn_natom_loss_list.print_avg_atom_loss()     
-        tst_natom_loss_list.print_avg_atom_loss()     
+        trn_natom_loss_list.print_avg_atom_loss(align_len)     
+        tst_natom_loss_list.print_avg_atom_loss(align_len)     
     print('')
 
     for epoch in range(1, n_epoch+1):
@@ -162,13 +163,13 @@ def train(model, g_reader, n_epoch=1000, test_reader=None, *,
             print(f"  {epoch:<8d}  {np.sqrt(np.abs(trn_loss[-1])):>.2e}  {np.sqrt(np.abs(tst_loss[-1])):>.2e}"
                   f"  {scheduler.get_last_lr()[0]:>.2e}  {trn_time:>8.2f}  {tst_time:8.2f}",end='')
             for loss_term in trn_loss[:-1]:
-                print(f"{loss_term:>18.4e}",end='')
+                print(f"{loss_term:>{align_len}.4e}",end='')
             if display_detail_test and epoch%(display_detail_test*display_epoch) == 0:
                 for loss_term in tst_loss[:-1]:
-                    print(f"{loss_term:>18.4e}",end='')
+                    print(f"{loss_term:>{align_len}.4e}",end='')
             if display_natom_loss:
-                trn_natom_loss_list.print_avg_atom_loss()
-                tst_natom_loss_list.print_avg_atom_loss()                 
+                trn_natom_loss_list.print_avg_atom_loss(align_len)
+                tst_natom_loss_list.print_avg_atom_loss(align_len)                 
             print('')
             if ckpt_file:
                 model.save(ckpt_file)