add check_completeness

Signed-off-by: Jinzhe Zeng <[email protected]>

deepmd_utils/model_format/network.py

@@ -3,6 +3,7 @@
 
 See issue #2982 for more information.
 """
+import itertools
 import json
 from typing import (
     Dict,
@@ -438,10 +439,12 @@ class NetworkCollection:
     def __init__(
         self,
         ndim: int,
+        ntypes: int,
         network_type: str = "network",
         networks: Dict[Union[str, tuple], Union[NativeNet, dict]] = {},
     ):
         self.ndim = ndim
+        self.ntypes = ntypes
         if network_type == "network":
             self.network_type = NativeNet
         elif network_type == "embedding_network":
@@ -452,6 +455,18 @@ def __init__(
         for kk, vv in networks.items():
             self[kk] = vv
 
+    def check_completeness(self):
+        """Check whether the collection is complete.
+
+        Raises
+        ------
+        RuntimeError
+            If the collection is incomplete.
+        """
+        for tt in itertools.product(range(self.ntypes), repeat=self.ndim):
+            if tuple(tt) not in self._networks:
+                raise RuntimeError(f"network for {tt} not found")
+
     def _convert_key(self, key):
         if isinstance(key, tuple):
             pass
@@ -491,6 +506,7 @@ def serialize(self) -> dict:
             raise NotImplementedError(self.network_type)
         return {
             "ndim": self.ndim,
+            "ntypes": self.ntypes,
             "network_type": network_type_name,
             "networks": {
                 ("_".join(["type"] + [str(tt) for tt in key])): value.serialize()

source/tests/test_model_format_utils.py

@@ -136,9 +136,14 @@ def setUp(self) -> None:
         }
 
     def test_two_dim(self):
-        networks = NetworkCollection(ndim=2)
+        networks = NetworkCollection(ndim=2, ntypes=2)
         networks[(0, 0)] = self.network
+        networks[(1, 1)] = self.network
         networks[(0, 1)] = self.network
+        with self.assertRaises(RuntimeError):
+            networks.check_completeness()
+        networks[(1, 0)] = self.network
+        networks.check_completeness()
         np.testing.assert_equal(
             networks.serialize(),
             NetworkCollection.deserialize(networks.serialize()).serialize(),
@@ -148,9 +153,12 @@ def test_two_dim(self):
         )
 
     def test_one_dim(self):
-        networks = NetworkCollection(ndim=1)
+        networks = NetworkCollection(ndim=1, ntypes=2)
         networks[(0,)] = self.network
+        with self.assertRaises(RuntimeError):
+            networks.check_completeness()
         networks[(1,)] = self.network
+        networks.check_completeness()
         np.testing.assert_equal(
             networks.serialize(),
             NetworkCollection.deserialize(networks.serialize()).serialize(),
@@ -160,8 +168,9 @@ def test_one_dim(self):
         )
 
     def test_zero_dim(self):
-        networks = NetworkCollection(ndim=0)
+        networks = NetworkCollection(ndim=0, ntypes=2)
         networks[()] = self.network
+        networks.check_completeness()
         np.testing.assert_equal(
             networks.serialize(),
             NetworkCollection.deserialize(networks.serialize()).serialize(),