chore: refactor

deepmodeling · Mar 8, 2024 · 831158c · 831158c
1 parent 3de4709
commit 831158c
Show file tree

Hide file tree

Showing 4 changed files with 24 additions and 8 deletions.
diff --git a/deepmd/dpmodel/fitting/general_fitting.py b/deepmd/dpmodel/fitting/general_fitting.py
@@ -191,8 +191,6 @@ def __setitem__(self, key, value):
             self.aparam_inv_std = value
         elif key in ["scale"]:
             self.scale = value
-        elif key in ["constant_matrix"]:
-            self.constant_matrix = value
         else:
             raise KeyError(key)
 
@@ -209,8 +207,6 @@ def __getitem__(self, key):
             return self.aparam_inv_std
         elif key in ["scale"]:
             return self.scale
-        elif key in ["constant_matrix"]:
-            return self.constant_matrix
         else:
             raise KeyError(key)
 

diff --git a/deepmd/dpmodel/fitting/polarizability_fitting.py b/deepmd/dpmodel/fitting/polarizability_fitting.py
@@ -173,6 +173,18 @@ def _net_out_dim(self):
             else self.embedding_width * self.embedding_width
         )
 
+    def __setitem__(self, key, value):
+        if key in ["constant_matrix"]:
+            self.constant_matrix = value
+        else:
+            super().__setitem__(key, value)
+
+    def __getitem__(self, key):
+        if key in ["constant_matrix"]:
+            return self.constant_matrix
+        else:
+            super().__getitem__(key, value)
+
     def serialize(self) -> dict:
         data = super().serialize()
         data["type"] = "polar"

diff --git a/deepmd/pt/model/task/fitting.py b/deepmd/pt/model/task/fitting.py
@@ -415,8 +415,6 @@ def __setitem__(self, key, value):
             self.aparam_inv_std = value
         elif key in ["scale"]:
             self.scale = value
-        elif key in ["constant_matrix"]:
-            self.constant_matrix = value
         else:
             raise KeyError(key)
 
@@ -433,8 +431,6 @@ def __getitem__(self, key):
             return self.aparam_inv_std
         elif key in ["scale"]:
             return self.scale
-        elif key in ["constant_matrix"]:
-            return self.constant_matrix
         else:
             raise KeyError(key)
 

diff --git a/deepmd/pt/model/task/polarizability.py b/deepmd/pt/model/task/polarizability.py
@@ -152,6 +152,18 @@ def _net_out_dim(self):
             else self.embedding_width * self.embedding_width
         )
 
+    def __setitem__(self, key, value):
+        if key in ["constant_matrix"]:
+            self.constant_matrix = value
+        else:
+            super().__setitem__(key, value)
+
+    def __getitem__(self, key):
+        if key in ["constant_matrix"]:
+            return self.constant_matrix
+        else:
+            super().__getitem__(key, value)
+
     def serialize(self) -> dict:
         data = super().serialize()
         data["type"] = "polar"