chore: avoid unnecessary deepcopy (#4327)

This PR saves about 1 minute in each CI job. See devel vs this branch:

![image](https://github.com/user-attachments/assets/13fd74b6-bec0-4e44-9f5e-54c06be24531)

![image](https://github.com/user-attachments/assets/7987ab4a-3b4b-4585-9503-cccdfce81a26)


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

## Summary by CodeRabbit

- **New Features**
- Enhanced data handling in various models by replacing deep copy with
shallow copy in the `deserialize` method, improving performance.
- Added `exclude_types` parameter in multiple fitting classes, allowing
for the exclusion of specific atom types during fitting.

- **Bug Fixes**
- Improved robustness in output definition retrieval by adding fallback
mechanisms in certain models.

- **Documentation**
- Updated comments and docstrings for clarity in methods across several
classes.

- **Refactor**
- Streamlined data handling processes in multiple models by eliminating
unnecessary deep copying.
- Adjusted version compatibility checks in several `deserialize`
methods.

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

Signed-off-by: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>

deepmd/dpmodel/atomic_model/dp_atomic_model.py

@@ -1,5 +1,4 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
-import copy
 from typing import (
     Optional,
 )
@@ -208,7 +207,7 @@ def serialize(self) -> dict:
 
     @classmethod
     def deserialize(cls, data) -> "DPAtomicModel":
-        data = copy.deepcopy(data)
+        data = data.copy()
         check_version_compatibility(data.pop("@version", 1), 2, 2)
         data.pop("@class")
         data.pop("type")

deepmd/dpmodel/atomic_model/linear_atomic_model.py

@@ -1,5 +1,4 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
-import copy
 from typing import (
     Optional,
     Union,
@@ -305,7 +304,7 @@ def serialize(self) -> dict:
 
     @classmethod
     def deserialize(cls, data: dict) -> "LinearEnergyAtomicModel":
-        data = copy.deepcopy(data)
+        data = data.copy()
         check_version_compatibility(data.pop("@version", 2), 2, 2)
         data.pop("@class", None)
         data.pop("type", None)
@@ -418,7 +417,7 @@ def serialize(self) -> dict:
 
     @classmethod
     def deserialize(cls, data) -> "DPZBLLinearEnergyAtomicModel":
-        data = copy.deepcopy(data)
+        data = data.copy()
         check_version_compatibility(data.pop("@version", 1), 2, 2)
         models = [
             BaseAtomicModel.get_class_by_type(model["type"]).deserialize(model)

deepmd/dpmodel/atomic_model/pairtab_atomic_model.py

@@ -1,5 +1,4 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
-import copy
 from typing import (
     Optional,
     Union,
@@ -174,7 +173,7 @@ def serialize(self) -> dict:
 
     @classmethod
     def deserialize(cls, data) -> "PairTabAtomicModel":
-        data = copy.deepcopy(data)
+        data = data.copy()
         check_version_compatibility(data.pop("@version", 1), 2, 2)
         data.pop("@class")
         data.pop("type")

deepmd/dpmodel/descriptor/se_e2_a.py

@@ -1,5 +1,4 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
-import copy
 import itertools
 from typing import (
     Any,
@@ -460,7 +459,7 @@ def serialize(self) -> dict:
     @classmethod
     def deserialize(cls, data: dict) -> "DescrptSeA":
         """Deserialize from dict."""
-        data = copy.deepcopy(data)
+        data = data.copy()
         check_version_compatibility(data.pop("@version", 1), 2, 1)
         data.pop("@class", None)
         data.pop("type", None)

deepmd/dpmodel/descriptor/se_r.py

@@ -1,5 +1,4 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
-import copy
 from typing import (
     Any,
     Optional,
@@ -388,7 +387,7 @@ def serialize(self) -> dict:
     @classmethod
     def deserialize(cls, data: dict) -> "DescrptSeR":
         """Deserialize from dict."""
-        data = copy.deepcopy(data)
+        data = data.copy()
         check_version_compatibility(data.pop("@version", 1), 2, 1)
         data.pop("@class", None)
         data.pop("type", None)

deepmd/dpmodel/descriptor/se_t.py

@@ -1,5 +1,4 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
-import copy
 import itertools
 from typing import (
     Optional,
@@ -388,7 +387,7 @@ def serialize(self) -> dict:
     @classmethod
     def deserialize(cls, data: dict) -> "DescrptSeT":
         """Deserialize from dict."""
-        data = copy.deepcopy(data)
+        data = data.copy()
         check_version_compatibility(data.pop("@version", 1), 2, 1)
         data.pop("@class", None)
         data.pop("type", None)

deepmd/dpmodel/fitting/dipole_fitting.py

@@ -1,5 +1,4 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
-import copy
 from typing import (
     Any,
     Optional,
@@ -156,7 +155,7 @@ def serialize(self) -> dict:
 
     @classmethod
     def deserialize(cls, data: dict) -> "GeneralFitting":
-        data = copy.deepcopy(data)
+        data = data.copy()
         check_version_compatibility(data.pop("@version", 1), 2, 1)
         var_name = data.pop("var_name", None)
         assert var_name == "dipole"

deepmd/dpmodel/fitting/dos_fitting.py

@@ -1,5 +1,4 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
-import copy
 from typing import (
     TYPE_CHECKING,
     Optional,
@@ -73,7 +72,7 @@ def __init__(
 
     @classmethod
     def deserialize(cls, data: dict) -> "GeneralFitting":
-        data = copy.deepcopy(data)
+        data = data.copy()
         check_version_compatibility(data.pop("@version", 1), 2, 1)
         data["numb_dos"] = data.pop("dim_out")
         data.pop("tot_ener_zero", None)

deepmd/dpmodel/fitting/ener_fitting.py

@@ -1,5 +1,4 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
-import copy
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -73,7 +72,7 @@ def __init__(
 
     @classmethod
     def deserialize(cls, data: dict) -> "GeneralFitting":
-        data = copy.deepcopy(data)
+        data = data.copy()
         check_version_compatibility(data.pop("@version", 1), 2, 1)
         data.pop("var_name")
         data.pop("dim_out")

deepmd/dpmodel/fitting/general_fitting.py

@@ -1,5 +1,4 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
-import copy
 from abc import (
     abstractmethod,
 )
@@ -320,7 +319,7 @@ def serialize(self) -> dict:
 
     @classmethod
     def deserialize(cls, data: dict) -> "GeneralFitting":
-        data = copy.deepcopy(data)
+        data = data.copy()
         data.pop("@class")
         data.pop("type")
         variables = data.pop("@variables")

deepmd/dpmodel/fitting/invar_fitting.py

@@ -1,5 +1,4 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
-import copy
 from typing import (
     Any,
     Optional,
@@ -179,7 +178,7 @@ def serialize(self) -> dict:
 
     @classmethod
     def deserialize(cls, data: dict) -> "GeneralFitting":
-        data = copy.deepcopy(data)
+        data = data.copy()
         check_version_compatibility(data.pop("@version", 1), 2, 1)
         return super().deserialize(data)
 

deepmd/dpmodel/fitting/polarizability_fitting.py

@@ -1,5 +1,4 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
-import copy
 from typing import (
     Any,
     Optional,
@@ -197,7 +196,7 @@ def serialize(self) -> dict:
 
     @classmethod
     def deserialize(cls, data: dict) -> "GeneralFitting":
-        data = copy.deepcopy(data)
+        data = data.copy()
         check_version_compatibility(data.pop("@version", 1), 3, 1)
         var_name = data.pop("var_name", None)
         assert var_name == "polar"

deepmd/dpmodel/fitting/property_fitting.py

@@ -1,5 +1,4 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
-import copy
 from typing import (
     Optional,
     Union,
@@ -111,7 +110,7 @@ def __init__(
 
     @classmethod
     def deserialize(cls, data: dict) -> "PropertyFittingNet":
-        data = copy.deepcopy(data)
+        data = data.copy()
         check_version_compatibility(data.pop("@version"), 2, 1)
         data.pop("dim_out")
         data.pop("var_name")

deepmd/dpmodel/utils/network.py

@@ -4,7 +4,6 @@
 See issue #2982 for more information.
 """
 
-import copy
 import itertools
 from typing import (
     Callable,
@@ -135,7 +134,7 @@ def deserialize(cls, data: dict) -> "NativeLayer":
         data : dict
             The dict to deserialize from.
         """
-        data = copy.deepcopy(data)
+        data = data.copy()
         check_version_compatibility(data.pop("@version", 1), 1, 1)
         data.pop("@class", None)
         variables = data.pop("@variables")
@@ -404,7 +403,7 @@ def deserialize(cls, data: dict) -> "LayerNorm":
         data : dict
             The dict to deserialize from.
         """
-        data = copy.deepcopy(data)
+        data = data.copy()
         check_version_compatibility(data.pop("@version", 1), 1, 1)
         data.pop("@class", None)
         variables = data.pop("@variables")
@@ -673,7 +672,7 @@ def deserialize(cls, data: dict) -> "EmbeddingNet":
             data : dict
                 The dict to deserialize from.
             """
-            data = copy.deepcopy(data)
+            data = data.copy()
             check_version_compatibility(data.pop("@version", 1), 2, 1)
             data.pop("@class", None)
             layers = data.pop("layers")
@@ -778,7 +777,7 @@ def deserialize(cls, data: dict) -> "FittingNet":
             data : dict
                 The dict to deserialize from.
             """
-            data = copy.deepcopy(data)
+            data = data.copy()
             check_version_compatibility(data.pop("@version", 1), 1, 1)
             data.pop("@class", None)
             layers = data.pop("layers")

deepmd/pt/model/atomic_model/base_atomic_model.py

@@ -1,6 +1,5 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 
-import copy
 import logging
 from typing import (
     Callable,
@@ -331,7 +330,7 @@ def serialize(self) -> dict:
 
     @classmethod
     def deserialize(cls, data: dict) -> "BaseAtomicModel":
-        data = copy.deepcopy(data)
+        data = data.copy()
         variables = data.pop("@variables", None)
         variables = (
             {"out_bias": None, "out_std": None} if variables is None else variables

deepmd/pt/model/atomic_model/dp_atomic_model.py

@@ -1,5 +1,4 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
-import copy
 import functools
 import logging
 from typing import (
@@ -149,7 +148,7 @@ def serialize(self) -> dict:
 
     @classmethod
     def deserialize(cls, data) -> "DPAtomicModel":
-        data = copy.deepcopy(data)
+        data = data.copy()
         check_version_compatibility(data.pop("@version", 1), 2, 1)
         data.pop("@class", None)
         data.pop("type", None)

deepmd/pt/model/atomic_model/linear_atomic_model.py

@@ -1,5 +1,4 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
-import copy
 from typing import (
     Callable,
     Optional,
@@ -369,7 +368,7 @@ def serialize(self) -> dict:
 
     @classmethod
     def deserialize(cls, data: dict) -> "LinearEnergyAtomicModel":
-        data = copy.deepcopy(data)
+        data = data.copy()
         check_version_compatibility(data.pop("@version", 2), 2, 1)
         data.pop("@class", None)
         data.pop("type", None)
@@ -563,7 +562,7 @@ def serialize(self) -> dict:
 
     @classmethod
     def deserialize(cls, data) -> "DPZBLLinearEnergyAtomicModel":
-        data = copy.deepcopy(data)
+        data = data.copy()
         check_version_compatibility(data.pop("@version", 1), 2, 1)
         models = [
             BaseAtomicModel.get_class_by_type(model["type"]).deserialize(model)

deepmd/pt/model/atomic_model/pairtab_atomic_model.py

@@ -1,5 +1,4 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
-import copy
 from typing import (
     Callable,
     Optional,
@@ -195,7 +194,7 @@ def serialize(self) -> dict:
 
     @classmethod
     def deserialize(cls, data) -> "PairTabAtomicModel":
-        data = copy.deepcopy(data)
+        data = data.copy()
         check_version_compatibility(data.pop("@version", 1), 2, 1)
         tab = PairTab.deserialize(data.pop("tab"))
         data.pop("@class", None)

deepmd/pt/model/model/dipole_model.py

@@ -1,7 +1,4 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
-from copy import (
-    deepcopy,
-)
 from typing import (
     Optional,
 )
@@ -40,19 +37,19 @@ def __init__(
     def translated_output_def(self):
         out_def_data = self.model_output_def().get_data()
         output_def = {
-            "dipole": deepcopy(out_def_data["dipole"]),
-            "global_dipole": deepcopy(out_def_data["dipole_redu"]),
+            "dipole": out_def_data["dipole"],
+            "global_dipole": out_def_data["dipole_redu"],
         }
         if self.do_grad_r("dipole"):
-            output_def["force"] = deepcopy(out_def_data["dipole_derv_r"])
+            output_def["force"] = out_def_data["dipole_derv_r"]
             output_def["force"].squeeze(-2)
         if self.do_grad_c("dipole"):
-            output_def["virial"] = deepcopy(out_def_data["dipole_derv_c_redu"])
+            output_def["virial"] = out_def_data["dipole_derv_c_redu"]
             output_def["virial"].squeeze(-2)
-            output_def["atom_virial"] = deepcopy(out_def_data["dipole_derv_c"])
+            output_def["atom_virial"] = out_def_data["dipole_derv_c"]
             output_def["atom_virial"].squeeze(-3)
         if "mask" in out_def_data:
-            output_def["mask"] = deepcopy(out_def_data["mask"])
+            output_def["mask"] = out_def_data["mask"]
         return output_def
 
     def forward(

deepmd/pt/model/model/dos_model.py

@@ -1,7 +1,4 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
-from copy import (
-    deepcopy,
-)
 from typing import (
     Optional,
 )
@@ -40,11 +37,11 @@ def __init__(
     def translated_output_def(self):
         out_def_data = self.model_output_def().get_data()
         output_def = {
-            "atom_dos": deepcopy(out_def_data["dos"]),
-            "dos": deepcopy(out_def_data["dos_redu"]),
+            "atom_dos": out_def_data["dos"],
+            "dos": out_def_data["dos_redu"],
         }
         if "mask" in out_def_data:
-            output_def["mask"] = deepcopy(out_def_data["mask"])
+            output_def["mask"] = out_def_data["mask"]
         return output_def
 
     def forward(