Handle custom data types in mixed systems (#855)

## Summary
- generalize mixed-system loader to split all registered data types
- test fparam and aparam round-trip through deepmd mixed format

## Testing
- `cd tests && pytest
test_deepmd_mixed.py::TestMixedSystemWithFparamAparam -q`


------

https://chatgpt.com/codex/tasks/task_b_688b28aaa2d48332bd3a9bc79849b7a0

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

* **New Features**
* Enhanced support for custom registered data types in system data
handling, allowing additional data arrays beyond standard types to be
processed and saved.
* Added comprehensive tests for new custom data types, verifying their
correct saving, loading, and integrity within mixed system files.

* **Tests**
* Introduced new test cases to ensure custom data types are properly
managed and validated in mixed system workflows.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Duo <[email protected]>

Author: 3 people

dpdata/deepmd/mixed.py

@@ -4,6 +4,8 @@
 
 import numpy as np
 
+import dpdata
+
 from .comp import dump as comp_dump
 from .comp import to_system_data as comp_to_system_data
 
@@ -27,10 +29,32 @@ def to_system_data(folder, type_map=None, labels=True):
         all_real_atom_types_concat = index_map[all_real_atom_types_concat]
     all_cells_concat = data["cells"]
     all_coords_concat = data["coords"]
+
+    # handle custom registered data types
     if labels:
-        all_eners_concat = data.get("energies")
-        all_forces_concat = data.get("forces")
-        all_virs_concat = data.get("virials")
+        dtypes = dpdata.system.LabeledSystem.DTYPES
+    else:
+        dtypes = dpdata.system.System.DTYPES
+    reserved = {
+        "atom_numbs",
+        "atom_names",
+        "atom_types",
+        "real_atom_names",
+        "real_atom_types",
+        "cells",
+        "coords",
+        "orig",
+        "nopbc",
+    }
+    extra_data = {}
+    for dtype in dtypes:
+        name = dtype.name
+        if name in reserved:
+            continue
+        if not (len(dtype.shape) and dtype.shape[0] == dpdata.system.Axis.NFRAMES):
+            continue
+        if name in data:
+            extra_data[name] = data.pop(name)
 
     data_list = []
     while True:
@@ -56,16 +80,12 @@ def to_system_data(folder, type_map=None, labels=True):
         all_cells_concat = all_cells_concat[rest_idx]
         temp_data["coords"] = all_coords_concat[temp_idx]
         all_coords_concat = all_coords_concat[rest_idx]
-        if labels:
-            if all_eners_concat is not None and all_eners_concat.size > 0:
-                temp_data["energies"] = all_eners_concat[temp_idx]
-                all_eners_concat = all_eners_concat[rest_idx]
-            if all_forces_concat is not None and all_forces_concat.size > 0:
-                temp_data["forces"] = all_forces_concat[temp_idx]
-                all_forces_concat = all_forces_concat[rest_idx]
-            if all_virs_concat is not None and all_virs_concat.size > 0:
-                temp_data["virials"] = all_virs_concat[temp_idx]
-                all_virs_concat = all_virs_concat[rest_idx]
+
+        for name in extra_data:
+            all_dtype_concat = extra_data[name]
+            temp_data[name] = all_dtype_concat[temp_idx]
+            extra_data[name] = all_dtype_concat[rest_idx]
+
         data_list.append(temp_data)
     return data_list
 

tests/test_deepmd_mixed.py

@@ -15,6 +15,11 @@
 )
 from context import dpdata
 
