From 00f55c1909a75212000f86cd7dffe5a9b2637c22 Mon Sep 17 00:00:00 2001
From: susumu-fujii <susumu-fujii>
Date: Tue, 18 Aug 2026 21:56:38 +0900
Subject: [PATCH] Fix NequIP atom type mapping for deployed models

---
 .../latgas_abinitio_interface/nequip.py       | 22 ++++++--
 tests/test_nequip.py                          | 56 +++++++++++++++++++
 2 files changed, 73 insertions(+), 5 deletions(-)
 create mode 100644 tests/test_nequip.py

diff --git a/abics/applications/latgas_abinitio_interface/nequip.py b/abics/applications/latgas_abinitio_interface/nequip.py
index 3d597027..5d107ebf 100644
--- a/abics/applications/latgas_abinitio_interface/nequip.py
+++ b/abics/applications/latgas_abinitio_interface/nequip.py
@@ -114,11 +114,23 @@ def from_directory(self, base_input_dir):
                 Path to the directory including base input files.
             """
             self.base_input_dir = base_input_dir
-            self.model = torch.jit.load(os.path.join(base_input_dir, "deployed.pth"))
-            yaml_file = os.path.join(base_input_dir, "input.yaml")
-            yaml_dic = Config.from_file(yaml_file)
-            self.element_list = yaml_dic["chemical_symbols"]
-            self.r_max = yaml_dic["r_max"]
+
+            metadata = {
+                "r_max": "",
+                "type_names": "",
+            }
+
+            self.model = torch.jit.load(
+                os.path.join(base_input_dir, "deployed.pth"),
+                _extra_files=metadata,
+            )
+
+            for key, value in metadata.items():
+                if isinstance(value, bytes):
+                    metadata[key] = value.decode()
+
+            self.element_list = metadata["type_names"].split()
+            self.r_max = float(metadata["r_max"])
 
         def update_info_by_structure(self, structure):
             """
diff --git a/tests/test_nequip.py b/tests/test_nequip.py
new file mode 100644
index 00000000..460af3ad
--- /dev/null
+++ b/tests/test_nequip.py
@@ -0,0 +1,56 @@
+import unittest
+from unittest.mock import patch
+
+from abics.applications.latgas_abinitio_interface.nequip import NequipSolver
+
+
+class TestNequipSolverInput(unittest.TestCase):
+
+    def test_deployed_type_names_take_precedence_over_chemical_symbols(self):
+        """
+        The atom type ordering used by a deployed NequIP model must come from
+        the deployed model metadata, not from input.yaml chemical_symbols.
+        """
+        import tempfile
+        from pathlib import Path
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            base_input_dir = Path(tmpdir)
+
+            # input.yaml intentionally has a different ordering.
+            (base_input_dir / "input.yaml").write_text(
+                "chemical_symbols:\n"
+                "  - Cu\n"
+                "  - V\n"
+                "  - Sn\n"
+                "r_max: 8.0\n"
+            )
+
+            # torch.jit.load only needs the file path to exist for this mock.
+            (base_input_dir / "deployed.pth").touch()
+
+            dummy_model = object()
+
+            def fake_jit_load(path, _extra_files=None):
+                self.assertEqual(Path(path), base_input_dir / "deployed.pth")
+
+                if _extra_files is not None:
+                    _extra_files["type_names"] = b"V Cu Sn"
+                    _extra_files["r_max"] = b"8.0"
+
+                return dummy_model
+
+            with patch(
+                "abics.applications.latgas_abinitio_interface.nequip.torch.jit.load",
+                side_effect=fake_jit_load,
+            ):
+                solver_input = NequipSolver.Input(ignore_species=None)
+                solver_input.from_directory(str(base_input_dir))
+
+            self.assertIs(solver_input.model, dummy_model)
+            self.assertEqual(solver_input.element_list, ["V", "Cu", "Sn"])
+            self.assertEqual(solver_input.r_max, 8.0)
+
+
+if __name__ == "__main__":
+    unittest.main()