refactor(phase3): migrate model and scf implementations to core

Author: ErjieWu

deepks/core/ml/eval/__init__.py

@@ -1,3 +1,4 @@
 """Evaluation components for DeepKS core ML layer."""
 
 from .evaluator import *  # noqa: F401,F403
+from .test import *  # noqa: F401,F403

deepks/core/ml/eval/evaluator.py

@@ -8,7 +8,7 @@
 except ImportError as e:
     sys.path.append(os.path.dirname(os.path.realpath(__file__)) + "/../../")
 from deepks.io.transforms.linalg import generalized_eigh
-from deepks.model.utils import get_density_matrix, cal_phi_loss, cal_v_delta, get_occ_func, make_loss, get_gedm, cal_vdr, loss_hr
+from deepks.core.ml.utils import get_density_matrix, cal_phi_loss, cal_v_delta, get_occ_func, make_loss, get_gedm, cal_vdr, loss_hr
 
 class Evaluator:
     def __init__(self,

deepks/core/ml/eval/test.py

@@ -0,0 +1,93 @@
+import os
+import numpy as np
+import torch
+import torch.nn as nn
+try:
+    import deepks
+except ImportError as e:
+    import sys
+    sys.path.append(os.path.dirname(os.path.realpath(__file__)) + "/../../")
+from deepks.default import DEVICE
+from deepks.core.ml.models.corrnet import CorrNet
+from deepks.io.readers.group_reader import GroupReader
+from deepks.utils import load_yaml, load_dirs, check_list
+
+
+def test(model, g_reader, dump_prefix="test", group=False):
+    model.eval()
+    loss_fn=nn.MSELoss()
+    label_list = []
+    pred_list = []
+
+    for i in range(g_reader.nsystems):
+        sample = g_reader.sample_all(i)
+        nframes = sample["lb_e"].shape[0]
+        for k, v in sample.items():
+            if isinstance(v, list):
+                sample[k] = [vv.to(DEVICE, non_blocking=True) for vv in v]
+            elif not torch.is_complex(v):
+                sample[k] = v.to(DEVICE, non_blocking=True)
+            else:
+                if k == "phialpha":
+                    sample[k] = v.to("cpu", dtype=torch.complex128, non_blocking=True)
+                else:
+                    sample[k] = v.to(DEVICE, dtype=torch.complex128, non_blocking=True)
+        label, data = sample["lb_e"], sample["eig"]
+        pred = model(data)
+        error = torch.sqrt(loss_fn(pred, label))
+
+        error_np = error.item()
+        label_np = label.cpu().numpy().reshape(nframes, -1).sum(axis=1)
+        pred_np = pred.detach().cpu().numpy().reshape(nframes, -1).sum(axis=1)
+        error_l1 = np.mean(np.abs(label_np - pred_np))
+        label_list.append(label_np)
+        pred_list.append(pred_np)
+
+        if not group and dump_prefix is not None:
+            nd = max(len(str(g_reader.nsystems)), 2)
+            dump_res = np.stack([label_np, pred_np], axis=1)
+            header = f"{g_reader.path_list[i]}\nmean l1 error: {error_l1}\nmean l2 error: {error_np}\nreal_ene  pred_ene"
+            filename = f"{dump_prefix}.{i:0{nd}}.out"
+            np.savetxt(filename, dump_res, header=header)
+            # print(f"system {i} finished")
+
+    all_label = np.concatenate(label_list, axis=0)
+    all_pred = np.concatenate(pred_list, axis=0)
+    all_err_l1 = np.mean(np.abs(all_label - all_pred))
+    all_err_l2 = np.sqrt(np.mean((all_label - all_pred) ** 2))
+    info = f"all systems mean l1 error: {all_err_l1}\nall systems mean l2 error: {all_err_l2}"
+    print(info)
+    if dump_prefix is not None and group:
+        np.savetxt(f"{dump_prefix}.out", np.stack([all_label, all_pred], axis=1), 
+                   header=info + "\nreal_ene  pred_ene")
+    return all_err_l1, all_err_l2
+
+
+def main(data_paths, model_file="model.pth", 
+         output_prefix='test', group=False,
+         e_name='l_e_delta', d_name=['dm_eig']):
+    data_paths = load_dirs(data_paths)
+    if len(d_name) == 1:
+        d_name = d_name[0]
+    g_reader = GroupReader(data_paths, e_name=e_name, d_name=d_name, 
+                           conv_filter=False, extra_label=True)
+    model_file = check_list(model_file)
+    for f in model_file:
+        print(f)
+        p = os.path.dirname(f)
+        model = CorrNet.load(f).double().to(DEVICE)
+        dump = os.path.join(p, output_prefix)
+        dir_name = os.path.dirname(dump)
+        if dir_name:
+            os.makedirs(dir_name, exist_ok=True)
+        if model.elem_table is not None:
+            elist, econst = model.elem_table
+            g_reader.collect_elems(elist)
+            g_reader.subtract_elem_const(econst)
+        test(model, g_reader, dump_prefix=dump, group=group)
+        g_reader.revert_elem_const()
+
+
+if __name__ == "__main__":
+    from deepks.main import test_cli as cli
+    cli()

deepks/core/ml/train/__init__.py

@@ -1 +1,3 @@
-"""Scaffold package for refactor architecture."""
+"""Training components for DeepKS core ML layer."""
+
+from .train import *  # noqa: F401,F403

deepks/core/ml/train/train.py

@@ -0,0 +1,249 @@
+import os
+import sys
+import numpy as np
+import torch
+import torch.optim as optim
+from time import time
+try:
+    import deepks
+except ImportError as e:
+    sys.path.append(os.path.dirname(os.path.realpath(__file__)) + "/../../")
+from deepks.default import DEVICE
+from deepks.core.ml.models.corrnet import CorrNet
+from deepks.io.readers.group_reader import GroupReader
+from deepks.utils import load_dirs, load_elem_table
+from deepks.core.ml.utils import preprocess, fit_elem_const, make_loss
+from deepks.core.ml.eval.evaluator import Evaluator, NatomLossList
+
+def train(model, g_reader, n_epoch=1000, test_reader=None, *,
+          energy_factor=1., force_factor=0., stress_factor=0., orbital_factor=0., v_delta_factor=0., v_delta_r_factor=0., phi_factor=0.,phi_occ=0, band_factor=0., band_occ=0, density_m_factor=0., density_m_occ=0, density_factor=0.,
+          energy_loss=None, force_loss=None, stress_loss=None, orbital_loss=None, v_delta_loss=None, v_delta_r_loss=None, phi_loss=None, band_loss=None, density_m_loss=None, grad_penalty=0.,
+          energy_per_atom=0, vd_divide_by_nlocal=False,
+          start_lr=0.001, decay_steps=100, decay_rate=0.96, stop_lr=None, decay_rate_iter=None,
+          weight_decay=0.,  fix_embedding=False,
+          display_epoch=100, display_detail_test=0, display_natom_loss=False, ckpt_file="model.pth",
+          graph_file=None, device=DEVICE):
+    
+    model = model.to(device)
+    model.eval()
+    print("# working on device:", device)
+    if test_reader is None:
+        test_reader = g_reader
+    # fix parameters if needed
+    if fix_embedding and model.embedder is not None:
+        model.embedder.requires_grad_(False)
+    # set up optimizer and lr scheduler
+    if decay_rate_iter is not None:
+        # decay_rate of start_lr for iterations, often start from iter.00
+        current_dir=os.getcwd()
+        current_iter=current_dir.split("/")[-2].split(".")[-1]
+        if current_iter != "init": # no need to change
+            current_iter=int(current_iter)
+            start_lr=start_lr*(decay_rate_iter**current_iter)
+            print(f"# resetting start_lr to {start_lr:.2e} because of decay_rate_iter")
+    optimizer = optim.Adam(model.parameters(), lr=start_lr, weight_decay=weight_decay)
+    if stop_lr is not None:
+        decay_rate = (stop_lr / start_lr) ** (1 / (n_epoch // decay_steps))
+        print(f"# resetting decay_rate: {decay_rate:.4f} "
+              + f"to satisfy stop_lr: {stop_lr:.2e}")
+    scheduler = optim.lr_scheduler.StepLR(optimizer, decay_steps, decay_rate)
+    # make evaluators for training
+    evaluator = Evaluator(energy_factor=energy_factor, force_factor=force_factor, 
+                          stress_factor=stress_factor, orbital_factor=orbital_factor,
+                          v_delta_factor=v_delta_factor, v_delta_r_factor=v_delta_r_factor,
+                          phi_factor=phi_factor, phi_occ=phi_occ,
+                          band_factor=band_factor, band_occ=band_occ,
+                          density_m_factor=density_m_factor, density_m_occ=density_m_occ,
+                          energy_lossfn=energy_loss, force_lossfn=force_loss,
+                          stress_lossfn=stress_loss, orbital_lossfn=orbital_loss,
+                          v_delta_lossfn=v_delta_loss, v_delta_r_lossfn=v_delta_r_loss, phi_lossfn=phi_loss,
+                          band_lossfn=band_loss, density_m_lossfn=density_m_loss,
+                          density_factor=density_factor, grad_penalty=grad_penalty, 
+                          energy_per_atom=energy_per_atom, vd_divide_by_nlocal=vd_divide_by_nlocal)
+    if not display_detail_test:
+        # make test evaluator that only returns l2loss of energy
+        test_eval = Evaluator(energy_factor=1., energy_lossfn=make_loss(), # default l2 loss 
+                            force_factor=0., density_factor=0., grad_penalty=0.,energy_per_atom=energy_per_atom)
+    else:
+        # make test evaluator that returns loss of every concerned items, but all with factor==1
+        to_one = lambda x: 0. if x == 0. else 1.
+        test_eval = Evaluator(energy_factor=to_one(energy_factor), force_factor=to_one(force_factor), 
+                            stress_factor=to_one(stress_factor), orbital_factor=to_one(orbital_factor),
+                            v_delta_factor=to_one(v_delta_factor), v_delta_r_factor=to_one(v_delta_r_factor),
+                            phi_factor=to_one(phi_factor), phi_occ=phi_occ,
+                            band_factor=to_one(band_factor), band_occ=band_occ,
+                            density_m_factor=to_one(density_m_factor), density_m_occ=density_m_occ,
+                            energy_lossfn=energy_loss, force_lossfn=force_loss,
+                            stress_lossfn=stress_loss, orbital_lossfn=orbital_loss,
+                            v_delta_lossfn=v_delta_loss, v_delta_r_lossfn=v_delta_r_loss, phi_lossfn=phi_loss,
+                            band_lossfn=band_loss, density_m_lossfn=density_m_loss,
+                            density_factor=to_one(density_factor), grad_penalty=grad_penalty,
+                            energy_per_atom=energy_per_atom, vd_divide_by_nlocal=vd_divide_by_nlocal)
+
+    print("# epoch      trn_err   tst_err        lr  trn_time  tst_time",end='')
+    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)
+    align_len=20
+    evaluator.print_head("trn_loss",data_keys,align_len)
+    if display_detail_test:
+        test_eval.print_head("tst_loss",data_keys,align_len)
+    # print("")
+
+    tic = time()
+    trn_natom_loss_list=NatomLossList()
+    tst_natom_loss_list=NatomLossList()
+    for batch in g_reader.sample_all_batch():
+        loss=evaluator(model,batch)
+        natom=batch["eig"].shape[1]
+        trn_natom_loss_list.add_loss(natom,loss)
+    trn_loss=trn_natom_loss_list.avg_loss()
+    for batch in test_reader.sample_all_batch():
+        loss=test_eval(model,batch)
+        natom=batch["eig"].shape[1]
+        tst_natom_loss_list.add_loss(natom,loss)
+    tst_loss=tst_natom_loss_list.avg_loss()    
+    # trn_loss = np.mean([[loss_term.item() for loss_term in evaluator(model, batch)]
+    #                 for batch in g_reader.sample_all_batch()],axis=0)
+    # tst_loss = np.mean([[loss_term.item() for loss_term in test_eval(model, batch)]
+    #                 for batch in test_reader.sample_all_batch()],axis=0)
+    tst_time = time() - tic
+    if display_natom_loss:
+        for natom in trn_natom_loss_list.natoms():
+            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,align_len)
+            else:
+                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:>{align_len}.4e}",end='')
+    if display_detail_test:
+        for loss_term in tst_loss[:-1]:
+            print(f"{loss_term:>{align_len}.4e}",end='')
+    if display_natom_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):
+        tic = time()
+        # loss_list = []
+        trn_natom_loss_list.clear_loss()
+        tst_natom_loss_list.clear_loss()
+        for sample in g_reader:
+            model.train()
+            optimizer.zero_grad()
+            loss = evaluator(model, sample)
+            loss[-1].backward()
+            # print("vdr_pred grad:",evaluator.vdr_pred.grad)
+            # print("e_loss grad:",evaluator.e_loss.grad)
+            # print("vdr_loss grad:",evaluator.vdr_loss.grad)
+            # print("tot_loss grad:",evaluator.tot_loss.grad)
+            optimizer.step()
+            # loss_list.append([loss_term.item() for loss_term in loss])
+            natom=sample["eig"].shape[1]
+            trn_natom_loss_list.add_loss(natom,loss)
+        scheduler.step()
+
+        if epoch % display_epoch == 0:
+            model.eval()
+            # trn_loss = np.mean(loss_list,axis=0)
+            trn_loss=trn_natom_loss_list.avg_loss()
+            trn_time = time() - tic
+            tic = time()
+            # tst_loss = np.mean([[loss_term.item() for loss_term in test_eval(model, batch)]
+            #                 for batch in test_reader.sample_all_batch()],axis=0)
+            for batch in test_reader.sample_all_batch():
+                loss=test_eval(model,batch)
+                natom=batch["eig"].shape[1]
+                tst_natom_loss_list.add_loss(natom,loss)
+            tst_loss=tst_natom_loss_list.avg_loss()  
+            tst_time = time() - tic
+            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:>{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:>{align_len}.4e}",end='')
+            if display_natom_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)
+
+    if ckpt_file:
+        model.save(ckpt_file)
+    if graph_file:
+        model.compile_save(graph_file)
+    
+
+def main(train_paths, test_paths=None,
+         restart=None, ckpt_file=None, 
+         model_args=None, data_args=None, 
+         preprocess_args=None, train_args=None, 
+         proj_basis=None, fit_elem=False, 
+         seed=None, device=None):
+   
+    if seed is None: 
+        seed = np.random.randint(0, 2**32)
+    print(f'# using seed: {seed}')
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+
+    if model_args is None: model_args = {}
+    if data_args is None: data_args = {}
+    if preprocess_args is None: preprocess_args = {}
+    if train_args is None: train_args = {}
+    if proj_basis is not None:
+        model_args["proj_basis"] = proj_basis
+    if ckpt_file is not None:
+        train_args["ckpt_file"] = ckpt_file
+    if device is not None:
+        train_args["device"] = device
+
+    train_paths = load_dirs(train_paths)
+    # print(f'# training with {len(train_paths)} system(s)')
+    g_reader = GroupReader(train_paths, **data_args)
+    if test_paths is not None:
+        test_paths = load_dirs(test_paths)
+        # print(f'# testing with {len(test_paths)} system(s)')
+        test_reader = GroupReader(test_paths, **data_args)
+    else:
+        print('# testing with training set')
+        test_reader = None
+
+    if restart is not None:
+        model = CorrNet.load(restart)
+        if model.elem_table is not None:
+            fit_elem_const(g_reader, test_reader, model.elem_table)
+    else:
+        input_dim = g_reader.ndesc
+        if model_args.get("input_dim", input_dim) != input_dim:
+            print(f"# `input_dim` in `model_args` does not match data",
+                  f"({input_dim}).", "Use the one in data.", file=sys.stderr)
+        model_args["input_dim"] = input_dim
+        if fit_elem:
+            elem_table = model_args.get("elem_table", None)
+            if isinstance(elem_table, str):
+                elem_table = load_elem_table(elem_table)
+            elem_table = fit_elem_const(g_reader, test_reader, elem_table)
+            model_args["elem_table"] = elem_table
+        model = CorrNet(**model_args).double()
+        
+    preprocess(model, g_reader, **preprocess_args)
+    # start=time()
+    train(model, g_reader, test_reader=test_reader, **train_args)
+    # end=time()
+    # print("all train time:",end-start)
+
+
+if __name__ == "__main__":
+    from deepks.main import train_cli as cli
+    cli()