+from dpdata.data_type import (
+    Axis,
+    DataType,
+)
+
 
 class TestMixedMultiSystemsDumpLoad(
     unittest.TestCase, CompLabeledMultiSys, MultiSystems, MSAllIsNoPBC
@@ -455,5 +460,140 @@ def tearDown(self):
             shutil.rmtree("tmp.deepmd.mixed.single")
 
 
-if __name__ == "__main__":
-    unittest.main()
+class TestMixedSystemWithFparamAparam(
+    unittest.TestCase, CompLabeledMultiSys, MultiSystems, MSAllIsNoPBC
+):
+    def setUp(self):
+        self.places = 6
+        self.e_places = 6
+        self.f_places = 6
+        self.v_places = 6
+
+        new_datatypes = [
+            DataType(
+                "fparam",
+                np.ndarray,
+                shape=(Axis.NFRAMES, 2),
+                required=False,
+            ),
+            DataType(
+                "aparam",
+                np.ndarray,
+                shape=(Axis.NFRAMES, Axis.NATOMS, 3),
+                required=False,
+            ),
+        ]
+
+        for datatype in new_datatypes:
+            dpdata.System.register_data_type(datatype)
+            dpdata.LabeledSystem.register_data_type(datatype)
+
+        # C1H4
+        system_1 = dpdata.LabeledSystem(
+            "gaussian/methane.gaussianlog", fmt="gaussian/log"
+        )
+
+        # C1H3
+        system_2 = dpdata.LabeledSystem(
+            "gaussian/methane_sub.gaussianlog", fmt="gaussian/log"
+        )
+
+        tmp_data_1 = system_1.data.copy()
+        nframes_1 = tmp_data_1["coords"].shape[0]
+        natoms_1 = tmp_data_1["atom_types"].shape[0]
+        tmp_data_1["fparam"] = np.random.random([nframes_1, 2])
+        tmp_data_1["aparam"] = np.random.random([nframes_1, natoms_1, 3])
+        system_1_with_params = dpdata.LabeledSystem(data=tmp_data_1)
+
+        tmp_data_2 = system_2.data.copy()
+        nframes_2 = tmp_data_2["coords"].shape[0]
+        natoms_2 = tmp_data_2["atom_types"].shape[0]
+        tmp_data_2["fparam"] = np.random.random([nframes_2, 2])
+        tmp_data_2["aparam"] = np.random.random([nframes_2, natoms_2, 3])
+        system_2_with_params = dpdata.LabeledSystem(data=tmp_data_2)
+
+        tmp_data_3 = system_1.data.copy()
+        nframes_3 = tmp_data_3["coords"].shape[0]
+        tmp_data_3["atom_numbs"] = [1, 1, 1, 2]
+        tmp_data_3["atom_names"] = ["C", "H", "A", "B"]
+        tmp_data_3["atom_types"] = np.array([0, 1, 2, 3, 3])
+        natoms_3 = len(tmp_data_3["atom_types"])
+        tmp_data_3["fparam"] = np.random.random([nframes_3, 2])
+        tmp_data_3["aparam"] = np.random.random([nframes_3, natoms_3, 3])
+        # C1H1A1B2 with params
+        system_3_with_params = dpdata.LabeledSystem(data=tmp_data_3)
+
+        self.ms = dpdata.MultiSystems(
+            system_1_with_params, system_2_with_params, system_3_with_params
+        )
+
+        self.ms.to_deepmd_npy_mixed("tmp.deepmd.fparam.aparam")
+        self.place_holder_ms = dpdata.MultiSystems()
+        self.place_holder_ms.from_deepmd_npy(
+            "tmp.deepmd.fparam.aparam", fmt="deepmd/npy"
+        )
+        self.systems = dpdata.MultiSystems()
+        self.systems.from_deepmd_npy_mixed(
+            "tmp.deepmd.fparam.aparam", fmt="deepmd/npy/mixed"
+        )
+
+        self.ms_1 = self.ms
+        self.ms_2 = self.systems
+
+        mixed_sets = glob("tmp.deepmd.fparam.aparam/*/set.*")
+        for i in mixed_sets:
+            self.assertEqual(
+                os.path.exists(os.path.join(i, "real_atom_types.npy")), True
+            )
+
+        self.system_names = ["C1H4A0B0", "C1H3A0B0", "C1H1A1B2"]
+        self.system_sizes = {"C1H4A0B0": 1, "C1H3A0B0": 1, "C1H1A1B2": 1}
+        self.atom_names = ["C", "H", "A", "B"]
+
+    def tearDown(self):
+        if os.path.exists("tmp.deepmd.fparam.aparam"):
+            shutil.rmtree("tmp.deepmd.fparam.aparam")
+
+    def test_len(self):
+        self.assertEqual(len(self.ms), 3)
+        self.assertEqual(len(self.systems), 3)
+
+    def test_get_nframes(self):
+        self.assertEqual(self.ms.get_nframes(), 3)
+        self.assertEqual(self.systems.get_nframes(), 3)
+
+    def test_str(self):
+        self.assertEqual(str(self.ms), "MultiSystems (3 systems containing 3 frames)")
+        self.assertEqual(
+            str(self.systems), "MultiSystems (3 systems containing 3 frames)"
+        )
+
+    def test_fparam_exists(self):
+        for formula in self.system_names:
+            if formula in self.ms.systems:
+                self.assertTrue("fparam" in self.ms[formula].data)
+            if formula in self.systems.systems:
+                self.assertTrue("fparam" in self.systems[formula].data)
+
+        for formula in self.system_names:
+            if formula in self.ms.systems and formula in self.systems.systems:
+                np.testing.assert_almost_equal(
+                    self.ms[formula].data["fparam"],
+                    self.systems[formula].data["fparam"],
+                    decimal=self.places,
+                )
+
+    def test_aparam_exists(self):
+        for formula in self.system_names:
+            if formula in self.ms.systems:
+                self.assertTrue("aparam" in self.ms[formula].data)
+            if formula in self.systems.systems:
+                self.assertTrue("aparam" in self.systems[formula].data)
+
+        for formula in self.system_names:
+            if formula in self.ms.systems and formula in self.systems.systems:
+                np.testing.assert_almost_equal(
+                    self.ms[formula].data["aparam"],
+                    self.systems[formula].data["aparam"],
+                    decimal=self.places,
+                )