fix: select CALYPSO models by backend

Fixes #1906

Coding-Agent: Codex

Codex-Version: codex-cli 0.149.0

Model: gpt-5.6-sol

Reasoning-Effort: xhigh

Author: njzjz-bot

dpgen/generator/lib/calypso_check_outcar.py

@@ -1,5 +1,6 @@
 #!/usr/bin/env python3
 
+import glob
 import os
 
 import numpy as np
@@ -10,6 +11,14 @@
 """
 
 
+def find_model_path():
+    """Find the first backend-specific model forwarded by DP-GEN."""
+    models = sorted(glob.glob(os.path.join("..", "graph.*")))
+    if not models:
+        raise FileNotFoundError("No graph model was forwarded for CALYPSO recovery")
+    return models[0]
+
+
 def Get_Element_Num(elements):
     """Using the Atoms.symples to Know Element&Num."""
     element = []
@@ -86,7 +95,7 @@ def check():
     from ase.io import read
     from deepmd.calculator import DP
 
-    calc = DP(model="../graph.000.pb")  # init the model before iteration
+    calc = DP(model=find_model_path())  # initialize one model before iteration
 
     to_be_opti = read("POSCAR")
     to_be_opti.calc = calc

dpgen/generator/lib/calypso_run_opt.py

@@ -1,5 +1,6 @@
 #!/usr/bin/env python3
 
+import glob
 import os
 import time
 
@@ -15,6 +16,14 @@
 """
 
 
+def find_model_path():
+    """Find the first backend-specific model forwarded by DP-GEN."""
+    models = sorted(glob.glob(os.path.join("..", "graph.*")))
+    if not models:
+        raise FileNotFoundError("No graph model was forwarded for CALYPSO optimization")
+    return models[0]
+
+
 def Get_Element_Num(elements):
     """Using the Atoms.symples to Know Element&Num."""
     element = []
@@ -112,7 +121,7 @@ def read_stress_fmax():
 
 def run_opt(fmax, stress):
     """Using the ASE&DP to Optimize Configures."""
-    calc = DP(model="../graph.000.pb")  # init the model before iteration
+    calc = DP(model=find_model_path())  # initialize one model before iteration
     os.system("mv OUTCAR OUTCAR-last")
 
     print("Start to Optimize Structures by DP----------")

dpgen/generator/lib/run_calypso.py

@@ -29,6 +29,25 @@
 calypso_model_devi_name = "model_devi_results"
 
 
+def _get_calypso_models(path, jdata):
+    """Return CALYPSO model files for the configured training backend."""
+    backend = jdata.get("train_backend", "tensorflow")
+    suffixes = {
+        "tensorflow": ".pb",
+        "pytorch": ".pth",
+        "jax": ".savedmodel",
+    }
+    try:
+        suffix = suffixes[backend]
+    except KeyError as exc:
+        supported = ", ".join(sorted(suffixes))
+        raise ValueError(
+            f"CALYPSO does not support training backend {backend!r}; "
+            f"supported backends are: {supported}."
+        ) from exc
+    return sorted(glob.glob(os.path.join(path, f"graph*{suffix}")))
+
+
 def gen_structures(
     iter_index, jdata, mdata, caly_run_path, current_idx, length_of_caly_runopt_list
 ):
@@ -50,7 +69,7 @@ def gen_structures(
     calypso_path = mdata.get("model_devi_calypso_path")
     # calypso_input_path = jdata.get('calypso_input_path')
 
-    all_models = glob.glob(os.path.join(calypso_run_opt_path, "graph*pb"))
+    all_models = _get_calypso_models(calypso_run_opt_path, jdata)
     model_names = [os.path.basename(ii) for ii in all_models]
 
     deepmdkit_python = mdata.get("model_devi_deepmdkit_python")
@@ -492,7 +511,7 @@ def run_calypso_model_devi(iter_index, jdata, mdata):
         elif lines[-1].strip().strip("\n") == "3":
             # Model Devi
             _calypso_run_opt_path = os.path.abspath(caly_run_opt_list[0])
-            all_models = glob.glob(os.path.join(_calypso_run_opt_path, "graph*pb"))
+            all_models = _get_calypso_models(_calypso_run_opt_path, jdata)
             cwd = os.getcwd()
             os.chdir(calypso_model_devi_path)
             args = " ".join(

tests/generator/test_calypso.py

@@ -1,5 +1,6 @@
 import os
 import sys
+import tempfile
 import unittest
 from pathlib import Path
 
@@ -8,6 +9,8 @@
 sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
 __package__ = "generator"
 
+from dpgen.generator.lib.run_calypso import _get_calypso_models
+
 from .context import (
     _parse_calypso_dis_mtx,
     _parse_calypso_input,
@@ -174,6 +177,24 @@ def test_parse_calypso_input(self):
         )
         os.remove("input.dat")
 
+    def test_backend_specific_model_selection(self):
+        with tempfile.TemporaryDirectory() as tmpdir:
+            model_dir = Path(tmpdir)
+            for suffix in ("pb", "pth", "savedmodel"):
+                (model_dir / f"graph.000.{suffix}").touch()
+
+            cases = {
+                "tensorflow": ".pb",
+                "pytorch": ".pth",
+                "jax": ".savedmodel",
+            }
+            for backend, suffix in cases.items():
+                with self.subTest(backend=backend):
+                    models = _get_calypso_models(
+                        str(model_dir), {"train_backend": backend}
+                    )
+                    self.assertEqual(models, [str(model_dir / f"graph.000{suffix}")])
+
 
 if __name__ == "__main__":
     unittest.main(verbosity=2)