Fix some bugs in tests.

Author: ErjieWu

deepks/io/input/dispatcher.py

@@ -31,19 +31,6 @@ def dispatch_command(config):
         scf_soft = config.get('scf_soft', 'pyscf')
         backend = get_scf_backend(scf_soft)
         backend.collect_stats(**config)
-    elif task_type == 'scf_task':
-        # Low-level per-task SCF runner (called by BatchTask in iterate workflow)
-        from deepks.physics.backends.pyscf.run import main as scf_run_main
-        # Strip orchestration-only keys that pyscf.run.main() doesn't understand
-        scf_config = {k: v for k, v in config.items()
-                      if k not in ('type', 'scf_soft')}
-        scf_run_main(**scf_config)
-    elif task_type == 'train_task':
-        # Low-level training runner (called by BatchTask in iterate workflow)
-        from deepks.ml.train.train import main as train_main
-        train_config = {k: v for k, v in config.items()
-                        if k not in ('type',)}
-        train_main(**train_config)
     elif task_type == 'iterate':
         # Use new iterate workflow
         from deepks.workflows.iterate import run_iterate_workflow

deepks/workflows/iterate/template.py

@@ -87,7 +87,7 @@ def make_scf_task(*, workdir=".",
 
     # --- build per-task override dict ---
     TASK_YAML = "_scf_task.yaml"
-    overrides = {"type": "scf_task"}
+    overrides = {"type": "scf", "scf_soft": "pyscf"}
     if sys_str is not None:
         overrides["systems"] = sys_str
     if model_file:
@@ -275,21 +275,19 @@ def make_train_task(*, workdir=".",
                     python="python", **task_args):
     """Create a training task as a BatchTask running via the unified 'deepks' CLI.
 
-    All training parameters are folded into a per-task YAML (_train_task.yaml)
-    with type: train_task.  The YAML is written to disk during preprocess()
-    (not via a shell one-liner), so the command is simply:
-        deepks _train_task.yaml
+    The base train_input.yaml (from share) is read at construction time, merged
+    with runtime overrides (train_paths, test_paths, type, restart, etc.), and
+    written back to train_input.yaml in the workdir during preprocess().  No
+    separate _train_task.yaml is created.
     """
-    link_share = task_args.pop("link_share_files", [])
+    from deepks.io.utils import load_yaml, deep_update
+
     link_prev  = task_args.pop("link_prev_files",  [])
+    link_share = task_args.pop("link_share_files", [])
     forward_files  = task_args.pop("forward_files",  [])
     backward_files = task_args.pop("backward_files", [])
 
-    # --- file links ---
-    if arg_file and source_arg is not None:
-        link_share.append((source_arg, arg_file))
-    if arg_file:
-        forward_files.append(arg_file)
+    # --- file links (everything except train_input.yaml, which we write ourselves) ---
     if restart_model and source_model is not None:
         link_prev.append((source_model, restart_model))
         forward_files.append(restart_model)
@@ -305,24 +303,27 @@ def make_train_task(*, workdir=".",
     if save_model:
         backward_files.append(save_model)
 
-    # --- build per-task YAML overrides ---
-    TASK_YAML = "_train_task.yaml"
-    overrides = {"type": "train_task"}
+    # --- read base train_input.yaml from source_arg (share folder) if available ---
+    base_config = {}
+    if source_arg is not None and os.path.exists(source_arg):
+        base_config = load_yaml(source_arg) or {}
+
+    # --- build runtime overrides and merge into base ---
+    overrides = {"type": "train"}
     if data_train:
-        overrides["train_paths"] = os.path.join(data_train, "*")
+        overrides["systems_train"] = os.path.join(data_train, "*")
     if data_test:
-        overrides["test_paths"] = os.path.join(data_test, "*")
+        overrides["systems_test"] = os.path.join(data_test, "*")
     if restart_model:
         overrides["restart"] = restart_model
     if proj_basis:
         overrides["proj_basis"] = proj_basis
     if save_model:
-        overrides.setdefault("train_args", {})["ckpt_file"] = save_model
+        overrides["ckpt_file"] = save_model
 
-    # Write the YAML at construction time; BatchTask.preprocess() writes it to
-    # disk before the shell command runs — no fragile python -c one-liner needed.
-    task_yaml_content = dump_yaml_str(overrides)
-    command = f"{SCF_CMD} {TASK_YAML}"
+    merged = deep_update(dict(base_config), overrides)
+    task_yaml_content = dump_yaml_str(merged)
+    command = f"{SCF_CMD} {arg_file}"
 
     return BatchTask(
         command,
@@ -336,7 +337,7 @@ def make_train_task(*, workdir=".",
         link_prev_files=link_prev,
         forward_files=forward_files,
         backward_files=backward_files,
-        write_files={TASK_YAML: task_yaml_content},
+        write_files={arg_file: task_yaml_content},
         **task_args
     )
 

deepks/workflows/scf/execute.py

@@ -262,3 +262,24 @@ def build_batch_task(sys_paths, sys_names, abacus_path, run_cmd,
         forward_files=forward_files,
         backward_files=backward_files
     )
+
+
+def execute_scf_tasks_pyscf(prepare_task, config):
+    """Execute SCF via PySCF by calling pyscf.run.main() directly."""
+    from deepks.physics.backends.pyscf.run import main as pyscf_main
+    from deepks.orchestration.workflow.task import PythonTask
+    from deepks.orchestration.workflow.workflow import Sequence
+
+    call_kwargs = {
+        k: v for k, v in config.items()
+        if k not in ('type', 'scf_soft')
+    }
+    run_task = PythonTask(
+        pyscf_main,
+        call_kwargs=call_kwargs,
+        outlog=config.get('outlog', 'log.scf'),
+        errlog='err',
+        workdir='.'
+    )
+    workflow = Sequence([prepare_task, run_task], workdir='.')
+    workflow.run()

deepks/workflows/scf/prepare.py

@@ -159,10 +159,7 @@ def prepare_scf_tasks(config):
     if scf_soft.lower() == 'abacus':
         return prepare_scf_tasks_abacus(config)
     elif scf_soft.lower() == 'pyscf':
-        # PySCF preparation (to be implemented or kept as is)
-        raise NotImplementedError(
-            "PySCF workflow not yet implemented in new architecture"
-        )
+        return prepare_scf_tasks_pyscf(config)
     else:
         raise ValueError(f"Unknown SCF backend: {scf_soft}")
 
@@ -218,3 +215,9 @@ def prepare_scf_tasks_abacus(config):
     )
 
     return task
+
+
+def prepare_scf_tasks_pyscf(config):
+    """PySCF SCF - no separate preparation needed; return a BlankTask."""
+    from deepks.orchestration.workflow.task import BlankTask
+    return BlankTask(workdir='.')

deepks/workflows/train/train.py

@@ -69,7 +69,10 @@ def train_model(train_reader: GroupReader,
 
     # Train model
     print('# starting training')
-    train_function(model, train_reader, test_reader=test_reader, **train_args)
+    import inspect
+    valid_train_keys = set(inspect.signature(train_function).parameters.keys())
+    filtered_train_args = {k: v for k, v in train_args.items() if k in valid_train_keys}
+    train_function(model, train_reader, test_reader=test_reader, **filtered_train_args)
 
     # Collect training statistics
     train_stats = {

deepks/workflows/train/workflow.py

@@ -9,6 +9,8 @@
 from .evaluate import evaluate_model
 from .types import TrainResult
 from dataclasses import asdict
+import sys
+from contextlib import redirect_stdout
 
 
 def run_train_workflow(config):
@@ -47,9 +49,22 @@ def run_train_workflow(config):
     # Stage 2: Train - Train the model
     model, train_stats = train_model(train_data, test_data, model_config)
 
-    # Stage 3: Evaluate - Evaluate model performance
+    # Stage 3: Evaluate - Evaluate model performance and write log.test
     metrics = evaluate_model(model, test_data, config)
 
+    # Reproduce legacy log.test: run test() with stdout redirected
+    ckpt_file = config.get('ckpt_file', 'model.pth')
+    test_log = config.get('test_log', 'log.test')
+    if test_data is not None:
+        from deepks.ml.eval.test import test as run_test
+        # Reconstruct the header line that GroupReader prints at construction time
+        data_keys = list(dict.fromkeys(test_data.readers[0].sample_all().keys()))
+        header_line = f'# load {test_data.nsystems} systems with fields {data_keys}'
+        with open(test_log, 'w', 1) as f_test, redirect_stdout(f_test):
+            print(header_line)
+            print(ckpt_file)
+            run_test(model, test_data, dump_prefix=None)
+
     result = TrainResult(
         model_path=config.get('ckpt_file', 'model.pth'),
         metrics=metrics,

tests/fixtures/legacy_integral_full/05_iter/01_abacus_local/iter_input.yaml

@@ -90,10 +90,10 @@ scf_soft: abacus
 scf_abacus:
   #INPUT args
   ntype: 2
-  ecutwfc: 50
-  scf_thr: 1e-5
+  ecutwfc: 100
+  scf_thr: 1e-7
   scf_nmax: 50
-  dft_functional: "lda"
+  dft_functional: "pbe"
   gamma_only: 1
   #STRU args
   orb_files:
@@ -115,10 +115,10 @@ init_scf_abacus:
   proj_file: 
     - "../../data/jle.orb"
   ntype: 2
-  ecutwfc: 50
-  scf_thr: 1e-5
+  ecutwfc: 100
+  scf_thr: 1e-7
   scf_nmax: 50
-  dft_functional: "lda"
+  dft_functional: "pbe"
   gamma_only: 1
   lattice_constant: 1.8897261255
   coord_type: "Cartesian"

tests/fixtures/legacy_integral_full/05_iter/01_abacus_local/test_iterate.py

@@ -44,7 +44,7 @@ def _remove_generated() -> None:
 
 
 def _fake_make_scf_abacus(*args, **kwargs):
-    from deepks.orchestration.scheduler.task import PythonTask
+    from deepks.orchestration.workflow.task import PythonTask
     workdir = kwargs.get("workdir", "00.scf")
 
     def _run_mock_scf():
@@ -55,7 +55,7 @@ def _run_mock_scf():
 
 
 def _fake_make_train(*args, **kwargs):
-    from deepks.orchestration.scheduler.task import PythonTask
+    from deepks.orchestration.workflow.task import PythonTask
     workdir = kwargs.get("workdir", "01.train")
 
     def _run_mock_train():
@@ -66,13 +66,13 @@ def _run_mock_train():
 
 @pytest.fixture(autouse=True)
 def _prepare_runtime_tree(monkeypatch):
-    _remove_generated()
+    # _remove_generated()
     if not ABACUS_AVAILABLE:
         from deepks.workflows.iterate import scf_step, train_step
         monkeypatch.setattr(scf_step, "make_scf_abacus", _fake_make_scf_abacus)
         monkeypatch.setattr(train_step, "make_train", _fake_make_train)
     yield
-    _remove_generated()
+    # _remove_generated()
 
 
 def run_iter():