diff --git a/.github/workflows/test_python.yml b/.github/workflows/test_python.yml index ebcbaaa81d..58359349bd 100644 --- a/.github/workflows/test_python.yml +++ b/.github/workflows/test_python.yml @@ -19,7 +19,7 @@ jobs: fail-fast: false matrix: group: [1, 2, 3, 4, 5, 6] - python: ["3.9", "3.12"] + python: ["3.10", "3.12"] steps: - uses: actions/checkout@v6 @@ -85,7 +85,7 @@ jobs: strategy: fail-fast: false matrix: - python: ["3.9", "3.12"] + python: ["3.10", "3.12"] needs: testpython steps: - name: Get durations from cache diff --git a/AGENTS.md b/AGENTS.md index 9b07cbd9e0..c629a08def 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -94,7 +94,7 @@ source/ # C++ source code and tests ### Build Dependencies and Setup -- **Python 3.9+** required +- **Python 3.10+** required - **Virtual environment** strongly recommended: `uv venv venv && source venv/bin/activate` - **Backend dependencies**: TensorFlow, PyTorch, JAX, or Paddle (install before building) - **Build tools**: CMake, C++ compiler, scikit-build-core diff --git a/backend/dynamic_metadata.py b/backend/dynamic_metadata.py index a66e9a2759..e7763cac84 100644 --- a/backend/dynamic_metadata.py +++ b/backend/dynamic_metadata.py @@ -3,9 +3,6 @@ from pathlib import ( Path, ) -from typing import ( - Optional, -) from .find_pytorch import ( get_pt_requirement, @@ -31,7 +28,7 @@ def __dir__() -> list[str]: def dynamic_metadata( field: str, - settings: Optional[dict[str, object]] = None, + settings: dict[str, object] | None = None, ): assert field in ["optional-dependencies", "entry-points", "scripts"] _, _, find_libpython_requires, extra_scripts, tf_version, pt_version = ( diff --git a/backend/find_paddle.py b/backend/find_paddle.py index 247526c717..c9b0319504 100644 --- a/backend/find_paddle.py +++ b/backend/find_paddle.py @@ -17,14 +17,10 @@ from sysconfig import ( get_path, ) -from typing import ( - Optional, - Union, -) @lru_cache -def find_paddle() -> tuple[Optional[str], list[str]]: +def find_paddle() -> tuple[str | None, list[str]]: """Find PaddlePadle library. Tries to find PaddlePadle in the order of: @@ -111,7 +107,7 @@ def get_pd_requirement(pd_version: str = "") -> dict: @lru_cache -def get_pd_version(pd_path: Optional[Union[str, Path]]) -> str: +def get_pd_version(pd_path: str | Path | None) -> str: """Get Paddle version from a Paddle Python library path. Parameters diff --git a/backend/find_pytorch.py b/backend/find_pytorch.py index df52e63219..ae7c733dcc 100644 --- a/backend/find_pytorch.py +++ b/backend/find_pytorch.py @@ -17,10 +17,6 @@ from sysconfig import ( get_path, ) -from typing import ( - Optional, - Union, -) from packaging.specifiers import ( SpecifierSet, @@ -35,7 +31,7 @@ @lru_cache -def find_pytorch() -> tuple[Optional[str], list[str]]: +def find_pytorch() -> tuple[str | None, list[str]]: """Find PyTorch library. Tries to find PyTorch in the order of: @@ -150,7 +146,7 @@ def get_pt_requirement(pt_version: str = "") -> dict: @lru_cache -def get_pt_version(pt_path: Optional[Union[str, Path]]) -> str: +def get_pt_version(pt_path: str | Path | None) -> str: """Get TF version from a TF Python library path. Parameters diff --git a/backend/find_tensorflow.py b/backend/find_tensorflow.py index 457e7a726c..4b508a4d14 100644 --- a/backend/find_tensorflow.py +++ b/backend/find_tensorflow.py @@ -17,10 +17,6 @@ from sysconfig import ( get_path, ) -from typing import ( - Optional, - Union, -) from packaging.specifiers import ( SpecifierSet, @@ -32,7 +28,7 @@ @lru_cache -def find_tensorflow() -> tuple[Optional[str], list[str]]: +def find_tensorflow() -> tuple[str | None, list[str]]: """Find TensorFlow library. Tries to find TensorFlow in the order of: @@ -208,7 +204,7 @@ def get_tf_requirement(tf_version: str = "") -> dict: @lru_cache -def get_tf_version(tf_path: Optional[Union[str, Path]]) -> str: +def get_tf_version(tf_path: str | Path | None) -> str: """Get TF version from a TF Python library path. Parameters diff --git a/deepmd/backend/backend.py b/deepmd/backend/backend.py index 3263169f6f..58dcfe427d 100644 --- a/deepmd/backend/backend.py +++ b/deepmd/backend/backend.py @@ -2,13 +2,15 @@ from abc import ( abstractmethod, ) +from collections.abc import ( + Callable, +) from enum import ( Flag, auto, ) from typing import ( TYPE_CHECKING, - Callable, ClassVar, ) diff --git a/deepmd/backend/dpmodel.py b/deepmd/backend/dpmodel.py index 7c21b256ae..31585aa7a6 100644 --- a/deepmd/backend/dpmodel.py +++ b/deepmd/backend/dpmodel.py @@ -1,7 +1,9 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +from collections.abc import ( + Callable, +) from typing import ( TYPE_CHECKING, - Callable, ClassVar, ) diff --git a/deepmd/backend/jax.py b/deepmd/backend/jax.py index 7a714c2090..9c0055b4f2 100644 --- a/deepmd/backend/jax.py +++ b/deepmd/backend/jax.py @@ -1,10 +1,12 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +from collections.abc import ( + Callable, +) from importlib.util import ( find_spec, ) from typing import ( TYPE_CHECKING, - Callable, ClassVar, ) diff --git a/deepmd/backend/paddle.py b/deepmd/backend/paddle.py index b1f664e76a..670130e86a 100644 --- a/deepmd/backend/paddle.py +++ b/deepmd/backend/paddle.py @@ -1,10 +1,12 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +from collections.abc import ( + Callable, +) from importlib.util import ( find_spec, ) from typing import ( TYPE_CHECKING, - Callable, ClassVar, ) diff --git a/deepmd/backend/pytorch.py b/deepmd/backend/pytorch.py index f5b0dd92b2..d155ef1f41 100644 --- a/deepmd/backend/pytorch.py +++ b/deepmd/backend/pytorch.py @@ -1,10 +1,12 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +from collections.abc import ( + Callable, +) from importlib.util import ( find_spec, ) from typing import ( TYPE_CHECKING, - Callable, ClassVar, ) diff --git a/deepmd/backend/suffix.py b/deepmd/backend/suffix.py index e77aecb5d9..3d1602da0e 100644 --- a/deepmd/backend/suffix.py +++ b/deepmd/backend/suffix.py @@ -4,10 +4,6 @@ from pathlib import ( Path, ) -from typing import ( - Optional, - Union, -) from deepmd.backend.backend import ( Backend, @@ -16,9 +12,9 @@ def format_model_suffix( filename: str, - feature: Optional[Backend.Feature] = None, - preferred_backend: Optional[Union[str, type["Backend"]]] = None, - strict_prefer: Optional[bool] = None, + feature: Backend.Feature | None = None, + preferred_backend: str | type["Backend"] | None = None, + strict_prefer: bool | None = None, ) -> str: """Check and format the suffixes of a filename. diff --git a/deepmd/backend/tensorflow.py b/deepmd/backend/tensorflow.py index 6b73d7c469..244b4d9980 100644 --- a/deepmd/backend/tensorflow.py +++ b/deepmd/backend/tensorflow.py @@ -1,10 +1,12 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +from collections.abc import ( + Callable, +) from importlib.util import ( find_spec, ) from typing import ( TYPE_CHECKING, - Callable, ClassVar, ) diff --git a/deepmd/calculator.py b/deepmd/calculator.py index 83ebd73e4c..dca2343ce0 100644 --- a/deepmd/calculator.py +++ b/deepmd/calculator.py @@ -88,9 +88,9 @@ def __init__( self, model: Union[str, "Path"], label: str = "DP", - type_dict: Optional[dict[str, int]] = None, + type_dict: dict[str, int] | None = None, neighbor_list: Optional["NeighborList"] = None, - head: Optional[str] = None, + head: str | None = None, **kwargs: Any, ) -> None: Calculator.__init__(self, label=label, **kwargs) diff --git a/deepmd/common.py b/deepmd/common.py index 26b655f876..98cf2461bd 100644 --- a/deepmd/common.py +++ b/deepmd/common.py @@ -15,14 +15,13 @@ TYPE_CHECKING, Any, TypeVar, - Union, get_args, ) try: from typing import Literal # python >=3.8 except ImportError: - from typing_extensions import Literal # type: ignore + from typing import Literal # type: ignore import numpy as np import yaml @@ -159,7 +158,7 @@ def j_deprecated( return jdata[key] -def j_loader(filename: Union[str, Path]) -> dict[str, Any]: +def j_loader(filename: str | Path) -> dict[str, Any]: """Load yaml or json settings file. Parameters @@ -188,7 +187,7 @@ def j_loader(filename: Union[str, Path]) -> dict[str, Any]: raise TypeError("config file must be json, or yaml/yml") -def expand_sys_str(root_dir: Union[str, Path]) -> list[str]: +def expand_sys_str(root_dir: str | Path) -> list[str]: """Recursively iterate over directories taking those that contain `type.raw` file. Parameters diff --git a/deepmd/dpmodel/array_api.py b/deepmd/dpmodel/array_api.py index 6b52ba7f3e..2db7550eee 100644 --- a/deepmd/dpmodel/array_api.py +++ b/deepmd/dpmodel/array_api.py @@ -1,11 +1,11 @@ # SPDX-License-Identifier: LGPL-3.0-or-later """Utilities for the array API.""" +from collections.abc import ( + Callable, +) from typing import ( Any, - Callable, - Optional, - Union, ) import array_api_compat @@ -15,7 +15,7 @@ ) # Type alias for array_api compatible arrays -Array = Union[np.ndarray, Any] # Any to support JAX, PyTorch, etc. arrays +Array = np.ndarray | Any # Any to support JAX, PyTorch, etc. arrays def support_array_api(version: str) -> Callable: @@ -125,7 +125,7 @@ def xp_add_at(x: Array, indices: Array, values: Array) -> Array: return x -def xp_bincount(x: Array, weights: Optional[Array] = None, minlength: int = 0) -> Array: +def xp_bincount(x: Array, weights: Array | None = None, minlength: int = 0) -> Array: """Counts the number of occurrences of each value in x.""" xp = array_api_compat.array_namespace(x) if array_api_compat.is_numpy_array(x) or array_api_compat.is_jax_array(x): diff --git a/deepmd/dpmodel/atomic_model/base_atomic_model.py b/deepmd/dpmodel/atomic_model/base_atomic_model.py index f9b9f0a15e..6fb48eb659 100644 --- a/deepmd/dpmodel/atomic_model/base_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/base_atomic_model.py @@ -2,7 +2,6 @@ import math from typing import ( Any, - Optional, ) import array_api_compat @@ -45,8 +44,8 @@ def __init__( type_map: list[str], atom_exclude_types: list[int] = [], pair_exclude_types: list[tuple[int, int]] = [], - rcond: Optional[float] = None, - preset_out_bias: Optional[dict[str, Array]] = None, + rcond: float | None = None, + preset_out_bias: dict[str, Array] | None = None, ) -> None: super().__init__() self.type_map = type_map @@ -133,7 +132,7 @@ def atomic_output_def(self) -> FittingOutputDef: ) def change_type_map( - self, type_map: list[str], model_with_new_type_stat: Optional[Any] = None + self, type_map: list[str], model_with_new_type_stat: Any | None = None ) -> None: """Change the type related params to new ones, according to `type_map` and the original one in the model. If there are new types in `type_map`, statistics will be updated accordingly to `model_with_new_type_stat` for these new types. @@ -154,9 +153,9 @@ def forward_common_atomic( extended_coord: Array, extended_atype: Array, nlist: Array, - mapping: Optional[Array] = None, - fparam: Optional[Array] = None, - aparam: Optional[Array] = None, + mapping: Array | None = None, + fparam: Array | None = None, + aparam: Array | None = None, ) -> dict[str, Array]: """Common interface for atomic inference. @@ -230,9 +229,9 @@ def call( extended_coord: Array, extended_atype: Array, nlist: Array, - mapping: Optional[Array] = None, - fparam: Optional[Array] = None, - aparam: Optional[Array] = None, + mapping: Array | None = None, + fparam: Array | None = None, + aparam: Array | None = None, ) -> dict[str, Array]: return self.forward_common_atomic( extended_coord, diff --git a/deepmd/dpmodel/atomic_model/dp_atomic_model.py b/deepmd/dpmodel/atomic_model/dp_atomic_model.py index 60db302667..56fc2b42c7 100644 --- a/deepmd/dpmodel/atomic_model/dp_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/dp_atomic_model.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( Any, - Optional, ) from deepmd.dpmodel.array_api import ( @@ -130,9 +129,9 @@ def forward_atomic( extended_coord: Array, extended_atype: Array, nlist: Array, - mapping: Optional[Array] = None, - fparam: Optional[Array] = None, - aparam: Optional[Array] = None, + mapping: Array | None = None, + fparam: Array | None = None, + aparam: Array | None = None, ) -> dict[str, Array]: """Models' atomic predictions. @@ -177,7 +176,7 @@ def forward_atomic( return ret def change_type_map( - self, type_map: list[str], model_with_new_type_stat: Optional[Any] = None + self, type_map: list[str], model_with_new_type_stat: Any | None = None ) -> None: """Change the type related params to new ones, according to `type_map` and the original one in the model. If there are new types in `type_map`, statistics will be updated accordingly to `model_with_new_type_stat` for these new types. diff --git a/deepmd/dpmodel/atomic_model/linear_atomic_model.py b/deepmd/dpmodel/atomic_model/linear_atomic_model.py index ed63bb2db7..3fa1119418 100644 --- a/deepmd/dpmodel/atomic_model/linear_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/linear_atomic_model.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( Any, - Optional, - Union, ) import array_api_compat @@ -115,7 +113,7 @@ def get_type_map(self) -> list[str]: return self.type_map def change_type_map( - self, type_map: list[str], model_with_new_type_stat: Optional[Any] = None + self, type_map: list[str], model_with_new_type_stat: Any | None = None ) -> None: """Change the type related params to new ones, according to `type_map` and the original one in the model. If there are new types in `type_map`, statistics will be updated accordingly to `model_with_new_type_stat` for these new types. @@ -150,7 +148,7 @@ def get_model_nsels(self) -> list[int]: """Get the processed sels for each individual models. Not distinguishing types.""" return [model.get_nsel() for model in self.models] - def get_model_sels(self) -> list[Union[int, list[int]]]: + def get_model_sels(self) -> list[int | list[int]]: """Get the sels for each individual models.""" return [model.get_sel() for model in self.models] @@ -199,9 +197,9 @@ def forward_atomic( extended_coord: Array, extended_atype: Array, nlist: Array, - mapping: Optional[Array] = None, - fparam: Optional[Array] = None, - aparam: Optional[Array] = None, + mapping: Array | None = None, + fparam: Array | None = None, + aparam: Array | None = None, ) -> dict[str, Array]: """Return atomic prediction. @@ -401,7 +399,7 @@ def __init__( sw_rmin: float, sw_rmax: float, type_map: list[str], - smin_alpha: Optional[float] = 0.1, + smin_alpha: float | None = 0.1, **kwargs: Any, ) -> None: models = [dp_model, zbl_model] diff --git a/deepmd/dpmodel/atomic_model/make_base_atomic_model.py b/deepmd/dpmodel/atomic_model/make_base_atomic_model.py index fac18c2744..3e48e88c87 100644 --- a/deepmd/dpmodel/atomic_model/make_base_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/make_base_atomic_model.py @@ -5,7 +5,6 @@ ) from typing import ( Any, - Optional, ) from deepmd.dpmodel.output_def import ( @@ -136,9 +135,9 @@ def fwd( extended_coord: t_tensor, extended_atype: t_tensor, nlist: t_tensor, - mapping: Optional[t_tensor] = None, - fparam: Optional[t_tensor] = None, - aparam: Optional[t_tensor] = None, + mapping: t_tensor | None = None, + fparam: t_tensor | None = None, + aparam: t_tensor | None = None, ) -> dict[str, t_tensor]: pass @@ -153,7 +152,7 @@ def deserialize(cls, data: dict) -> Any: @abstractmethod def change_type_map( - self, type_map: list[str], model_with_new_type_stat: Optional[Any] = None + self, type_map: list[str], model_with_new_type_stat: Any | None = None ) -> None: pass @@ -206,7 +205,7 @@ def make_atom_mask( def do_grad_r( self, - var_name: Optional[str] = None, + var_name: str | None = None, ) -> bool: """Tell if the output variable `var_name` is r_differentiable. if var_name is None, returns if any of the variable is r_differentiable. @@ -223,7 +222,7 @@ def do_grad_r( def do_grad_c( self, - var_name: Optional[str] = None, + var_name: str | None = None, ) -> bool: """Tell if the output variable `var_name` is c_differentiable. if var_name is None, returns if any of the variable is c_differentiable. diff --git a/deepmd/dpmodel/atomic_model/pairtab_atomic_model.py b/deepmd/dpmodel/atomic_model/pairtab_atomic_model.py index 54a3712912..5385d4c56c 100644 --- a/deepmd/dpmodel/atomic_model/pairtab_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/pairtab_atomic_model.py @@ -2,8 +2,6 @@ from typing import ( Any, NoReturn, - Optional, - Union, ) import array_api_compat @@ -64,10 +62,10 @@ def __init__( self, tab_file: str, rcut: float, - sel: Union[int, list[int]], + sel: int | list[int], type_map: list[str], - rcond: Optional[float] = None, - atom_ener: Optional[list[float]] = None, + rcond: float | None = None, + atom_ener: list[float] | None = None, **kwargs: Any, ) -> None: super().__init__(type_map, **kwargs) @@ -204,9 +202,9 @@ def forward_atomic( extended_coord: Array, extended_atype: Array, nlist: Array, - mapping: Optional[Array] = None, - fparam: Optional[Array] = None, - aparam: Optional[Array] = None, + mapping: Array | None = None, + fparam: Array | None = None, + aparam: Array | None = None, ) -> dict[str, Array]: xp = array_api_compat.array_namespace(extended_coord, extended_atype, nlist) nframes, nloc, nnei = nlist.shape diff --git a/deepmd/dpmodel/common.py b/deepmd/dpmodel/common.py index c1b766012c..bd6f7dac49 100644 --- a/deepmd/dpmodel/common.py +++ b/deepmd/dpmodel/common.py @@ -3,13 +3,15 @@ ABC, abstractmethod, ) +from collections.abc import ( + Callable, +) from functools import ( wraps, ) from typing import ( TYPE_CHECKING, Any, - Callable, Optional, overload, ) @@ -101,7 +103,7 @@ def __call__(self, *args: Any, **kwargs: Any) -> "Array": return self.call(*args, **kwargs) -def to_numpy_array(x: Optional["Array"]) -> Optional[np.ndarray]: +def to_numpy_array(x: Optional["Array"]) -> np.ndarray | None: """Convert an array to a NumPy array. Parameters diff --git a/deepmd/dpmodel/descriptor/descriptor.py b/deepmd/dpmodel/descriptor/descriptor.py index 417104c8c1..9b0e067972 100644 --- a/deepmd/dpmodel/descriptor/descriptor.py +++ b/deepmd/dpmodel/descriptor/descriptor.py @@ -4,12 +4,12 @@ ABC, abstractmethod, ) +from collections.abc import ( + Callable, +) from typing import ( Any, - Callable, NoReturn, - Optional, - Union, ) import numpy as np @@ -86,8 +86,8 @@ def get_dim_emb(self) -> int: def compute_input_stats( self, - merged: Union[Callable[[], list[dict]], list[dict]], - path: Optional[DPPath] = None, + merged: Callable[[], list[dict]] | list[dict], + path: DPPath | None = None, ) -> NoReturn: """ Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data. @@ -127,9 +127,9 @@ def call( nlist: Array, extended_coord: Array, extended_atype: Array, - extended_atype_embd: Optional[Array] = None, - mapping: Optional[Array] = None, - type_embedding: Optional[Array] = None, + extended_atype_embd: Array | None = None, + mapping: Array | None = None, + type_embedding: Array | None = None, ) -> Any: """Calculate DescriptorBlock.""" pass diff --git a/deepmd/dpmodel/descriptor/dpa1.py b/deepmd/dpmodel/descriptor/dpa1.py index 5fc04ddc30..5228ba55b2 100644 --- a/deepmd/dpmodel/descriptor/dpa1.py +++ b/deepmd/dpmodel/descriptor/dpa1.py @@ -1,7 +1,9 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import math -from typing import ( +from collections.abc import ( Callable, +) +from typing import ( NoReturn, Optional, Union, @@ -244,7 +246,7 @@ def __init__( self, rcut: float, rcut_smth: float, - sel: Union[list[int], int], + sel: list[int] | int, ntypes: int, neuron: list[int] = [25, 50, 100], axis_neuron: int = 8, @@ -264,18 +266,18 @@ def __init__( precision: str = DEFAULT_PRECISION, scaling_factor: float = 1.0, normalize: bool = True, - temperature: Optional[float] = None, + temperature: float | None = None, trainable_ln: bool = True, - ln_eps: Optional[float] = 1e-5, + ln_eps: float | None = 1e-5, smooth_type_embedding: bool = True, concat_output_tebd: bool = True, spin: None = None, - stripped_type_embedding: Optional[bool] = None, + stripped_type_embedding: bool | None = None, use_econf_tebd: bool = False, use_tebd_bias: bool = False, - type_map: Optional[list[str]] = None, + type_map: list[str] | None = None, # consistent with argcheck, not used though - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, ) -> None: ## seed, uniform_seed, not included. # Ensure compatibility with the deprecated stripped_type_embedding option. @@ -419,8 +421,8 @@ def dim_emb(self) -> int: def compute_input_stats( self, - merged: Union[Callable[[], list[dict]], list[dict]], - path: Optional[DPPath] = None, + merged: Callable[[], list[dict]] | list[dict], + path: DPPath | None = None, ) -> None: """ Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data. @@ -488,7 +490,7 @@ def call( coord_ext: Array, atype_ext: Array, nlist: Array, - mapping: Optional[Array] = None, + mapping: Array | None = None, ) -> Array: """Compute the descriptor. @@ -638,7 +640,7 @@ def deserialize(cls, data: dict) -> "DescrptDPA1": def update_sel( cls, train_data: DeepmdDataSystem, - type_map: Optional[list[str]], + type_map: list[str] | None, local_jdata: dict, ) -> tuple[Array, Array]: """Update the selection and perform neighbor statistics. @@ -673,7 +675,7 @@ def __init__( self, rcut: float, rcut_smth: float, - sel: Union[list[int], int], + sel: list[int] | int, ntypes: int, neuron: list[int] = [25, 50, 100], axis_neuron: int = 8, @@ -692,11 +694,11 @@ def __init__( precision: str = DEFAULT_PRECISION, scaling_factor: float = 1.0, normalize: bool = True, - temperature: Optional[float] = None, + temperature: float | None = None, trainable_ln: bool = True, - ln_eps: Optional[float] = 1e-5, + ln_eps: float | None = 1e-5, smooth: bool = True, - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, trainable: bool = True, ) -> None: self.rcut = rcut @@ -873,8 +875,8 @@ def dim_emb(self) -> int: def compute_input_stats( self, - merged: Union[Callable[[], list[dict]], list[dict]], - path: Optional[DPPath] = None, + merged: Callable[[], list[dict]] | list[dict], + path: DPPath | None = None, ) -> None: """ Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data. @@ -954,9 +956,9 @@ def call( nlist: Array, coord_ext: Array, atype_ext: Array, - atype_embd_ext: Optional[Array] = None, - mapping: Optional[Array] = None, - type_embedding: Optional[Array] = None, + atype_embd_ext: Array | None = None, + mapping: Array | None = None, + type_embedding: Array | None = None, ) -> tuple[Array, Array]: xp = array_api_compat.array_namespace(nlist, coord_ext, atype_ext) # nf x nloc x nnei x 4 @@ -1190,12 +1192,12 @@ def __init__( do_mask: bool = False, scaling_factor: float = 1.0, normalize: bool = True, - temperature: Optional[float] = None, + temperature: float | None = None, trainable_ln: bool = True, ln_eps: float = 1e-5, smooth: bool = True, precision: str = DEFAULT_PRECISION, - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, trainable: bool = True, ) -> None: """Construct a neighbor-wise attention net.""" @@ -1239,8 +1241,8 @@ def call( self, input_G: Array, nei_mask: Array, - input_r: Optional[Array] = None, - sw: Optional[Array] = None, + input_r: Array | None = None, + sw: Array | None = None, ) -> Array: out = input_G for layer in self.attention_layers: @@ -1322,12 +1324,12 @@ def __init__( do_mask: bool = False, scaling_factor: float = 1.0, normalize: bool = True, - temperature: Optional[float] = None, + temperature: float | None = None, trainable_ln: bool = True, ln_eps: float = 1e-5, smooth: bool = True, precision: str = DEFAULT_PRECISION, - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, trainable: bool = True, ) -> None: """Construct a neighbor-wise attention layer.""" @@ -1369,8 +1371,8 @@ def call( self, x: Array, nei_mask: Array, - input_r: Optional[Array] = None, - sw: Optional[Array] = None, + input_r: Array | None = None, + sw: Array | None = None, ) -> Array: residual = x x, _ = self.attention_layer(x, nei_mask, input_r=input_r, sw=sw) @@ -1431,11 +1433,11 @@ def __init__( do_mask: bool = False, scaling_factor: float = 1.0, normalize: bool = True, - temperature: Optional[float] = None, + temperature: float | None = None, bias: bool = True, smooth: bool = True, precision: str = DEFAULT_PRECISION, - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, trainable: bool = True, ) -> None: """Construct a multi-head neighbor-wise attention net.""" @@ -1482,8 +1484,8 @@ def call( self, query: Array, nei_mask: Array, - input_r: Optional[Array] = None, - sw: Optional[Array] = None, + input_r: Array | None = None, + sw: Array | None = None, attnw_shift: float = 20.0, ) -> tuple[Array, Array]: xp = array_api_compat.array_namespace(query, nei_mask) diff --git a/deepmd/dpmodel/descriptor/dpa2.py b/deepmd/dpmodel/descriptor/dpa2.py index 75bf519984..e5e02d312c 100644 --- a/deepmd/dpmodel/descriptor/dpa2.py +++ b/deepmd/dpmodel/descriptor/dpa2.py @@ -1,10 +1,10 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +from collections.abc import ( + Callable, +) from typing import ( Any, - Callable, NoReturn, - Optional, - Union, ) import array_api_compat @@ -215,7 +215,7 @@ def __init__( use_sqrt_nnei: bool = True, g1_out_conv: bool = True, g1_out_mlp: bool = True, - ln_eps: Optional[float] = 1e-5, + ln_eps: float | None = 1e-5, ) -> None: r"""The constructor for the RepformerArgs class which defines the parameters of the repformer block in DPA2 descriptor. @@ -373,9 +373,9 @@ def __init__( self, ntypes: int, # args for repinit - repinit: Union[RepinitArgs, dict], + repinit: RepinitArgs | dict, # args for repformer - repformer: Union[RepformerArgs, dict], + repformer: RepformerArgs | dict, # kwargs for descriptor concat_output_tebd: bool = True, precision: str = "float64", @@ -383,11 +383,11 @@ def __init__( exclude_types: list[tuple[int, int]] = [], env_protection: float = 0.0, trainable: bool = True, - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, add_tebd_to_repinit_out: bool = False, use_econf_tebd: bool = False, use_tebd_bias: bool = False, - type_map: Optional[list[str]] = None, + type_map: list[str] | None = None, ) -> None: r"""The DPA-2 descriptor[1]_. @@ -443,7 +443,7 @@ def __init__( Comput Mater 10, 293 (2024). https://doi.org/10.1038/s41524-024-01493-2 """ - def init_subclass_params(sub_data: Union[dict, Any], sub_class: type) -> Any: + def init_subclass_params(sub_data: dict | Any, sub_class: type) -> Any: if isinstance(sub_data, dict): return sub_class(**sub_data) elif isinstance(sub_data, sub_class): @@ -748,8 +748,8 @@ def dim_emb(self) -> int: def compute_input_stats( self, - merged: Union[Callable[[], list[dict]], list[dict]], - path: Optional[DPPath] = None, + merged: Callable[[], list[dict]] | list[dict], + path: DPPath | None = None, ) -> None: """ Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data. @@ -806,7 +806,7 @@ def call( coord_ext: Array, atype_ext: Array, nlist: Array, - mapping: Optional[Array] = None, + mapping: Array | None = None, ) -> tuple[Array, Array]: """Compute the descriptor. @@ -1073,7 +1073,7 @@ def deserialize(cls, data: dict) -> "DescrptDPA2": def update_sel( cls, train_data: DeepmdDataSystem, - type_map: Optional[list[str]], + type_map: list[str] | None, local_jdata: dict, ) -> tuple[Array, Array]: """Update the selection and perform neighbor statistics. diff --git a/deepmd/dpmodel/descriptor/dpa3.py b/deepmd/dpmodel/descriptor/dpa3.py index a54591339f..47a4fb1478 100644 --- a/deepmd/dpmodel/descriptor/dpa3.py +++ b/deepmd/dpmodel/descriptor/dpa3.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( Any, - Optional, - Union, ) import array_api_compat @@ -297,7 +295,7 @@ def __init__( self, ntypes: int, # args for repflow - repflow: Union[RepFlowArgs, dict], + repflow: RepFlowArgs | dict, # kwargs for descriptor concat_output_tebd: bool = False, activation_function: str = "silu", @@ -305,15 +303,15 @@ def __init__( exclude_types: list[tuple[int, int]] = [], env_protection: float = 0.0, trainable: bool = True, - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, use_econf_tebd: bool = False, use_tebd_bias: bool = False, use_loc_mapping: bool = True, - type_map: Optional[list[str]] = None, + type_map: list[str] | None = None, ) -> None: super().__init__() - def init_subclass_params(sub_data: Union[dict, Any], sub_class: type) -> Any: + def init_subclass_params(sub_data: dict | Any, sub_class: type) -> Any: if isinstance(sub_data, dict): return sub_class(**sub_data) elif isinstance(sub_data, sub_class): @@ -502,7 +500,7 @@ def dim_emb(self) -> int: return self.get_dim_emb() def compute_input_stats( - self, merged: list[dict], path: Optional[DPPath] = None + self, merged: list[dict], path: DPPath | None = None ) -> None: """Update mean and stddev for descriptor elements.""" descrpt_list = [self.repflows] @@ -532,7 +530,7 @@ def call( coord_ext: Array, atype_ext: Array, nlist: Array, - mapping: Optional[Array] = None, + mapping: Array | None = None, ) -> tuple[Array, Array]: """Compute the descriptor. @@ -663,7 +661,7 @@ def deserialize(cls, data: dict) -> "DescrptDPA3": def update_sel( cls, train_data: DeepmdDataSystem, - type_map: Optional[list[str]], + type_map: list[str] | None, local_jdata: dict, ) -> tuple[Array, Array]: """Update the selection and perform neighbor statistics. diff --git a/deepmd/dpmodel/descriptor/hybrid.py b/deepmd/dpmodel/descriptor/hybrid.py index 083adf4240..3682e17c4d 100644 --- a/deepmd/dpmodel/descriptor/hybrid.py +++ b/deepmd/dpmodel/descriptor/hybrid.py @@ -3,8 +3,6 @@ from typing import ( Any, NoReturn, - Optional, - Union, ) import array_api_compat @@ -46,9 +44,9 @@ class DescrptHybrid(BaseDescriptor, NativeOP): def __init__( self, - list: list[Union[BaseDescriptor, dict[str, Any]]], - type_map: Optional[list[str]] = None, - ntypes: Optional[int] = None, # to be compat with input + list: list[BaseDescriptor | dict[str, Any]], + type_map: list[str] | None = None, + ntypes: int | None = None, # to be compat with input ) -> None: super().__init__() # warning: list is conflict with built-in list @@ -196,7 +194,7 @@ def change_type_map( ) def compute_input_stats( - self, merged: list[dict], path: Optional[DPPath] = None + self, merged: list[dict], path: DPPath | None = None ) -> None: """Update mean and stddev for descriptor elements.""" for descrpt in self.descrpt_list: @@ -204,8 +202,8 @@ def compute_input_stats( def set_stat_mean_and_stddev( self, - mean: list[Union[np.ndarray, list[Array]]], - stddev: list[Union[np.ndarray, list[Array]]], + mean: list[np.ndarray | list[Array]], + stddev: list[np.ndarray | list[Array]], ) -> None: """Update mean and stddev for descriptor.""" for ii, descrpt in enumerate(self.descrpt_list): @@ -214,8 +212,8 @@ def set_stat_mean_and_stddev( def get_stat_mean_and_stddev( self, ) -> tuple[ - list[Union[Array, list[Array]]], - list[Union[Array, list[Array]]], + list[Array | list[Array]], + list[Array | list[Array]], ]: """Get mean and stddev for descriptor.""" mean_list = [] @@ -263,13 +261,13 @@ def call( coord_ext: Array, atype_ext: Array, nlist: Array, - mapping: Optional[Array] = None, + mapping: Array | None = None, ) -> tuple[ Array, - Optional[Array], - Optional[Array], - Optional[Array], - Optional[Array], + Array | None, + Array | None, + Array | None, + Array | None, ]: """Compute the descriptor. @@ -333,7 +331,7 @@ def call( def update_sel( cls, train_data: DeepmdDataSystem, - type_map: Optional[list[str]], + type_map: list[str] | None, local_jdata: dict, ) -> tuple[Array, Array]: """Update the selection and perform neighbor statistics. diff --git a/deepmd/dpmodel/descriptor/make_base_descriptor.py b/deepmd/dpmodel/descriptor/make_base_descriptor.py index e867ecdaa9..f87ca2c5b6 100644 --- a/deepmd/dpmodel/descriptor/make_base_descriptor.py +++ b/deepmd/dpmodel/descriptor/make_base_descriptor.py @@ -3,12 +3,12 @@ ABC, abstractmethod, ) +from collections.abc import ( + Callable, +) from typing import ( Any, - Callable, NoReturn, - Optional, - Union, ) from deepmd.common import ( @@ -129,7 +129,7 @@ def share_params( @abstractmethod def change_type_map( - self, type_map: list[str], model_with_new_type_stat: Optional[Any] = None + self, type_map: list[str], model_with_new_type_stat: Any | None = None ) -> None: """Change the type related params to new ones, according to `type_map` and the original one in the model. If there are new types in `type_map`, statistics will be updated accordingly to `model_with_new_type_stat` for these new types. @@ -148,8 +148,8 @@ def get_stat_mean_and_stddev(self) -> Any: def compute_input_stats( self, - merged: Union[Callable[[], list[dict]], list[dict]], - path: Optional[DPPath] = None, + merged: Callable[[], list[dict]] | list[dict], + path: DPPath | None = None, ) -> NoReturn: """Update mean and stddev for descriptor elements.""" raise NotImplementedError @@ -185,7 +185,7 @@ def fwd( extended_coord: Array, extended_atype: Array, nlist: Array, - mapping: Optional[Array] = None, + mapping: Array | None = None, ) -> Array: """Calculate descriptor.""" pass @@ -218,9 +218,9 @@ def deserialize(cls, data: dict) -> "BD": def update_sel( cls, train_data: DeepmdDataSystem, - type_map: Optional[list[str]], + type_map: list[str] | None, local_jdata: dict, - ) -> tuple[dict, Optional[float]]: + ) -> tuple[dict, float | None]: """Update the selection and perform neighbor statistics. Parameters diff --git a/deepmd/dpmodel/descriptor/repflows.py b/deepmd/dpmodel/descriptor/repflows.py index 407bf95351..706fc690e4 100644 --- a/deepmd/dpmodel/descriptor/repflows.py +++ b/deepmd/dpmodel/descriptor/repflows.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( +from collections.abc import ( Callable, - Optional, - Union, ) import array_api_compat @@ -207,7 +205,7 @@ def __init__( use_dynamic_sel: bool = False, sel_reduce_factor: float = 10.0, use_loc_mapping: bool = True, - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, trainable: bool = True, ) -> None: super().__init__() @@ -421,8 +419,8 @@ def dim_emb(self) -> int: def compute_input_stats( self, - merged: Union[Callable[[], list[dict]], list[dict]], - path: Optional[DPPath] = None, + merged: Callable[[], list[dict]] | list[dict], + path: DPPath | None = None, ) -> None: """ Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data. @@ -479,8 +477,8 @@ def call( nlist: Array, coord_ext: Array, atype_ext: Array, - atype_embd_ext: Optional[Array] = None, - mapping: Optional[Array] = None, + atype_embd_ext: Array | None = None, + mapping: Array | None = None, ) -> tuple[Array, Array]: xp = array_api_compat.array_namespace(nlist, coord_ext, atype_ext) nframes, nloc, nnei = nlist.shape @@ -874,7 +872,7 @@ def __init__( update_residual: float = 0.1, update_residual_init: str = "const", precision: str = "float64", - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, trainable: bool = True, ) -> None: super().__init__() diff --git a/deepmd/dpmodel/descriptor/repformers.py b/deepmd/dpmodel/descriptor/repformers.py index 9b5b21c1ea..79d4f9228f 100644 --- a/deepmd/dpmodel/descriptor/repformers.py +++ b/deepmd/dpmodel/descriptor/repformers.py @@ -1,9 +1,9 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +from collections.abc import ( + Callable, +) from typing import ( Any, - Callable, - Optional, - Union, ) import array_api_compat @@ -206,8 +206,8 @@ def __init__( use_sqrt_nnei: bool = True, g1_out_conv: bool = True, g1_out_mlp: bool = True, - ln_eps: Optional[float] = 1e-5, - seed: Optional[Union[int, list[int]]] = None, + ln_eps: float | None = 1e-5, + seed: int | list[int] | None = None, trainable: bool = True, ) -> None: super().__init__() @@ -383,8 +383,8 @@ def dim_emb(self) -> int: def compute_input_stats( self, - merged: Union[Callable[[], list[dict]], list[dict]], - path: Optional[DPPath] = None, + merged: Callable[[], list[dict]] | list[dict], + path: DPPath | None = None, ) -> None: """ Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data. @@ -441,9 +441,9 @@ def call( nlist: Array, coord_ext: Array, atype_ext: Array, - atype_embd_ext: Optional[Array] = None, - mapping: Optional[Array] = None, - type_embedding: Optional[Array] = None, + atype_embd_ext: Array | None = None, + mapping: Array | None = None, + type_embedding: Array | None = None, ) -> Array: xp = array_api_compat.array_namespace(nlist, coord_ext, atype_ext) exclude_mask = self.emask.build_type_exclude_mask(nlist, atype_ext) @@ -592,7 +592,7 @@ def get_residual( _mode: str = "norm", trainable: bool = True, precision: str = "float64", - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, ) -> Array: """ Get residual tensor for one update vector. @@ -856,7 +856,7 @@ def __init__( smooth: bool = True, attnw_shift: float = 20.0, precision: str = "float64", - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, trainable: bool = True, ) -> None: """Return neighbor-wise multi-head self-attention maps, with gate mechanism.""" @@ -981,7 +981,7 @@ def __init__( input_dim: int, head_num: int, precision: str = "float64", - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, trainable: bool = True, ) -> None: super().__init__() @@ -1072,7 +1072,7 @@ def __init__( input_dim: int, head_num: int, precision: str = "float64", - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, trainable: bool = True, ) -> None: super().__init__() @@ -1153,7 +1153,7 @@ def __init__( smooth: bool = True, attnw_shift: float = 20.0, precision: str = "float64", - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, trainable: bool = True, ) -> None: super().__init__() @@ -1318,8 +1318,8 @@ def __init__( use_sqrt_nnei: bool = True, g1_out_conv: bool = True, g1_out_mlp: bool = True, - ln_eps: Optional[float] = 1e-5, - seed: Optional[Union[int, list[int]]] = None, + ln_eps: float | None = 1e-5, + seed: int | list[int] | None = None, trainable: bool = True, ) -> None: super().__init__() diff --git a/deepmd/dpmodel/descriptor/se_atten_v2.py b/deepmd/dpmodel/descriptor/se_atten_v2.py index f6c497d151..99074fb652 100644 --- a/deepmd/dpmodel/descriptor/se_atten_v2.py +++ b/deepmd/dpmodel/descriptor/se_atten_v2.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( Any, - Optional, - Union, ) import numpy as np @@ -39,7 +37,7 @@ def __init__( self, rcut: float, rcut_smth: float, - sel: Union[list[int], int], + sel: list[int] | int, ntypes: int, neuron: list[int] = [25, 50, 100], axis_neuron: int = 8, @@ -58,17 +56,17 @@ def __init__( precision: str = DEFAULT_PRECISION, scaling_factor: float = 1.0, normalize: bool = True, - temperature: Optional[float] = None, + temperature: float | None = None, trainable_ln: bool = True, - ln_eps: Optional[float] = 1e-5, + ln_eps: float | None = 1e-5, concat_output_tebd: bool = True, - spin: Optional[Any] = None, - stripped_type_embedding: Optional[bool] = None, + spin: Any | None = None, + stripped_type_embedding: bool | None = None, use_econf_tebd: bool = False, use_tebd_bias: bool = False, - type_map: Optional[list[str]] = None, + type_map: list[str] | None = None, # consistent with argcheck, not used though - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, ) -> None: DescrptDPA1.__init__( self, diff --git a/deepmd/dpmodel/descriptor/se_e2_a.py b/deepmd/dpmodel/descriptor/se_e2_a.py index 7cdfa963ee..c09a6cbdc3 100644 --- a/deepmd/dpmodel/descriptor/se_e2_a.py +++ b/deepmd/dpmodel/descriptor/se_e2_a.py @@ -1,11 +1,11 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import itertools +from collections.abc import ( + Callable, +) from typing import ( Any, - Callable, NoReturn, - Optional, - Union, ) import array_api_compat @@ -165,11 +165,11 @@ def __init__( set_davg_zero: bool = False, activation_function: str = "tanh", precision: str = DEFAULT_PRECISION, - spin: Optional[Any] = None, - type_map: Optional[list[str]] = None, - ntypes: Optional[int] = None, # to be compat with input + spin: Any | None = None, + type_map: list[str] | None = None, + ntypes: int | None = None, # to be compat with input # consistent with argcheck, not used though - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, ) -> None: del ntypes ## seed, uniform_seed, not included. @@ -295,7 +295,7 @@ def share_params( raise NotImplementedError def change_type_map( - self, type_map: list[str], model_with_new_type_stat: Optional[Any] = None + self, type_map: list[str], model_with_new_type_stat: Any | None = None ) -> None: """Change the type related params to new ones, according to `type_map` and the original one in the model. If there are new types in `type_map`, statistics will be updated accordingly to `model_with_new_type_stat` for these new types. @@ -316,8 +316,8 @@ def get_type_map(self) -> list[str]: def compute_input_stats( self, - merged: Union[Callable[[], list[dict]], list[dict]], - path: Optional[DPPath] = None, + merged: Callable[[], list[dict]] | list[dict], + path: DPPath | None = None, ) -> None: """ Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data. @@ -392,7 +392,7 @@ def call( coord_ext: Array, atype_ext: Array, nlist: Array, - mapping: Optional[Array] = None, + mapping: Array | None = None, ) -> Array: """Compute the descriptor. @@ -522,7 +522,7 @@ def deserialize(cls, data: dict) -> "DescrptSeA": def update_sel( cls, train_data: DeepmdDataSystem, - type_map: Optional[list[str]], + type_map: list[str] | None, local_jdata: dict, ) -> tuple[Array, Array]: """Update the selection and perform neighbor statistics. @@ -557,7 +557,7 @@ def call( coord_ext: Array, atype_ext: Array, nlist: Array, - mapping: Optional[Array] = None, + mapping: Array | None = None, ) -> Array: """Compute the descriptor. diff --git a/deepmd/dpmodel/descriptor/se_r.py b/deepmd/dpmodel/descriptor/se_r.py index 4287083442..6decd91a23 100644 --- a/deepmd/dpmodel/descriptor/se_r.py +++ b/deepmd/dpmodel/descriptor/se_r.py @@ -1,10 +1,10 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +from collections.abc import ( + Callable, +) from typing import ( Any, - Callable, NoReturn, - Optional, - Union, ) import array_api_compat @@ -123,11 +123,11 @@ def __init__( set_davg_zero: bool = False, activation_function: str = "tanh", precision: str = DEFAULT_PRECISION, - spin: Optional[Any] = None, - type_map: Optional[list[str]] = None, - ntypes: Optional[int] = None, # to be compat with input + spin: Any | None = None, + type_map: list[str] | None = None, + ntypes: int | None = None, # to be compat with input # consistent with argcheck, not used though - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, ) -> None: del ntypes ## seed, uniform_seed, not included. @@ -254,7 +254,7 @@ def share_params( raise NotImplementedError def change_type_map( - self, type_map: list[str], model_with_new_type_stat: Optional[Any] = None + self, type_map: list[str], model_with_new_type_stat: Any | None = None ) -> None: """Change the type related params to new ones, according to `type_map` and the original one in the model. If there are new types in `type_map`, statistics will be updated accordingly to `model_with_new_type_stat` for these new types. @@ -275,8 +275,8 @@ def get_type_map(self) -> list[str]: def compute_input_stats( self, - merged: Union[Callable[[], list[dict]], list[dict]], - path: Optional[DPPath] = None, + merged: Callable[[], list[dict]] | list[dict], + path: DPPath | None = None, ) -> None: """ Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data. @@ -344,7 +344,7 @@ def call( coord_ext: Array, atype_ext: Array, nlist: Array, - mapping: Optional[Array] = None, + mapping: Array | None = None, ) -> Array: """Compute the descriptor. @@ -459,7 +459,7 @@ def deserialize(cls, data: dict) -> "DescrptSeR": def update_sel( cls, train_data: DeepmdDataSystem, - type_map: Optional[list[str]], + type_map: list[str] | None, local_jdata: dict, ) -> tuple[Array, Array]: """Update the selection and perform neighbor statistics. diff --git a/deepmd/dpmodel/descriptor/se_t.py b/deepmd/dpmodel/descriptor/se_t.py index cfeb5d7735..863187dd4c 100644 --- a/deepmd/dpmodel/descriptor/se_t.py +++ b/deepmd/dpmodel/descriptor/se_t.py @@ -1,11 +1,11 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import itertools +from collections.abc import ( + Callable, +) from typing import ( Any, - Callable, NoReturn, - Optional, - Union, ) import array_api_compat @@ -111,9 +111,9 @@ def __init__( exclude_types: list[tuple[int, int]] = [], precision: str = DEFAULT_PRECISION, trainable: bool = True, - seed: Optional[Union[int, list[int]]] = None, - type_map: Optional[list[str]] = None, - ntypes: Optional[int] = None, # to be compat with input + seed: int | list[int] | None = None, + type_map: list[str] | None = None, + ntypes: int | None = None, # to be compat with input ) -> None: del ntypes self.rcut = rcut @@ -256,8 +256,8 @@ def get_type_map(self) -> list[str]: def compute_input_stats( self, - merged: Union[Callable[[], list[dict]], list[dict]], - path: Optional[DPPath] = None, + merged: Callable[[], list[dict]] | list[dict], + path: DPPath | None = None, ) -> None: """ Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data. @@ -320,7 +320,7 @@ def call( coord_ext: Array, atype_ext: Array, nlist: Array, - mapping: Optional[Array] = None, + mapping: Array | None = None, ) -> tuple[Array, Array]: """Compute the descriptor. @@ -458,7 +458,7 @@ def deserialize(cls, data: dict) -> "DescrptSeT": def update_sel( cls, train_data: DeepmdDataSystem, - type_map: Optional[list[str]], + type_map: list[str] | None, local_jdata: dict, ) -> tuple[Array, Array]: """Update the selection and perform neighbor statistics. diff --git a/deepmd/dpmodel/descriptor/se_t_tebd.py b/deepmd/dpmodel/descriptor/se_t_tebd.py index b9e0e62531..e118d5abd4 100644 --- a/deepmd/dpmodel/descriptor/se_t_tebd.py +++ b/deepmd/dpmodel/descriptor/se_t_tebd.py @@ -1,9 +1,10 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( +from collections.abc import ( Callable, +) +from typing import ( NoReturn, Optional, - Union, ) import array_api_compat @@ -123,7 +124,7 @@ def __init__( self, rcut: float, rcut_smth: float, - sel: Union[list[int], int], + sel: list[int] | int, ntypes: int, neuron: list = [2, 4, 8], tebd_dim: int = 8, @@ -135,8 +136,8 @@ def __init__( exclude_types: list[tuple[int, int]] = [], precision: str = "float64", trainable: bool = True, - seed: Optional[Union[int, list[int]]] = None, - type_map: Optional[list[str]] = None, + seed: int | list[int] | None = None, + type_map: list[str] | None = None, concat_output_tebd: bool = True, use_econf_tebd: bool = False, use_tebd_bias: bool = False, @@ -258,8 +259,8 @@ def dim_emb(self) -> int: def compute_input_stats( self, - merged: Union[Callable[[], list[dict]], list[dict]], - path: Optional[DPPath] = None, + merged: Callable[[], list[dict]] | list[dict], + path: DPPath | None = None, ) -> None: """ Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data. @@ -327,7 +328,7 @@ def call( coord_ext: Array, atype_ext: Array, nlist: Array, - mapping: Optional[Array] = None, + mapping: Array | None = None, ) -> tuple[Array, Array]: """Compute the descriptor. @@ -456,7 +457,7 @@ def deserialize(cls, data: dict) -> "DescrptSeTTebd": def update_sel( cls, train_data: DeepmdDataSystem, - type_map: Optional[list[str]], + type_map: list[str] | None, local_jdata: dict, ) -> tuple[Array, Array]: """Update the selection and perform neighbor statistics. @@ -491,7 +492,7 @@ def __init__( self, rcut: float, rcut_smth: float, - sel: Union[list[int], int], + sel: list[int] | int, ntypes: int, neuron: list = [25, 50, 100], tebd_dim: int = 8, @@ -503,7 +504,7 @@ def __init__( exclude_types: list[tuple[int, int]] = [], env_protection: float = 0.0, smooth: bool = True, - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, trainable: bool = True, ) -> None: self.rcut = rcut @@ -659,8 +660,8 @@ def dim_emb(self) -> int: def compute_input_stats( self, - merged: Union[Callable[[], list[dict]], list[dict]], - path: Optional[DPPath] = None, + merged: Callable[[], list[dict]] | list[dict], + path: DPPath | None = None, ) -> None: """ Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data. @@ -736,9 +737,9 @@ def call( nlist: Array, coord_ext: Array, atype_ext: Array, - atype_embd_ext: Optional[Array] = None, - mapping: Optional[Array] = None, - type_embedding: Optional[Array] = None, + atype_embd_ext: Array | None = None, + mapping: Array | None = None, + type_embedding: Array | None = None, ) -> tuple[Array, Array]: xp = array_api_compat.array_namespace(nlist, coord_ext, atype_ext) # nf x nloc x nnei x 4 diff --git a/deepmd/dpmodel/fitting/dipole_fitting.py b/deepmd/dpmodel/fitting/dipole_fitting.py index e6bea408f8..2f10d7b487 100644 --- a/deepmd/dpmodel/fitting/dipole_fitting.py +++ b/deepmd/dpmodel/fitting/dipole_fitting.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( Any, - Optional, - Union, ) import array_api_compat @@ -101,21 +99,21 @@ def __init__( numb_fparam: int = 0, numb_aparam: int = 0, dim_case_embd: int = 0, - rcond: Optional[float] = None, + rcond: float | None = None, tot_ener_zero: bool = False, - trainable: Optional[list[bool]] = None, + trainable: list[bool] | None = None, activation_function: str = "tanh", precision: str = DEFAULT_PRECISION, - layer_name: Optional[list[Optional[str]]] = None, + layer_name: list[str | None] | None = None, use_aparam_as_mask: bool = False, spin: Any = None, mixed_types: bool = False, exclude_types: list[int] = [], r_differentiable: bool = True, c_differentiable: bool = True, - type_map: Optional[list[str]] = None, - seed: Optional[Union[int, list[int]]] = None, - default_fparam: Optional[list[float]] = None, + type_map: list[str] | None = None, + seed: int | list[int] | None = None, + default_fparam: list[float] | None = None, ) -> None: if tot_ener_zero: raise NotImplementedError("tot_ener_zero is not implemented") @@ -191,11 +189,11 @@ def call( self, descriptor: Array, atype: Array, - gr: Optional[Array] = None, - g2: Optional[Array] = None, - h2: Optional[Array] = None, - fparam: Optional[Array] = None, - aparam: Optional[Array] = None, + gr: Array | None = None, + g2: Array | None = None, + h2: Array | None = None, + fparam: Array | None = None, + aparam: Array | None = None, ) -> dict[str, Array]: """Calculate the fitting. diff --git a/deepmd/dpmodel/fitting/dos_fitting.py b/deepmd/dpmodel/fitting/dos_fitting.py index b444e8ae13..803f31b30f 100644 --- a/deepmd/dpmodel/fitting/dos_fitting.py +++ b/deepmd/dpmodel/fitting/dos_fitting.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( TYPE_CHECKING, - Optional, - Union, ) import numpy as np @@ -40,16 +38,16 @@ def __init__( numb_fparam: int = 0, numb_aparam: int = 0, dim_case_embd: int = 0, - bias_dos: Optional[Array] = None, - rcond: Optional[float] = None, - trainable: Union[bool, list[bool]] = True, + bias_dos: Array | None = None, + rcond: float | None = None, + trainable: bool | list[bool] = True, activation_function: str = "tanh", precision: str = DEFAULT_PRECISION, mixed_types: bool = False, exclude_types: list[int] = [], - type_map: Optional[list[str]] = None, - seed: Optional[Union[int, list[int]]] = None, - default_fparam: Optional[list] = None, + type_map: list[str] | None = None, + seed: int | list[int] | None = None, + default_fparam: list | None = None, ) -> None: if bias_dos is not None: self.bias_dos = bias_dos diff --git a/deepmd/dpmodel/fitting/ener_fitting.py b/deepmd/dpmodel/fitting/ener_fitting.py index 794c074485..713aef5117 100644 --- a/deepmd/dpmodel/fitting/ener_fitting.py +++ b/deepmd/dpmodel/fitting/ener_fitting.py @@ -2,8 +2,6 @@ from typing import ( TYPE_CHECKING, Any, - Optional, - Union, ) from deepmd.dpmodel.common import ( @@ -33,20 +31,20 @@ def __init__( numb_fparam: int = 0, numb_aparam: int = 0, dim_case_embd: int = 0, - rcond: Optional[float] = None, + rcond: float | None = None, tot_ener_zero: bool = False, - trainable: Optional[list[bool]] = None, - atom_ener: Optional[list[float]] = None, + trainable: list[bool] | None = None, + atom_ener: list[float] | None = None, activation_function: str = "tanh", precision: str = DEFAULT_PRECISION, - layer_name: Optional[list[Optional[str]]] = None, + layer_name: list[str | None] | None = None, use_aparam_as_mask: bool = False, spin: Any = None, mixed_types: bool = False, exclude_types: list[int] = [], - type_map: Optional[list[str]] = None, - seed: Optional[Union[int, list[int]]] = None, - default_fparam: Optional[list] = None, + type_map: list[str] | None = None, + seed: int | list[int] | None = None, + default_fparam: list | None = None, ) -> None: super().__init__( var_name="energy", diff --git a/deepmd/dpmodel/fitting/general_fitting.py b/deepmd/dpmodel/fitting/general_fitting.py index 10a746cbcb..c61ab234b1 100644 --- a/deepmd/dpmodel/fitting/general_fitting.py +++ b/deepmd/dpmodel/fitting/general_fitting.py @@ -2,11 +2,11 @@ from abc import ( abstractmethod, ) +from collections.abc import ( + Callable, +) from typing import ( Any, - Callable, - Optional, - Union, ) import array_api_compat @@ -113,21 +113,21 @@ def __init__( numb_fparam: int = 0, numb_aparam: int = 0, dim_case_embd: int = 0, - bias_atom_e: Optional[Array] = None, - rcond: Optional[float] = None, + bias_atom_e: Array | None = None, + rcond: float | None = None, tot_ener_zero: bool = False, - trainable: Optional[list[bool]] = None, + trainable: list[bool] | None = None, activation_function: str = "tanh", precision: str = DEFAULT_PRECISION, - layer_name: Optional[list[Optional[str]]] = None, + layer_name: list[str | None] | None = None, use_aparam_as_mask: bool = False, spin: Any = None, mixed_types: bool = True, exclude_types: list[int] = [], - remove_vaccum_contribution: Optional[list[bool]] = None, - type_map: Optional[list[str]] = None, - seed: Optional[Union[int, list[int]]] = None, - default_fparam: Optional[list[float]] = None, + remove_vaccum_contribution: list[bool] | None = None, + type_map: list[str] | None = None, + seed: int | list[int] | None = None, + default_fparam: list[float] | None = None, ) -> None: self.var_name = var_name self.ntypes = ntypes @@ -224,7 +224,7 @@ def __init__( def compute_input_stats( self, - merged: Union[Callable[[], list[dict]], list[dict]], + merged: Callable[[], list[dict]] | list[dict], protection: float = 1e-2, ) -> None: """ @@ -325,7 +325,7 @@ def set_case_embd(self, case_idx: int) -> None: self.case_embd = np.eye(self.dim_case_embd, dtype=self.prec)[case_idx] def change_type_map( - self, type_map: list[str], model_with_new_type_stat: Optional[Any] = None + self, type_map: list[str], model_with_new_type_stat: Any | None = None ) -> None: """Change the type related params to new ones, according to `type_map` and the original one in the model. If there are new types in `type_map`, statistics will be updated accordingly to `model_with_new_type_stat` for these new types. @@ -447,11 +447,11 @@ def _call_common( self, descriptor: Array, atype: Array, - gr: Optional[Array] = None, - g2: Optional[Array] = None, - h2: Optional[Array] = None, - fparam: Optional[Array] = None, - aparam: Optional[Array] = None, + gr: Array | None = None, + g2: Array | None = None, + h2: Array | None = None, + fparam: Array | None = None, + aparam: Array | None = None, ) -> dict[str, Array]: """Calculate the fitting. diff --git a/deepmd/dpmodel/fitting/invar_fitting.py b/deepmd/dpmodel/fitting/invar_fitting.py index 15ecacbf56..f771f927cd 100644 --- a/deepmd/dpmodel/fitting/invar_fitting.py +++ b/deepmd/dpmodel/fitting/invar_fitting.py @@ -2,8 +2,6 @@ from typing import ( Any, NoReturn, - Optional, - Union, ) from deepmd.dpmodel import ( @@ -128,21 +126,21 @@ def __init__( numb_fparam: int = 0, numb_aparam: int = 0, dim_case_embd: int = 0, - bias_atom: Optional[Array] = None, - rcond: Optional[float] = None, + bias_atom: Array | None = None, + rcond: float | None = None, tot_ener_zero: bool = False, - trainable: Optional[list[bool]] = None, - atom_ener: Optional[list[float]] = None, + trainable: list[bool] | None = None, + atom_ener: list[float] | None = None, activation_function: str = "tanh", precision: str = DEFAULT_PRECISION, - layer_name: Optional[list[Optional[str]]] = None, + layer_name: list[str | None] | None = None, use_aparam_as_mask: bool = False, spin: Any = None, mixed_types: bool = True, exclude_types: list[int] = [], - type_map: Optional[list[str]] = None, - seed: Optional[Union[int, list[int]]] = None, - default_fparam: Optional[list[float]] = None, + type_map: list[str] | None = None, + seed: int | list[int] | None = None, + default_fparam: list[float] | None = None, ) -> None: if tot_ener_zero: raise NotImplementedError("tot_ener_zero is not implemented") @@ -220,11 +218,11 @@ def call( self, descriptor: Array, atype: Array, - gr: Optional[Array] = None, - g2: Optional[Array] = None, - h2: Optional[Array] = None, - fparam: Optional[Array] = None, - aparam: Optional[Array] = None, + gr: Array | None = None, + g2: Array | None = None, + h2: Array | None = None, + fparam: Array | None = None, + aparam: Array | None = None, ) -> dict[str, Array]: """Calculate the fitting. diff --git a/deepmd/dpmodel/fitting/make_base_fitting.py b/deepmd/dpmodel/fitting/make_base_fitting.py index be9c5edb1f..7b65a150b2 100644 --- a/deepmd/dpmodel/fitting/make_base_fitting.py +++ b/deepmd/dpmodel/fitting/make_base_fitting.py @@ -6,7 +6,6 @@ from typing import ( Any, NoReturn, - Optional, ) from deepmd.common import ( @@ -55,11 +54,11 @@ def fwd( self, descriptor: t_tensor, atype: t_tensor, - gr: Optional[t_tensor] = None, - g2: Optional[t_tensor] = None, - h2: Optional[t_tensor] = None, - fparam: Optional[t_tensor] = None, - aparam: Optional[t_tensor] = None, + gr: t_tensor | None = None, + g2: t_tensor | None = None, + h2: t_tensor | None = None, + fparam: t_tensor | None = None, + aparam: t_tensor | None = None, ) -> dict[str, t_tensor]: """Calculate fitting.""" pass @@ -75,7 +74,7 @@ def get_type_map(self) -> list[str]: @abstractmethod def change_type_map( - self, type_map: list[str], model_with_new_type_stat: Optional[Any] = None + self, type_map: list[str], model_with_new_type_stat: Any | None = None ) -> None: """Change the type related params to new ones, according to `type_map` and the original one in the model. If there are new types in `type_map`, statistics will be updated accordingly to `model_with_new_type_stat` for these new types. diff --git a/deepmd/dpmodel/fitting/polarizability_fitting.py b/deepmd/dpmodel/fitting/polarizability_fitting.py index 04a19b394c..f3e6318ba5 100644 --- a/deepmd/dpmodel/fitting/polarizability_fitting.py +++ b/deepmd/dpmodel/fitting/polarizability_fitting.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( Any, - Optional, - Union, ) import array_api_compat @@ -108,22 +106,22 @@ def __init__( numb_fparam: int = 0, numb_aparam: int = 0, dim_case_embd: int = 0, - rcond: Optional[float] = None, + rcond: float | None = None, tot_ener_zero: bool = False, - trainable: Optional[list[bool]] = None, + trainable: list[bool] | None = None, activation_function: str = "tanh", precision: str = DEFAULT_PRECISION, - layer_name: Optional[list[Optional[str]]] = None, + layer_name: list[str | None] | None = None, use_aparam_as_mask: bool = False, spin: Any = None, mixed_types: bool = False, exclude_types: list[int] = [], fit_diag: bool = True, - scale: Optional[list[float]] = None, + scale: list[float] | None = None, shift_diag: bool = True, - type_map: Optional[list[str]] = None, - seed: Optional[Union[int, list[int]]] = None, - default_fparam: Optional[list[float]] = None, + type_map: list[str] | None = None, + seed: int | list[int] | None = None, + default_fparam: list[float] | None = None, ) -> None: if tot_ener_zero: raise NotImplementedError("tot_ener_zero is not implemented") @@ -257,11 +255,11 @@ def call( self, descriptor: Array, atype: Array, - gr: Optional[Array] = None, - g2: Optional[Array] = None, - h2: Optional[Array] = None, - fparam: Optional[Array] = None, - aparam: Optional[Array] = None, + gr: Array | None = None, + g2: Array | None = None, + h2: Array | None = None, + fparam: Array | None = None, + aparam: Array | None = None, ) -> dict[str, Array]: """Calculate the fitting. diff --git a/deepmd/dpmodel/fitting/property_fitting.py b/deepmd/dpmodel/fitting/property_fitting.py index b4e8a4d10c..4df91825b1 100644 --- a/deepmd/dpmodel/fitting/property_fitting.py +++ b/deepmd/dpmodel/fitting/property_fitting.py @@ -1,8 +1,4 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( - Optional, - Union, -) from deepmd.dpmodel.array_api import ( Array, @@ -77,9 +73,9 @@ def __init__( dim_descrpt: int, task_dim: int = 1, neuron: list[int] = [128, 128, 128], - bias_atom_p: Optional[Array] = None, - rcond: Optional[float] = None, - trainable: Union[bool, list[bool]] = True, + bias_atom_p: Array | None = None, + rcond: float | None = None, + trainable: bool | list[bool] = True, intensive: bool = False, property_name: str = "property", resnet_dt: bool = True, @@ -90,10 +86,10 @@ def __init__( precision: str = DEFAULT_PRECISION, mixed_types: bool = True, exclude_types: list[int] = [], - type_map: Optional[list[str]] = None, - default_fparam: Optional[list] = None, + type_map: list[str] | None = None, + default_fparam: list | None = None, # not used - seed: Optional[int] = None, + seed: int | None = None, ) -> None: self.task_dim = task_dim self.intensive = intensive diff --git a/deepmd/dpmodel/infer/deep_eval.py b/deepmd/dpmodel/infer/deep_eval.py index f460c6062e..605fe72d62 100644 --- a/deepmd/dpmodel/infer/deep_eval.py +++ b/deepmd/dpmodel/infer/deep_eval.py @@ -1,11 +1,12 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import json +from collections.abc import ( + Callable, +) from typing import ( TYPE_CHECKING, Any, - Callable, Optional, - Union, ) import numpy as np @@ -80,7 +81,7 @@ def __init__( model_file: str, output_def: ModelOutputDef, *args: Any, - auto_batch_size: Union[bool, int, AutoBatchSize] = True, + auto_batch_size: bool | int | AutoBatchSize = True, neighbor_list: Optional["ase.neighborlist.NewPrimitiveNeighborList"] = None, **kwargs: Any, ) -> None: @@ -168,11 +169,11 @@ def get_ntypes_spin(self) -> int: def eval( self, coords: Array, - cells: Optional[Array], + cells: Array | None, atom_types: Array, atomic: bool = False, - fparam: Optional[Array] = None, - aparam: Optional[Array] = None, + fparam: Array | None = None, + aparam: Array | None = None, **kwargs: Any, ) -> dict[str, Array]: """Evaluate the energy, force and virial by using this DP. @@ -309,10 +310,10 @@ def _get_natoms_and_nframes( def _eval_model( self, coords: Array, - cells: Optional[Array], + cells: Array | None, atom_types: Array, - fparam: Optional[Array], - aparam: Optional[Array], + fparam: Array | None, + aparam: Array | None, request_defs: list[OutputVariableDef], ) -> dict[str, Array]: model = self.dp diff --git a/deepmd/dpmodel/loss/ener.py b/deepmd/dpmodel/loss/ener.py index 55e6c90a4e..2bf103c249 100644 --- a/deepmd/dpmodel/loss/ener.py +++ b/deepmd/dpmodel/loss/ener.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( Any, - Optional, ) import array_api_compat @@ -44,7 +43,7 @@ def __init__( limit_pref_ae: float = 0.0, start_pref_pf: float = 0.0, limit_pref_pf: float = 0.0, - relative_f: Optional[float] = None, + relative_f: float | None = None, enable_atom_ener_coeff: bool = False, start_pref_gf: float = 0.0, limit_pref_gf: float = 0.0, diff --git a/deepmd/dpmodel/model/base_model.py b/deepmd/dpmodel/model/base_model.py index f7a56437a4..163cd62387 100644 --- a/deepmd/dpmodel/model/base_model.py +++ b/deepmd/dpmodel/model/base_model.py @@ -7,7 +7,6 @@ ) from typing import ( Any, - Optional, ) from deepmd.utils.data_system import ( @@ -133,7 +132,7 @@ def deserialize(cls, data: dict) -> "BaseBaseModel": model_def_script: str """The model definition script.""" - min_nbor_dist: Optional[float] + min_nbor_dist: float | None """The minimum distance between two atoms. Used for model compression. None when skipping neighbor statistics. """ @@ -143,7 +142,7 @@ def get_model_def_script(self) -> str: """Get the model definition script.""" pass - def get_min_nbor_dist(self) -> Optional[float]: + def get_min_nbor_dist(self) -> float | None: """Get the minimum distance between two atoms.""" return self.min_nbor_dist @@ -163,9 +162,9 @@ def get_nsel(self) -> int: def update_sel( cls, train_data: DeepmdDataSystem, - type_map: Optional[list[str]], + type_map: list[str] | None, local_jdata: dict, - ) -> tuple[dict, Optional[float]]: + ) -> tuple[dict, float | None]: """Update the selection and perform neighbor statistics. Parameters diff --git a/deepmd/dpmodel/model/dp_model.py b/deepmd/dpmodel/model/dp_model.py index 9098d1c011..063533f2a7 100644 --- a/deepmd/dpmodel/model/dp_model.py +++ b/deepmd/dpmodel/model/dp_model.py @@ -1,10 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( - Optional, -) - from deepmd.dpmodel.descriptor.base_descriptor import ( BaseDescriptor, ) @@ -22,9 +18,9 @@ class DPModelCommon: def update_sel( cls, train_data: DeepmdDataSystem, - type_map: Optional[list[str]], + type_map: list[str] | None, local_jdata: dict, - ) -> tuple[dict, Optional[float]]: + ) -> tuple[dict, float | None]: """Update the selection and perform neighbor statistics. Parameters diff --git a/deepmd/dpmodel/model/dp_zbl_model.py b/deepmd/dpmodel/model/dp_zbl_model.py index f3f106f1c7..b5940f4707 100644 --- a/deepmd/dpmodel/model/dp_zbl_model.py +++ b/deepmd/dpmodel/model/dp_zbl_model.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( Any, - Optional, ) from deepmd.dpmodel.atomic_model.linear_atomic_model import ( @@ -39,9 +38,9 @@ def __init__( def update_sel( cls, train_data: DeepmdDataSystem, - type_map: Optional[list[str]], + type_map: list[str] | None, local_jdata: dict, - ) -> tuple[dict, Optional[float]]: + ) -> tuple[dict, float | None]: """Update the selection and perform neighbor statistics. Parameters diff --git a/deepmd/dpmodel/model/make_model.py b/deepmd/dpmodel/model/make_model.py index 74d5dfd4bb..3a88aac1e4 100644 --- a/deepmd/dpmodel/model/make_model.py +++ b/deepmd/dpmodel/model/make_model.py @@ -1,8 +1,9 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +from collections.abc import ( + Callable, +) from typing import ( Any, - Callable, - Optional, ) import array_api_compat @@ -51,8 +52,8 @@ def model_call_from_call_lower( np.ndarray, np.ndarray, np.ndarray, - Optional[np.ndarray], - Optional[np.ndarray], + np.ndarray | None, + np.ndarray | None, bool, ], dict[str, Array], @@ -63,9 +64,9 @@ def model_call_from_call_lower( model_output_def: ModelOutputDef, coord: Array, atype: Array, - box: Optional[Array] = None, - fparam: Optional[Array] = None, - aparam: Optional[Array] = None, + box: Array | None = None, + fparam: Array | None = None, + aparam: Array | None = None, do_atomic_virial: bool = False, ) -> dict[str, Array]: """Return model prediction from lower interface. @@ -163,7 +164,7 @@ def __init__( self, *args: Any, # underscore to prevent conflict with normal inputs - atomic_model_: Optional[T_AtomicModel] = None, + atomic_model_: T_AtomicModel | None = None, **kwargs: Any, ) -> None: BaseModel.__init__(self) @@ -224,9 +225,9 @@ def call( self, coord: Array, atype: Array, - box: Optional[Array] = None, - fparam: Optional[Array] = None, - aparam: Optional[Array] = None, + box: Array | None = None, + fparam: Array | None = None, + aparam: Array | None = None, do_atomic_virial: bool = False, ) -> dict[str, Array]: """Return model prediction. @@ -279,9 +280,9 @@ def call_lower( extended_coord: Array, extended_atype: Array, nlist: Array, - mapping: Optional[Array] = None, - fparam: Optional[Array] = None, - aparam: Optional[Array] = None, + mapping: Array | None = None, + fparam: Array | None = None, + aparam: Array | None = None, do_atomic_virial: bool = False, ) -> dict[str, Array]: """Return model prediction. Lower interface that takes @@ -341,9 +342,9 @@ def forward_common_atomic( extended_coord: Array, extended_atype: Array, nlist: Array, - mapping: Optional[Array] = None, - fparam: Optional[Array] = None, - aparam: Optional[Array] = None, + mapping: Array | None = None, + fparam: Array | None = None, + aparam: Array | None = None, do_atomic_virial: bool = False, ) -> dict[str, Array]: atomic_ret = self.atomic_model.forward_common_atomic( @@ -367,16 +368,16 @@ def forward_common_atomic( def input_type_cast( self, coord: Array, - box: Optional[Array] = None, - fparam: Optional[Array] = None, - aparam: Optional[Array] = None, - ) -> tuple[Array, Array, Optional[np.ndarray], Optional[np.ndarray], str]: + box: Array | None = None, + fparam: Array | None = None, + aparam: Array | None = None, + ) -> tuple[Array, Array, np.ndarray | None, np.ndarray | None, str]: """Cast the input data to global float type.""" input_prec = RESERVED_PRECISION_DICT[self.precision_dict[coord.dtype.name]] ### ### type checking would not pass jit, convert to coord prec anyway ### - _lst: list[Optional[np.ndarray]] = [ + _lst: list[np.ndarray | None] = [ vv.astype(coord.dtype) if vv is not None else None for vv in [box, fparam, aparam] ] @@ -520,7 +521,7 @@ def _format_nlist( def do_grad_r( self, - var_name: Optional[str] = None, + var_name: str | None = None, ) -> bool: """Tell if the output variable `var_name` is r_differentiable. if var_name is None, returns if any of the variable is r_differentiable. @@ -529,7 +530,7 @@ def do_grad_r( def do_grad_c( self, - var_name: Optional[str] = None, + var_name: str | None = None, ) -> bool: """Tell if the output variable `var_name` is c_differentiable. if var_name is None, returns if any of the variable is c_differentiable. diff --git a/deepmd/dpmodel/model/spin_model.py b/deepmd/dpmodel/model/spin_model.py index 7706a009fc..521978bdde 100644 --- a/deepmd/dpmodel/model/spin_model.py +++ b/deepmd/dpmodel/model/spin_model.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( Any, - Optional, ) import numpy as np @@ -59,7 +58,7 @@ def process_spin_input_lower( extended_atype: Array, extended_spin: Array, nlist: Array, - mapping: Optional[Array] = None, + mapping: Array | None = None, ) -> tuple[Array, Array]: """ Add `extended_spin` into `extended_coord` to generate virtual atoms, and extend `nlist` and `mapping`. @@ -276,7 +275,7 @@ def get_model_def_script(self) -> str: """Get the model definition script.""" return self.backbone_model.get_model_def_script() - def get_min_nbor_dist(self) -> Optional[float]: + def get_min_nbor_dist(self) -> float | None: """Get the minimum neighbor distance.""" return self.backbone_model.get_min_nbor_dist() @@ -339,9 +338,9 @@ def call( coord: Array, atype: Array, spin: Array, - box: Optional[Array] = None, - fparam: Optional[Array] = None, - aparam: Optional[Array] = None, + box: Array | None = None, + fparam: Array | None = None, + aparam: Array | None = None, do_atomic_virial: bool = False, ) -> dict[str, Array]: """Return model prediction. @@ -402,9 +401,9 @@ def call_lower( extended_atype: Array, extended_spin: Array, nlist: Array, - mapping: Optional[Array] = None, - fparam: Optional[Array] = None, - aparam: Optional[Array] = None, + mapping: Array | None = None, + fparam: Array | None = None, + aparam: Array | None = None, do_atomic_virial: bool = False, ) -> dict[str, Array]: """Return model prediction. Lower interface that takes diff --git a/deepmd/dpmodel/model/transform_output.py b/deepmd/dpmodel/model/transform_output.py index f35faf444e..d3315eda55 100644 --- a/deepmd/dpmodel/model/transform_output.py +++ b/deepmd/dpmodel/model/transform_output.py @@ -1,8 +1,5 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( - Optional, -) import array_api_compat import numpy as np @@ -29,7 +26,7 @@ def fit_output_to_model_output( fit_output_def: FittingOutputDef, coord_ext: Array, do_atomic_virial: bool = False, - mask: Optional[Array] = None, + mask: Array | None = None, ) -> dict[str, Array]: """Transform the output of the fitting network to the model output. diff --git a/deepmd/dpmodel/utils/env_mat.py b/deepmd/dpmodel/utils/env_mat.py index 2302e24c71..f2943d5e44 100644 --- a/deepmd/dpmodel/utils/env_mat.py +++ b/deepmd/dpmodel/utils/env_mat.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( Any, - Optional, ) import array_api_compat @@ -117,8 +116,8 @@ def call( coord_ext: Array, atype_ext: Array, nlist: Array, - davg: Optional[Array] = None, - dstd: Optional[Array] = None, + davg: Array | None = None, + dstd: Array | None = None, radial_only: bool = False, ) -> tuple[Array, Array, Array]: """Compute the environment matrix. diff --git a/deepmd/dpmodel/utils/env_mat_stat.py b/deepmd/dpmodel/utils/env_mat_stat.py index a26a99f2c2..b8befa0087 100644 --- a/deepmd/dpmodel/utils/env_mat_stat.py +++ b/deepmd/dpmodel/utils/env_mat_stat.py @@ -82,7 +82,7 @@ def __init__(self, descriptor: Union["Descriptor", "DescriptorBlock"]) -> None: ) # se_r=1, se_a=4 def iter( - self, data: list[dict[str, Union[np.ndarray, list[tuple[int, int]]]]] + self, data: list[dict[str, np.ndarray | list[tuple[int, int]]]] ) -> Iterator[dict[str, StatItem]]: """Get the iterator of the environment matrix. diff --git a/deepmd/dpmodel/utils/learning_rate.py b/deepmd/dpmodel/utils/learning_rate.py index 499c068a93..10f7ec8d04 100644 --- a/deepmd/dpmodel/utils/learning_rate.py +++ b/deepmd/dpmodel/utils/learning_rate.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( Any, - Optional, ) import numpy as np @@ -14,7 +13,7 @@ def __init__( stop_lr: float, decay_steps: int, stop_steps: int, - decay_rate: Optional[float] = None, + decay_rate: float | None = None, **kwargs: Any, ) -> None: """ diff --git a/deepmd/dpmodel/utils/neighbor_stat.py b/deepmd/dpmodel/utils/neighbor_stat.py index 289e047cf2..1bcc894624 100644 --- a/deepmd/dpmodel/utils/neighbor_stat.py +++ b/deepmd/dpmodel/utils/neighbor_stat.py @@ -2,9 +2,6 @@ from collections.abc import ( Iterator, ) -from typing import ( - Optional, -) import array_api_compat import numpy as np @@ -51,7 +48,7 @@ def call( self, coord: Array, atype: Array, - cell: Optional[Array], + cell: Array | None, ) -> tuple[Array, Array]: """Calculate the neareest neighbor distance between atoms, maximum nbor size of atoms and the output data range of the environment matrix. diff --git a/deepmd/dpmodel/utils/network.py b/deepmd/dpmodel/utils/network.py index d48c42ad08..e712adfdd8 100644 --- a/deepmd/dpmodel/utils/network.py +++ b/deepmd/dpmodel/utils/network.py @@ -5,12 +5,12 @@ """ import itertools +from collections.abc import ( + Callable, +) from typing import ( Any, - Callable, ClassVar, - Optional, - Union, ) import array_api_compat @@ -99,10 +99,10 @@ def __init__( num_out: int, bias: bool = True, use_timestep: bool = False, - activation_function: Optional[str] = None, + activation_function: str | None = None, resnet: bool = False, precision: str = DEFAULT_PRECISION, - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, trainable: bool = True, ) -> None: # trainable must be set before any array attribute is set @@ -207,7 +207,7 @@ def check_shape_consistency(self) -> None: def check_type_consistency(self) -> None: precision = self.precision - def check_var(var: Optional[Array]) -> None: + def check_var(var: Array | None) -> None: if var is not None: # array api standard doesn't provide a API to get the dtype name # this is really hacked @@ -425,7 +425,7 @@ def __init__( uni_init: bool = True, trainable: bool = True, precision: str = DEFAULT_PRECISION, - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, ) -> None: self.eps = eps self.uni_init = uni_init @@ -581,7 +581,7 @@ class NN(ModuleBase): The layers of the network. """ - def __init__(self, layers: Optional[list[dict]] = None) -> None: + def __init__(self, layers: list[dict] | None = None) -> None: super().__init__() if layers is None: layers = [] @@ -715,9 +715,9 @@ def __init__( activation_function: str = "tanh", resnet_dt: bool = False, precision: str = DEFAULT_PRECISION, - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, bias: bool = True, - trainable: Union[bool, list[bool]] = True, + trainable: bool | list[bool] = True, ) -> None: layers = [] i_in = in_dim @@ -827,8 +827,8 @@ def __init__( resnet_dt: bool = False, precision: str = DEFAULT_PRECISION, bias_out: bool = True, - seed: Optional[Union[int, list[int]]] = None, - trainable: Union[bool, list[bool]] = True, + seed: int | list[int] | None = None, + trainable: bool | list[bool] = True, ) -> None: if trainable is None: trainable = [True] * (len(neuron) + 1) @@ -937,7 +937,7 @@ def __init__( ndim: int, ntypes: int, network_type: str = "network", - networks: list[Union[NativeNet, dict]] = [], + networks: list[NativeNet | dict] = [], ) -> None: self.ndim = ndim self.ntypes = ntypes @@ -960,7 +960,7 @@ def check_completeness(self) -> None: if self[tuple(tt)] is None: raise RuntimeError(f"network for {tt} not found") - def _convert_key(self, key: Union[int, tuple]) -> int: + def _convert_key(self, key: int | tuple) -> int: if isinstance(key, int): idx = key else: @@ -975,10 +975,10 @@ def _convert_key(self, key: Union[int, tuple]) -> int: idx = sum([tt * self.ntypes**ii for ii, tt in enumerate(key)]) return idx - def __getitem__(self, key: Union[int, tuple]) -> Any: + def __getitem__(self, key: int | tuple) -> Any: return self._networks[self._convert_key(key)] - def __setitem__(self, key: Union[int, tuple], value: Any) -> None: + def __setitem__(self, key: int | tuple, value: Any) -> None: if value is None: pass elif isinstance(value, self.network_type): @@ -1029,7 +1029,7 @@ def aggregate( # noqa: ANN201 data, # noqa: ANN001 owners, # noqa: ANN001 average: bool = True, - num_owner: Optional[int] = None, + num_owner: int | None = None, ): """ Aggregate rows in data by specifying the owners. diff --git a/deepmd/dpmodel/utils/nlist.py b/deepmd/dpmodel/utils/nlist.py index 86b1353485..a43cf46403 100644 --- a/deepmd/dpmodel/utils/nlist.py +++ b/deepmd/dpmodel/utils/nlist.py @@ -1,8 +1,4 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( - Optional, - Union, -) import array_api_compat @@ -23,7 +19,7 @@ def extend_input_and_build_neighbor_list( rcut: float, sel: list[int], mixed_types: bool = False, - box: Optional[Array] = None, + box: Array | None = None, ) -> tuple[Array, Array]: xp = array_api_compat.array_namespace(coord, atype) nframes, nloc = atype.shape[:2] @@ -55,7 +51,7 @@ def build_neighbor_list( atype: Array, nloc: int, rcut: float, - sel: Union[int, list[int]], + sel: int | list[int], distinguish_types: bool = True, ) -> Array: """Build neighbor list for a single frame. keeps nsel neighbors. @@ -249,7 +245,7 @@ def build_multiple_neighbor_list( def extend_coord_with_ghosts( coord: Array, atype: Array, - cell: Optional[Array], + cell: Array | None, rcut: float, ) -> tuple[Array, Array]: """Extend the coordinates of the atoms by appending peridoc images. diff --git a/deepmd/dpmodel/utils/safe_gradient.py b/deepmd/dpmodel/utils/safe_gradient.py index 08ffa9bb10..e58ff1d613 100644 --- a/deepmd/dpmodel/utils/safe_gradient.py +++ b/deepmd/dpmodel/utils/safe_gradient.py @@ -7,7 +7,6 @@ from typing import ( Any, - Optional, ) import array_api_compat @@ -21,7 +20,7 @@ def safe_for_sqrt(x: Any) -> Any: def safe_for_vector_norm( - x: Any, /, *, axis: Optional[Any] = None, keepdims: bool = False, ord: Any = 2 + x: Any, /, *, axis: Any | None = None, keepdims: bool = False, ord: Any = 2 ) -> Any: """Safe version of sqrt that has a gradient of 0 at x = 0.""" xp = array_api_compat.array_namespace(x) diff --git a/deepmd/dpmodel/utils/seed.py b/deepmd/dpmodel/utils/seed.py index 165ff558b9..1603cc0164 100644 --- a/deepmd/dpmodel/utils/seed.py +++ b/deepmd/dpmodel/utils/seed.py @@ -1,7 +1,5 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( - Optional, - Union, overload, ) @@ -11,10 +9,10 @@ def child_seed(seed: None, idx: int) -> None: ... @overload -def child_seed(seed: Union[int, list[int]], idx: int) -> list[int]: ... +def child_seed(seed: int | list[int], idx: int) -> list[int]: ... -def child_seed(seed: Optional[Union[int, list[int]]], idx: int) -> Optional[list[int]]: +def child_seed(seed: int | list[int] | None, idx: int) -> list[int] | None: """Generate a child seed from a parent seed. Parameters diff --git a/deepmd/dpmodel/utils/serialization.py b/deepmd/dpmodel/utils/serialization.py index b765e2eca3..74d31ea589 100644 --- a/deepmd/dpmodel/utils/serialization.py +++ b/deepmd/dpmodel/utils/serialization.py @@ -1,12 +1,14 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import datetime import json +from collections.abc import ( + Callable, +) from pathlib import ( Path, ) from typing import ( Any, - Callable, ) import h5py diff --git a/deepmd/dpmodel/utils/type_embed.py b/deepmd/dpmodel/utils/type_embed.py index 33c70c5763..a1b698b698 100644 --- a/deepmd/dpmodel/utils/type_embed.py +++ b/deepmd/dpmodel/utils/type_embed.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( Any, - Optional, - Union, ) import array_api_compat @@ -65,11 +63,11 @@ def __init__( activation_function: str = "tanh", precision: str = "default", trainable: bool = True, - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, padding: bool = False, use_econf_tebd: bool = False, use_tebd_bias: bool = False, - type_map: Optional[list[str]] = None, + type_map: list[str] | None = None, ) -> None: self.ntypes = ntypes self.neuron = neuron diff --git a/deepmd/entrypoints/eval_desc.py b/deepmd/entrypoints/eval_desc.py index da47b6d065..a918b63edb 100644 --- a/deepmd/entrypoints/eval_desc.py +++ b/deepmd/entrypoints/eval_desc.py @@ -8,7 +8,6 @@ ) from typing import ( Any, - Optional, ) import numpy as np @@ -34,7 +33,7 @@ def eval_desc( system: str, datafile: str, output: str = "desc", - head: Optional[str] = None, + head: str | None = None, **kwargs: Any, ) -> None: """Evaluate descriptors for given systems. diff --git a/deepmd/entrypoints/neighbor_stat.py b/deepmd/entrypoints/neighbor_stat.py index d492b072ad..836f1e0174 100644 --- a/deepmd/entrypoints/neighbor_stat.py +++ b/deepmd/entrypoints/neighbor_stat.py @@ -2,7 +2,6 @@ import logging from typing import ( Any, - Optional, ) from deepmd.backend.backend import ( @@ -22,7 +21,7 @@ def neighbor_stat( *, system: str, rcut: float, - type_map: Optional[list[str]], + type_map: list[str] | None, mixed_type: bool = False, backend: str = "tensorflow", **kwargs: Any, diff --git a/deepmd/entrypoints/test.py b/deepmd/entrypoints/test.py index 16c097f743..4a0cb27cb1 100644 --- a/deepmd/entrypoints/test.py +++ b/deepmd/entrypoints/test.py @@ -8,7 +8,6 @@ from typing import ( TYPE_CHECKING, Any, - Optional, ) import numpy as np @@ -66,16 +65,16 @@ def test( *, model: str, - system: Optional[str], - datafile: Optional[str], - train_json: Optional[str] = None, - valid_json: Optional[str] = None, + system: str | None, + datafile: str | None, + train_json: str | None = None, + valid_json: str | None = None, numb_test: int, - rand_seed: Optional[int], + rand_seed: int | None, shuffle_test: bool, detail_file: str, atomic: bool, - head: Optional[str] = None, + head: str | None = None, **kwargs: Any, ) -> None: """Test model predictions. @@ -303,7 +302,7 @@ def test_ener( data: DeepmdData, system: str, numb_test: int, - detail_file: Optional[str], + detail_file: str | None, has_atom_ener: bool, append_detail: bool = False, ) -> tuple[list[np.ndarray], list[int]]: @@ -688,7 +687,7 @@ def test_dos( data: DeepmdData, system: str, numb_test: int, - detail_file: Optional[str], + detail_file: str | None, has_atom_dos: bool, append_detail: bool = False, ) -> tuple[list[np.ndarray], list[int]]: @@ -848,7 +847,7 @@ def test_property( data: DeepmdData, system: str, numb_test: int, - detail_file: Optional[str], + detail_file: str | None, has_atom_property: bool, append_detail: bool = False, ) -> tuple[list[np.ndarray], list[int]]: @@ -1042,7 +1041,7 @@ def test_wfc( dp: "DeepWFC", data: DeepmdData, numb_test: int, - detail_file: Optional[str], + detail_file: str | None, ) -> tuple[list[np.ndarray], list[int]]: """Test energy type model. @@ -1104,7 +1103,7 @@ def test_polar( dp: "DeepPolar", data: DeepmdData, numb_test: int, - detail_file: Optional[str], + detail_file: str | None, *, atomic: bool, ) -> tuple[list[np.ndarray], list[int]]: @@ -1246,7 +1245,7 @@ def test_dipole( dp: "DeepDipole", data: DeepmdData, numb_test: int, - detail_file: Optional[str], + detail_file: str | None, atomic: bool, ) -> tuple[list[np.ndarray], list[int]]: """Test energy type model. diff --git a/deepmd/infer/deep_dos.py b/deepmd/infer/deep_dos.py index 0d7ccee2b6..4e63a156c1 100644 --- a/deepmd/infer/deep_dos.py +++ b/deepmd/infer/deep_dos.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( Any, - Optional, - Union, ) import numpy as np @@ -61,11 +59,11 @@ def numb_dos(self) -> int: def eval( self, coords: np.ndarray, - cells: Optional[np.ndarray], - atom_types: Union[list[int], np.ndarray], + cells: np.ndarray | None, + atom_types: list[int] | np.ndarray, atomic: bool = False, - fparam: Optional[np.ndarray] = None, - aparam: Optional[np.ndarray] = None, + fparam: np.ndarray | None = None, + aparam: np.ndarray | None = None, mixed_type: bool = False, **kwargs: Any, ) -> tuple[np.ndarray, ...]: diff --git a/deepmd/infer/deep_eval.py b/deepmd/infer/deep_eval.py index 5f29f08330..0a8686a98b 100644 --- a/deepmd/infer/deep_eval.py +++ b/deepmd/infer/deep_eval.py @@ -8,7 +8,6 @@ Any, ClassVar, Optional, - Union, ) import numpy as np @@ -84,7 +83,7 @@ def __init__( model_file: str, output_def: ModelOutputDef, *args: Any, - auto_batch_size: Union[bool, int, AutoBatchSize] = True, + auto_batch_size: bool | int | AutoBatchSize = True, neighbor_list: Optional["ase.neighborlist.NewPrimitiveNeighborList"] = None, **kwargs: Any, ) -> None: @@ -102,11 +101,11 @@ def __new__( def eval( self, coords: np.ndarray, - cells: Optional[np.ndarray], + cells: np.ndarray | None, atom_types: np.ndarray, atomic: bool = False, - fparam: Optional[np.ndarray] = None, - aparam: Optional[np.ndarray] = None, + fparam: np.ndarray | None = None, + aparam: np.ndarray | None = None, **kwargs: Any, ) -> dict[str, np.ndarray]: """Evaluate the energy, force and virial by using this DP. @@ -173,11 +172,11 @@ def get_dim_aparam(self) -> int: def eval_descriptor( self, coords: np.ndarray, - cells: Optional[np.ndarray], + cells: np.ndarray | None, atom_types: np.ndarray, - fparam: Optional[np.ndarray] = None, - aparam: Optional[np.ndarray] = None, - efield: Optional[np.ndarray] = None, + fparam: np.ndarray | None = None, + aparam: np.ndarray | None = None, + efield: np.ndarray | None = None, mixed_type: bool = False, **kwargs: Any, ) -> np.ndarray: @@ -224,10 +223,10 @@ def eval_descriptor( def eval_fitting_last_layer( self, coords: np.ndarray, - cells: Optional[np.ndarray], + cells: np.ndarray | None, atom_types: np.ndarray, - fparam: Optional[np.ndarray] = None, - aparam: Optional[np.ndarray] = None, + fparam: np.ndarray | None = None, + aparam: np.ndarray | None = None, **kwargs: Any, ) -> np.ndarray: """Evaluate fitting before last layer by using this DP. @@ -400,7 +399,7 @@ def __init__( self, model_file: str, *args: Any, - auto_batch_size: Union[bool, int, AutoBatchSize] = True, + auto_batch_size: bool | int | AutoBatchSize = True, neighbor_list: Optional["ase.neighborlist.NewPrimitiveNeighborList"] = None, **kwargs: Any, ) -> None: @@ -471,10 +470,10 @@ def _expande_atype( def eval_descriptor( self, coords: np.ndarray, - cells: Optional[np.ndarray], + cells: np.ndarray | None, atom_types: np.ndarray, - fparam: Optional[np.ndarray] = None, - aparam: Optional[np.ndarray] = None, + fparam: np.ndarray | None = None, + aparam: np.ndarray | None = None, mixed_type: bool = False, **kwargs: Any, ) -> np.ndarray: @@ -538,10 +537,10 @@ def eval_descriptor( def eval_fitting_last_layer( self, coords: np.ndarray, - cells: Optional[np.ndarray], + cells: np.ndarray | None, atom_types: np.ndarray, - fparam: Optional[np.ndarray] = None, - aparam: Optional[np.ndarray] = None, + fparam: np.ndarray | None = None, + aparam: np.ndarray | None = None, mixed_type: bool = False, **kwargs: Any, ) -> np.ndarray: @@ -633,18 +632,18 @@ def eval_typeebd(self) -> np.ndarray: def _standard_input( self, - coords: Union[np.ndarray, list], - cells: Optional[Union[np.ndarray, list]], - atom_types: Union[np.ndarray, list], - fparam: Optional[Union[np.ndarray, list]], - aparam: Optional[Union[np.ndarray, list]], + coords: np.ndarray | list, + cells: np.ndarray | list | None, + atom_types: np.ndarray | list, + fparam: np.ndarray | list | None, + aparam: np.ndarray | list | None, mixed_type: bool, ) -> tuple[ np.ndarray, - Optional[np.ndarray], + np.ndarray | None, np.ndarray, - Optional[np.ndarray], - Optional[np.ndarray], + np.ndarray | None, + np.ndarray | None, ]: coords = np.array(coords) if cells is not None: diff --git a/deepmd/infer/deep_polar.py b/deepmd/infer/deep_polar.py index 52ffd78b26..9b74f43270 100644 --- a/deepmd/infer/deep_polar.py +++ b/deepmd/infer/deep_polar.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( Any, - Optional, - Union, ) import numpy as np @@ -50,11 +48,11 @@ def output_tensor_name(self) -> str: def eval( self, coords: np.ndarray, - cells: Optional[np.ndarray], - atom_types: Union[list[int], np.ndarray], + cells: np.ndarray | None, + atom_types: list[int] | np.ndarray, atomic: bool = False, - fparam: Optional[np.ndarray] = None, - aparam: Optional[np.ndarray] = None, + fparam: np.ndarray | None = None, + aparam: np.ndarray | None = None, mixed_type: bool = False, **kwargs: Any, ) -> np.ndarray: diff --git a/deepmd/infer/deep_pot.py b/deepmd/infer/deep_pot.py index 6e00a30f91..ce23c5981b 100644 --- a/deepmd/infer/deep_pot.py +++ b/deepmd/infer/deep_pot.py @@ -2,8 +2,6 @@ from typing import ( Any, Literal, - Optional, - Union, overload, ) @@ -93,45 +91,45 @@ def output_def_mag(self) -> ModelOutputDef: def eval( self, coords: np.ndarray, - cells: Optional[np.ndarray], - atom_types: Union[list[int], np.ndarray], + cells: np.ndarray | None, + atom_types: list[int] | np.ndarray, atomic: Literal[True], - fparam: Optional[np.ndarray], - aparam: Optional[np.ndarray], + fparam: np.ndarray | None, + aparam: np.ndarray | None, mixed_type: bool, **kwargs: Any, - ) -> Union[ - tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray], - tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray], - ]: + ) -> ( + tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray] + | tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray] + ): pass @overload def eval( self, coords: np.ndarray, - cells: Optional[np.ndarray], - atom_types: Union[list[int], np.ndarray], + cells: np.ndarray | None, + atom_types: list[int] | np.ndarray, atomic: Literal[False], - fparam: Optional[np.ndarray], - aparam: Optional[np.ndarray], + fparam: np.ndarray | None, + aparam: np.ndarray | None, mixed_type: bool, **kwargs: Any, - ) -> Union[ - tuple[np.ndarray, np.ndarray, np.ndarray], - tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray], - ]: + ) -> ( + tuple[np.ndarray, np.ndarray, np.ndarray] + | tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray] + ): pass @overload def eval( self, coords: np.ndarray, - cells: Optional[np.ndarray], - atom_types: Union[list[int], np.ndarray], + cells: np.ndarray | None, + atom_types: list[int] | np.ndarray, atomic: bool, - fparam: Optional[np.ndarray], - aparam: Optional[np.ndarray], + fparam: np.ndarray | None, + aparam: np.ndarray | None, mixed_type: bool, **kwargs: Any, ) -> tuple[np.ndarray, ...]: @@ -140,11 +138,11 @@ def eval( def eval( self, coords: np.ndarray, - cells: Optional[np.ndarray], - atom_types: Union[list[int], np.ndarray], + cells: np.ndarray | None, + atom_types: list[int] | np.ndarray, atomic: bool = False, - fparam: Optional[np.ndarray] = None, - aparam: Optional[np.ndarray] = None, + fparam: np.ndarray | None = None, + aparam: np.ndarray | None = None, mixed_type: bool = False, **kwargs: Any, ) -> tuple[np.ndarray, ...]: diff --git a/deepmd/infer/deep_property.py b/deepmd/infer/deep_property.py index 5944491cc0..5e35dcd781 100644 --- a/deepmd/infer/deep_property.py +++ b/deepmd/infer/deep_property.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( Any, - Optional, - Union, ) import numpy as np @@ -81,11 +79,11 @@ def task_dim(self) -> int: def eval( self, coords: np.ndarray, - cells: Optional[np.ndarray], - atom_types: Union[list[int], np.ndarray], + cells: np.ndarray | None, + atom_types: list[int] | np.ndarray, atomic: bool = False, - fparam: Optional[np.ndarray] = None, - aparam: Optional[np.ndarray] = None, + fparam: np.ndarray | None = None, + aparam: np.ndarray | None = None, mixed_type: bool = False, **kwargs: dict[str, Any], ) -> tuple[np.ndarray, ...]: diff --git a/deepmd/infer/deep_tensor.py b/deepmd/infer/deep_tensor.py index bb5bc12697..877ade6e19 100644 --- a/deepmd/infer/deep_tensor.py +++ b/deepmd/infer/deep_tensor.py @@ -2,10 +2,6 @@ from abc import ( abstractmethod, ) -from typing import ( - Optional, - Union, -) import numpy as np @@ -41,11 +37,11 @@ class DeepTensor(DeepEval): def eval( self, coords: np.ndarray, - cells: Optional[np.ndarray], - atom_types: Union[list[int], np.ndarray], + cells: np.ndarray | None, + atom_types: list[int] | np.ndarray, atomic: bool = True, - fparam: Optional[np.ndarray] = None, - aparam: Optional[np.ndarray] = None, + fparam: np.ndarray | None = None, + aparam: np.ndarray | None = None, mixed_type: bool = False, **kwargs: dict, ) -> np.ndarray: @@ -110,11 +106,11 @@ def eval( def eval_full( self, coords: np.ndarray, - cells: Optional[np.ndarray], + cells: np.ndarray | None, atom_types: np.ndarray, atomic: bool = False, - fparam: Optional[np.ndarray] = None, - aparam: Optional[np.ndarray] = None, + fparam: np.ndarray | None = None, + aparam: np.ndarray | None = None, mixed_type: bool = False, **kwargs: dict, ) -> tuple[np.ndarray, ...]: @@ -241,11 +237,11 @@ class OldDeepTensor(DeepTensor): def eval_full( self, coords: np.ndarray, - cells: Optional[np.ndarray], + cells: np.ndarray | None, atom_types: np.ndarray, atomic: bool = False, - fparam: Optional[np.ndarray] = None, - aparam: Optional[np.ndarray] = None, + fparam: np.ndarray | None = None, + aparam: np.ndarray | None = None, mixed_type: bool = False, **kwargs: dict, ) -> tuple[np.ndarray, ...]: diff --git a/deepmd/infer/model_devi.py b/deepmd/infer/model_devi.py index c025297fa1..31ee414cbd 100644 --- a/deepmd/infer/model_devi.py +++ b/deepmd/infer/model_devi.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( Any, - Optional, overload, ) @@ -20,14 +19,14 @@ try: from typing import Literal # python >=3.8 except ImportError: - from typing_extensions import Literal # type: ignore + from typing import Literal # type: ignore @overload def calc_model_devi_f( fs: np.ndarray, - real_f: Optional[np.ndarray] = None, - relative: Optional[float] = None, + real_f: np.ndarray | None = None, + relative: float | None = None, atomic: Literal[False] = ..., ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: ... @@ -35,8 +34,8 @@ def calc_model_devi_f( @overload def calc_model_devi_f( fs: np.ndarray, - real_f: Optional[np.ndarray] = None, - relative: Optional[float] = None, + real_f: np.ndarray | None = None, + relative: float | None = None, atomic: Literal[True] = ..., ) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: ... @@ -44,16 +43,16 @@ def calc_model_devi_f( @overload def calc_model_devi_f( fs: np.ndarray, - real_f: Optional[np.ndarray] = None, - relative: Optional[float] = None, + real_f: np.ndarray | None = None, + relative: float | None = None, atomic: bool = False, ) -> tuple[np.ndarray, ...]: ... def calc_model_devi_f( fs: np.ndarray, - real_f: Optional[np.ndarray] = None, - relative: Optional[float] = None, + real_f: np.ndarray | None = None, + relative: float | None = None, atomic: bool = False, ) -> tuple[np.ndarray, ...]: """Calculate model deviation of force. @@ -107,9 +106,7 @@ def calc_model_devi_f( return max_devi_f, min_devi_f, avg_devi_f -def calc_model_devi_e( - es: np.ndarray, real_e: Optional[np.ndarray] = None -) -> np.ndarray: +def calc_model_devi_e(es: np.ndarray, real_e: np.ndarray | None = None) -> np.ndarray: """Calculate model deviation of total energy per atom. Here we don't use the atomic energy, as the decomposition @@ -139,8 +136,8 @@ def calc_model_devi_e( def calc_model_devi_v( vs: np.ndarray, - real_v: Optional[np.ndarray] = None, - relative: Optional[float] = None, + real_v: np.ndarray | None = None, + relative: float | None = None, ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: """Calculate model deviation of virial. @@ -226,7 +223,7 @@ def write_model_devi_out( return devi -def _check_tmaps(tmaps: list[list[str]], ref_tmap: Optional[list[str]] = None) -> bool: +def _check_tmaps(tmaps: list[list[str]], ref_tmap: list[str] | None = None) -> bool: """Check whether type maps are identical.""" assert isinstance(tmaps, list) if ref_tmap is None: @@ -243,18 +240,18 @@ def _check_tmaps(tmaps: list[list[str]], ref_tmap: Optional[list[str]] = None) - def calc_model_devi( coord: np.ndarray, - box: Optional[np.ndarray], + box: np.ndarray | None, atype: np.ndarray, models: list[DeepPot], - fname: Optional[str] = None, + fname: str | None = None, frequency: int = 1, mixed_type: bool = False, - fparam: Optional[np.ndarray] = None, - aparam: Optional[np.ndarray] = None, - real_data: Optional[dict] = None, + fparam: np.ndarray | None = None, + aparam: np.ndarray | None = None, + real_data: dict | None = None, atomic: bool = False, - relative: Optional[float] = None, - relative_v: Optional[float] = None, + relative: float | None = None, + relative_v: float | None = None, ) -> np.ndarray: """Python interface to calculate model deviation. @@ -362,8 +359,8 @@ def make_model_devi( frequency: int, real_error: bool = False, atomic: bool = False, - relative: Optional[float] = None, - relative_v: Optional[float] = None, + relative: float | None = None, + relative_v: float | None = None, **kwargs: Any, ) -> None: """Make model deviation calculation. diff --git a/deepmd/jax/atomic_model/dp_atomic_model.py b/deepmd/jax/atomic_model/dp_atomic_model.py index adfc22c6fa..7227839f1f 100644 --- a/deepmd/jax/atomic_model/dp_atomic_model.py +++ b/deepmd/jax/atomic_model/dp_atomic_model.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( Any, - Optional, ) from deepmd.dpmodel.atomic_model.dp_atomic_model import DPAtomicModel as DPAtomicModelDP @@ -55,9 +54,9 @@ def forward_common_atomic( extended_coord: jnp.ndarray, extended_atype: jnp.ndarray, nlist: jnp.ndarray, - mapping: Optional[jnp.ndarray] = None, - fparam: Optional[jnp.ndarray] = None, - aparam: Optional[jnp.ndarray] = None, + mapping: jnp.ndarray | None = None, + fparam: jnp.ndarray | None = None, + aparam: jnp.ndarray | None = None, ) -> dict[str, jnp.ndarray]: return super().forward_common_atomic( extended_coord, diff --git a/deepmd/jax/atomic_model/linear_atomic_model.py b/deepmd/jax/atomic_model/linear_atomic_model.py index a8889e4346..1c183db7ac 100644 --- a/deepmd/jax/atomic_model/linear_atomic_model.py +++ b/deepmd/jax/atomic_model/linear_atomic_model.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( Any, - Optional, ) from packaging.version import ( @@ -59,9 +58,9 @@ def forward_common_atomic( extended_coord: jnp.ndarray, extended_atype: jnp.ndarray, nlist: jnp.ndarray, - mapping: Optional[jnp.ndarray] = None, - fparam: Optional[jnp.ndarray] = None, - aparam: Optional[jnp.ndarray] = None, + mapping: jnp.ndarray | None = None, + fparam: jnp.ndarray | None = None, + aparam: jnp.ndarray | None = None, ) -> dict[str, jnp.ndarray]: return super().forward_common_atomic( extended_coord, diff --git a/deepmd/jax/atomic_model/pairtab_atomic_model.py b/deepmd/jax/atomic_model/pairtab_atomic_model.py index 8cdd0f7ca6..7f18a6403c 100644 --- a/deepmd/jax/atomic_model/pairtab_atomic_model.py +++ b/deepmd/jax/atomic_model/pairtab_atomic_model.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( Any, - Optional, ) from packaging.version import ( @@ -44,9 +43,9 @@ def forward_common_atomic( extended_coord: jnp.ndarray, extended_atype: jnp.ndarray, nlist: jnp.ndarray, - mapping: Optional[jnp.ndarray] = None, - fparam: Optional[jnp.ndarray] = None, - aparam: Optional[jnp.ndarray] = None, + mapping: jnp.ndarray | None = None, + fparam: jnp.ndarray | None = None, + aparam: jnp.ndarray | None = None, ) -> dict[str, jnp.ndarray]: return super().forward_common_atomic( extended_coord, diff --git a/deepmd/jax/common.py b/deepmd/jax/common.py index 14ae1cad9d..27c7f8883d 100644 --- a/deepmd/jax/common.py +++ b/deepmd/jax/common.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( Any, - Optional, overload, ) @@ -24,7 +23,7 @@ def to_jax_array(array: np.ndarray) -> jnp.ndarray: ... def to_jax_array(array: None) -> None: ... -def to_jax_array(array: Optional[np.ndarray]) -> Optional[jnp.ndarray]: +def to_jax_array(array: np.ndarray | None) -> jnp.ndarray | None: """Convert a numpy array to a JAX array. Parameters diff --git a/deepmd/jax/infer/deep_eval.py b/deepmd/jax/infer/deep_eval.py index fbd8860c0c..ee605e4c07 100644 --- a/deepmd/jax/infer/deep_eval.py +++ b/deepmd/jax/infer/deep_eval.py @@ -1,11 +1,12 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import json +from collections.abc import ( + Callable, +) from typing import ( TYPE_CHECKING, Any, - Callable, Optional, - Union, ) import numpy as np @@ -83,7 +84,7 @@ def __init__( model_file: str, output_def: ModelOutputDef, *args: Any, - auto_batch_size: Union[bool, int, AutoBatchSize] = True, + auto_batch_size: bool | int | AutoBatchSize = True, neighbor_list: Optional["ase.neighborlist.NewPrimitiveNeighborList"] = None, **kwargs: Any, ) -> None: @@ -189,11 +190,11 @@ def get_ntypes_spin(self) -> int: def eval( self, coords: np.ndarray, - cells: Optional[np.ndarray], + cells: np.ndarray | None, atom_types: np.ndarray, atomic: bool = False, - fparam: Optional[np.ndarray] = None, - aparam: Optional[np.ndarray] = None, + fparam: np.ndarray | None = None, + aparam: np.ndarray | None = None, **kwargs: Any, ) -> dict[str, np.ndarray]: """Evaluate the energy, force and virial by using this DP. @@ -330,10 +331,10 @@ def _get_natoms_and_nframes( def _eval_model( self, coords: np.ndarray, - cells: Optional[np.ndarray], + cells: np.ndarray | None, atom_types: np.ndarray, - fparam: Optional[np.ndarray], - aparam: Optional[np.ndarray], + fparam: np.ndarray | None, + aparam: np.ndarray | None, request_defs: list[OutputVariableDef], ) -> tuple[np.ndarray, ...]: model = self.dp diff --git a/deepmd/jax/jax2tf/make_model.py b/deepmd/jax/jax2tf/make_model.py index 341fdf0d1f..3cd30de85a 100644 --- a/deepmd/jax/jax2tf/make_model.py +++ b/deepmd/jax/jax2tf/make_model.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( +from collections.abc import ( Callable, ) diff --git a/deepmd/jax/jax2tf/nlist.py b/deepmd/jax/jax2tf/nlist.py index f85526f1e9..c44a1196c8 100644 --- a/deepmd/jax/jax2tf/nlist.py +++ b/deepmd/jax/jax2tf/nlist.py @@ -1,7 +1,4 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( - Union, -) import tensorflow as tf import tensorflow.experimental.numpy as tnp @@ -17,7 +14,7 @@ def build_neighbor_list( atype: tnp.ndarray, nloc: int, rcut: float, - sel: Union[int, list[int]], + sel: int | list[int], distinguish_types: bool = True, ) -> tnp.ndarray: """Build neighbor list for a single frame. keeps nsel neighbors. diff --git a/deepmd/jax/jax2tf/serialization.py b/deepmd/jax/jax2tf/serialization.py index 096fc41e5a..e819ebf65a 100644 --- a/deepmd/jax/jax2tf/serialization.py +++ b/deepmd/jax/jax2tf/serialization.py @@ -1,8 +1,7 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import json -from typing import ( +from collections.abc import ( Callable, - Optional, ) import tensorflow as tf @@ -155,9 +154,9 @@ def make_call_whether_do_atomic_virial(do_atomic_virial: bool) -> Callable: def call( coord: tnp.ndarray, atype: tnp.ndarray, - box: Optional[tnp.ndarray] = None, - fparam: Optional[tnp.ndarray] = None, - aparam: Optional[tnp.ndarray] = None, + box: tnp.ndarray | None = None, + fparam: tnp.ndarray | None = None, + aparam: tnp.ndarray | None = None, ) -> dict[str, tnp.ndarray]: """Return model prediction. diff --git a/deepmd/jax/jax2tf/tfmodel.py b/deepmd/jax/jax2tf/tfmodel.py index 61c83fa028..85115547d7 100644 --- a/deepmd/jax/jax2tf/tfmodel.py +++ b/deepmd/jax/jax2tf/tfmodel.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( Any, - Optional, ) import jax.experimental.jax2tf as jax2tf @@ -75,9 +74,9 @@ def __call__( self, coord: jnp.ndarray, atype: jnp.ndarray, - box: Optional[jnp.ndarray] = None, - fparam: Optional[jnp.ndarray] = None, - aparam: Optional[jnp.ndarray] = None, + box: jnp.ndarray | None = None, + fparam: jnp.ndarray | None = None, + aparam: jnp.ndarray | None = None, do_atomic_virial: bool = False, ) -> Any: """Return model prediction. @@ -111,9 +110,9 @@ def call( self, coord: jnp.ndarray, atype: jnp.ndarray, - box: Optional[jnp.ndarray] = None, - fparam: Optional[jnp.ndarray] = None, - aparam: Optional[jnp.ndarray] = None, + box: jnp.ndarray | None = None, + fparam: jnp.ndarray | None = None, + aparam: jnp.ndarray | None = None, do_atomic_virial: bool = False, ) -> dict[str, jnp.ndarray]: """Return model prediction. @@ -175,9 +174,9 @@ def call_lower( extended_coord: jnp.ndarray, extended_atype: jnp.ndarray, nlist: jnp.ndarray, - mapping: Optional[jnp.ndarray] = None, - fparam: Optional[jnp.ndarray] = None, - aparam: Optional[jnp.ndarray] = None, + mapping: jnp.ndarray | None = None, + fparam: jnp.ndarray | None = None, + aparam: jnp.ndarray | None = None, do_atomic_virial: bool = False, ) -> dict[str, jnp.ndarray]: if do_atomic_virial: @@ -269,7 +268,7 @@ def get_model_def_script(self) -> str: """Get the model definition script.""" return self.model_def_script - def get_min_nbor_dist(self) -> Optional[float]: + def get_min_nbor_dist(self) -> float | None: """Get the minimum distance between two atoms.""" return self.min_nbor_dist @@ -291,9 +290,9 @@ def mixed_types(self) -> bool: def update_sel( cls, train_data: DeepmdDataSystem, - type_map: Optional[list[str]], + type_map: list[str] | None, local_jdata: dict, - ) -> tuple[dict, Optional[float]]: + ) -> tuple[dict, float | None]: """Update the selection and perform neighbor statistics. Parameters diff --git a/deepmd/jax/model/base_model.py b/deepmd/jax/model/base_model.py index 203da40d07..533181e250 100644 --- a/deepmd/jax/model/base_model.py +++ b/deepmd/jax/model/base_model.py @@ -1,7 +1,4 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( - Optional, -) from deepmd.dpmodel.model.base_model import ( make_base_model, @@ -24,9 +21,9 @@ def forward_common_atomic( extended_coord: jnp.ndarray, extended_atype: jnp.ndarray, nlist: jnp.ndarray, - mapping: Optional[jnp.ndarray] = None, - fparam: Optional[jnp.ndarray] = None, - aparam: Optional[jnp.ndarray] = None, + mapping: jnp.ndarray | None = None, + fparam: jnp.ndarray | None = None, + aparam: jnp.ndarray | None = None, do_atomic_virial: bool = False, ) -> dict[str, jnp.ndarray]: atomic_ret = self.atomic_model.forward_common_atomic( @@ -63,9 +60,9 @@ def eval_output( cc_ext: jnp.ndarray, extended_atype: jnp.ndarray, nlist: jnp.ndarray, - mapping: Optional[jnp.ndarray], - fparam: Optional[jnp.ndarray], - aparam: Optional[jnp.ndarray], + mapping: jnp.ndarray | None, + fparam: jnp.ndarray | None, + aparam: jnp.ndarray | None, *, _kk: str = kk, _atom_axis: int = atom_axis, @@ -120,9 +117,9 @@ def eval_ce( cc_ext: jnp.ndarray, extended_atype: jnp.ndarray, nlist: jnp.ndarray, - mapping: Optional[jnp.ndarray], - fparam: Optional[jnp.ndarray], - aparam: Optional[jnp.ndarray], + mapping: jnp.ndarray | None, + fparam: jnp.ndarray | None, + aparam: jnp.ndarray | None, *, _kk: str = kk, _atom_axis: int = atom_axis - 1, diff --git a/deepmd/jax/model/dp_model.py b/deepmd/jax/model/dp_model.py index ee98a689e4..5545b5505b 100644 --- a/deepmd/jax/model/dp_model.py +++ b/deepmd/jax/model/dp_model.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( Any, - Optional, ) from deepmd.dpmodel.model import ( @@ -52,9 +51,9 @@ def forward_common_atomic( extended_coord: jnp.ndarray, extended_atype: jnp.ndarray, nlist: jnp.ndarray, - mapping: Optional[jnp.ndarray] = None, - fparam: Optional[jnp.ndarray] = None, - aparam: Optional[jnp.ndarray] = None, + mapping: jnp.ndarray | None = None, + fparam: jnp.ndarray | None = None, + aparam: jnp.ndarray | None = None, do_atomic_virial: bool = False, ) -> dict[str, jnp.ndarray]: return forward_common_atomic( diff --git a/deepmd/jax/model/dp_zbl_model.py b/deepmd/jax/model/dp_zbl_model.py index 065dbc7aa7..d9be671a46 100644 --- a/deepmd/jax/model/dp_zbl_model.py +++ b/deepmd/jax/model/dp_zbl_model.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( Any, - Optional, ) from deepmd.dpmodel.model.dp_zbl_model import DPZBLModel as DPZBLModelDP @@ -34,9 +33,9 @@ def forward_common_atomic( extended_coord: jnp.ndarray, extended_atype: jnp.ndarray, nlist: jnp.ndarray, - mapping: Optional[jnp.ndarray] = None, - fparam: Optional[jnp.ndarray] = None, - aparam: Optional[jnp.ndarray] = None, + mapping: jnp.ndarray | None = None, + fparam: jnp.ndarray | None = None, + aparam: jnp.ndarray | None = None, do_atomic_virial: bool = False, ) -> dict[str, jnp.ndarray]: return forward_common_atomic( diff --git a/deepmd/jax/model/hlo.py b/deepmd/jax/model/hlo.py index cbeb915329..47959dd130 100644 --- a/deepmd/jax/model/hlo.py +++ b/deepmd/jax/model/hlo.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( Any, - Optional, ) from deepmd.dpmodel.model.make_model import ( @@ -57,7 +56,7 @@ def __init__( is_aparam_nall: bool, model_output_type: str, mixed_types: bool, - min_nbor_dist: Optional[float], + min_nbor_dist: float | None, sel: list[int], ) -> None: self._call_lower = jax_export.deserialize(stablehlo).call @@ -85,9 +84,9 @@ def __call__( self, coord: jnp.ndarray, atype: jnp.ndarray, - box: Optional[jnp.ndarray] = None, - fparam: Optional[jnp.ndarray] = None, - aparam: Optional[jnp.ndarray] = None, + box: jnp.ndarray | None = None, + fparam: jnp.ndarray | None = None, + aparam: jnp.ndarray | None = None, do_atomic_virial: bool = False, ) -> Any: """Return model prediction. @@ -121,9 +120,9 @@ def call( self, coord: jnp.ndarray, atype: jnp.ndarray, - box: Optional[jnp.ndarray] = None, - fparam: Optional[jnp.ndarray] = None, - aparam: Optional[jnp.ndarray] = None, + box: jnp.ndarray | None = None, + fparam: jnp.ndarray | None = None, + aparam: jnp.ndarray | None = None, do_atomic_virial: bool = False, ) -> dict[str, jnp.ndarray]: """Return model prediction. @@ -175,9 +174,9 @@ def call_lower( extended_coord: jnp.ndarray, extended_atype: jnp.ndarray, nlist: jnp.ndarray, - mapping: Optional[jnp.ndarray] = None, - fparam: Optional[jnp.ndarray] = None, - aparam: Optional[jnp.ndarray] = None, + mapping: jnp.ndarray | None = None, + fparam: jnp.ndarray | None = None, + aparam: jnp.ndarray | None = None, do_atomic_virial: bool = False, ) -> dict[str, jnp.ndarray]: if extended_coord.shape[1] > nlist.shape[1]: @@ -265,7 +264,7 @@ def get_model_def_script(self) -> str: """Get the model definition script.""" return self.model_def_script - def get_min_nbor_dist(self) -> Optional[float]: + def get_min_nbor_dist(self) -> float | None: """Get the minimum distance between two atoms.""" return self.min_nbor_dist @@ -287,9 +286,9 @@ def mixed_types(self) -> bool: def update_sel( cls, train_data: DeepmdDataSystem, - type_map: Optional[list[str]], + type_map: list[str] | None, local_jdata: dict, - ) -> tuple[dict, Optional[float]]: + ) -> tuple[dict, float | None]: """Update the selection and perform neighbor statistics. Parameters diff --git a/deepmd/jax/utils/neighbor_stat.py b/deepmd/jax/utils/neighbor_stat.py index ddfc4199a3..c462665626 100644 --- a/deepmd/jax/utils/neighbor_stat.py +++ b/deepmd/jax/utils/neighbor_stat.py @@ -2,9 +2,6 @@ from collections.abc import ( Iterator, ) -from typing import ( - Optional, -) import numpy as np @@ -81,7 +78,7 @@ def _execute( self, coord: np.ndarray, atype: np.ndarray, - cell: Optional[np.ndarray], + cell: np.ndarray | None, ) -> tuple[np.ndarray, np.ndarray]: """Execute the operation. diff --git a/deepmd/lmp.py b/deepmd/lmp.py index 7ac1570f0f..3daf269c92 100644 --- a/deepmd/lmp.py +++ b/deepmd/lmp.py @@ -9,9 +9,6 @@ from pathlib import ( Path, ) -from typing import ( - Optional, -) import torch # noqa: TID253 from packaging.version import ( @@ -34,7 +31,7 @@ find_libpython = None -def get_env(paths: list[Optional[str]]) -> str: +def get_env(paths: list[str | None]) -> str: """Get the environment variable from given paths.""" return ":".join(p for p in paths if p is not None) diff --git a/deepmd/loggers/loggers.py b/deepmd/loggers/loggers.py index a57d63be59..aa5c73aa6e 100644 --- a/deepmd/loggers/loggers.py +++ b/deepmd/loggers/loggers.py @@ -144,7 +144,7 @@ def setStream(self, stream: "_MPIFileStream") -> NoReturn: def set_log_handles( - level: int, log_path: Optional["Path"] = None, mpi_log: Optional[str] = None + level: int, log_path: Optional["Path"] = None, mpi_log: str | None = None ) -> None: """Set desired level for package loggers and add file handlers. diff --git a/deepmd/loggers/training.py b/deepmd/loggers/training.py index 5de7926460..c7fe94e24d 100644 --- a/deepmd/loggers/training.py +++ b/deepmd/loggers/training.py @@ -1,14 +1,11 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import datetime -from typing import ( - Optional, -) def format_training_message( batch: int, wall_time: float, - eta: Optional[int] = None, + eta: int | None = None, ) -> str: """Format a training message.""" msg = f"batch {batch:7d}: total wall time = {wall_time:.2f} s" @@ -21,7 +18,7 @@ def format_training_message_per_task( batch: int, task_name: str, rmse: dict[str, float], - learning_rate: Optional[float], + learning_rate: float | None, ) -> str: if task_name: task_name += ": " diff --git a/deepmd/main.py b/deepmd/main.py index d829f11ba2..62118ae3c6 100644 --- a/deepmd/main.py +++ b/deepmd/main.py @@ -15,7 +15,6 @@ ) from typing import ( Any, - Optional, ) from deepmd.backend.backend import ( @@ -69,7 +68,7 @@ def __call__( parser: argparse.ArgumentParser, namespace: argparse.Namespace, values: Any, - option_string: Optional[str] = None, + option_string: str | None = None, ) -> None: setattr(namespace, self.dest, BACKEND_TABLE[values]) @@ -87,7 +86,7 @@ def __call__( parser: argparse.ArgumentParser, namespace: argparse.Namespace, values: Any, - option_string: Optional[str] = None, + option_string: str | None = None, ) -> None: if self.call_count == 0: warnings.warn( @@ -946,7 +945,7 @@ def main_parser() -> argparse.ArgumentParser: return parser -def parse_args(args: Optional[list[str]] = None) -> argparse.Namespace: +def parse_args(args: list[str] | None = None) -> argparse.Namespace: """Parse arguments and convert argument strings to objects. Parameters @@ -970,7 +969,7 @@ def parse_args(args: Optional[list[str]] = None) -> argparse.Namespace: return parsed_args -def main(args: Optional[list[str]] = None) -> None: +def main(args: list[str] | None = None) -> None: """DeePMD-kit new entry point. Parameters diff --git a/deepmd/pd/entrypoints/main.py b/deepmd/pd/entrypoints/main.py index fe092111b1..f71a4806ff 100644 --- a/deepmd/pd/entrypoints/main.py +++ b/deepmd/pd/entrypoints/main.py @@ -8,8 +8,6 @@ ) from typing import ( Any, - Optional, - Union, ) import h5py @@ -82,13 +80,13 @@ def get_trainer( config: dict[str, Any], - init_model: Optional[str] = None, - restart_model: Optional[str] = None, - finetune_model: Optional[str] = None, + init_model: str | None = None, + restart_model: str | None = None, + finetune_model: str | None = None, force_load: bool = False, - init_frz_model: Optional[str] = None, - shared_links: Optional[dict[str, Any]] = None, - finetune_links: Optional[dict[str, Any]] = None, + init_frz_model: str | None = None, + shared_links: dict[str, Any] | None = None, + finetune_links: dict[str, Any] | None = None, ) -> training.Trainer: multi_task = "model_dict" in config.get("model", {}) @@ -102,8 +100,8 @@ def prepare_trainer_input_single( model_params_single: dict[str, Any], data_dict_single: dict[str, Any], rank: int = 0, - seed: Optional[int] = None, - ) -> tuple[DpLoaderSet, Optional[DpLoaderSet], Optional[DPPath]]: + seed: int | None = None, + ) -> tuple[DpLoaderSet, DpLoaderSet | None, DPPath | None]: training_dataset_params = data_dict_single["training_data"] validation_dataset_params = data_dict_single.get("validation_data", None) validation_systems = ( @@ -229,10 +227,10 @@ def get_backend_info(self) -> dict: def train( input_file: str, - init_model: Optional[str], - restart: Optional[str], - finetune: Optional[str], - init_frz_model: Optional[str], + init_model: str | None, + restart: str | None, + finetune: str | None, + init_frz_model: str | None, model_branch: str, skip_neighbor_stat: bool = False, use_pretrain_script: bool = False, @@ -347,7 +345,7 @@ def train( def freeze( model: str, output: str = "frozen_model.json", - head: Optional[str] = None, + head: str | None = None, do_atomic_virial: bool = False, ) -> None: paddle.set_flags( @@ -451,12 +449,12 @@ def freeze( def change_bias( input_file: str, mode: str = "change", - bias_value: Optional[list] = None, - datafile: Optional[str] = None, + bias_value: list | None = None, + datafile: str | None = None, system: str = ".", numb_batch: int = 0, - model_branch: Optional[str] = None, - output: Optional[str] = None, + model_branch: str | None = None, + output: str | None = None, ) -> None: if input_file.endswith(".pd"): old_state_dict = paddle.load(input_file) @@ -562,7 +560,7 @@ def change_bias( log.info(f"Saved model to {output_path}") -def main(args: Optional[Union[list[str], argparse.Namespace]] = None): +def main(args: list[str] | argparse.Namespace | None = None): if not isinstance(args, argparse.Namespace): FLAGS = parse_args(args=args) else: diff --git a/deepmd/pd/infer/deep_eval.py b/deepmd/pd/infer/deep_eval.py index 696531ed7f..6715d6d0a0 100644 --- a/deepmd/pd/infer/deep_eval.py +++ b/deepmd/pd/infer/deep_eval.py @@ -1,11 +1,12 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import logging +from collections.abc import ( + Callable, +) from typing import ( TYPE_CHECKING, Any, - Callable, Optional, - Union, ) import numpy as np @@ -106,9 +107,9 @@ def __init__( model_file: str, output_def: ModelOutputDef, *args: Any, - auto_batch_size: Union[bool, int, AutoBatchSize] = True, + auto_batch_size: bool | int | AutoBatchSize = True, neighbor_list: Optional["ase.neighborlist.NewPrimitiveNeighborList"] = None, - head: Optional[Union[str, int]] = None, + head: str | int | None = None, no_jit: bool = False, **kwargs: Any, ) -> None: @@ -314,11 +315,11 @@ def get_model_branch(self) -> tuple[dict[str, str], dict[str, dict[str, Any]]]: def eval( self, coords: np.ndarray, - cells: Optional[np.ndarray], + cells: np.ndarray | None, atom_types: np.ndarray, atomic: bool = False, - fparam: Optional[np.ndarray] = None, - aparam: Optional[np.ndarray] = None, + fparam: np.ndarray | None = None, + aparam: np.ndarray | None = None, **kwargs: Any, ) -> dict[str, np.ndarray]: """Evaluate the energy, force and virial by using this DP. @@ -468,10 +469,10 @@ def _get_natoms_and_nframes( def _eval_model( self, coords: np.ndarray, - cells: Optional[np.ndarray], + cells: np.ndarray | None, atom_types: np.ndarray, - fparam: Optional[np.ndarray], - aparam: Optional[np.ndarray], + fparam: np.ndarray | None, + aparam: np.ndarray | None, request_defs: list[OutputVariableDef], ): if not self.static_model: @@ -586,11 +587,11 @@ def _eval_model( def _eval_model_spin( self, coords: np.ndarray, - cells: Optional[np.ndarray], + cells: np.ndarray | None, atom_types: np.ndarray, spins: np.ndarray, - fparam: Optional[np.ndarray], - aparam: Optional[np.ndarray], + fparam: np.ndarray | None, + aparam: np.ndarray | None, request_defs: list[OutputVariableDef], ) -> tuple[np.ndarray, ...]: model = self.dp.to(DEVICE) @@ -782,10 +783,10 @@ def get_model(self) -> "BaseModel": def eval_descriptor( self, coords: np.ndarray, - cells: Optional[np.ndarray], + cells: np.ndarray | None, atom_types: np.ndarray, - fparam: Optional[np.ndarray] = None, - aparam: Optional[np.ndarray] = None, + fparam: np.ndarray | None = None, + aparam: np.ndarray | None = None, **kwargs: Any, ) -> np.ndarray: """Evaluate descriptors by using this DP. @@ -839,10 +840,10 @@ def eval_descriptor( def eval_fitting_last_layer( self, coords: np.ndarray, - cells: Optional[np.ndarray], + cells: np.ndarray | None, atom_types: np.ndarray, - fparam: Optional[np.ndarray] = None, - aparam: Optional[np.ndarray] = None, + fparam: np.ndarray | None = None, + aparam: np.ndarray | None = None, **kwargs: Any, ) -> np.ndarray: """Evaluate fitting before last layer by using this DP. diff --git a/deepmd/pd/loss/ener.py b/deepmd/pd/loss/ener.py index 09ec5ff49e..fbd806b26d 100644 --- a/deepmd/pd/loss/ener.py +++ b/deepmd/pd/loss/ener.py @@ -1,7 +1,4 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( - Optional, -) import paddle import paddle.nn.functional as F @@ -46,7 +43,7 @@ def __init__( limit_pref_ae: float = 0.0, start_pref_pf: float = 0.0, limit_pref_pf: float = 0.0, - relative_f: Optional[float] = None, + relative_f: float | None = None, enable_atom_ener_coeff: bool = False, start_pref_gf: float = 0.0, limit_pref_gf: float = 0.0, diff --git a/deepmd/pd/model/atomic_model/base_atomic_model.py b/deepmd/pd/model/atomic_model/base_atomic_model.py index e4bf23bf53..87cb18f6fc 100644 --- a/deepmd/pd/model/atomic_model/base_atomic_model.py +++ b/deepmd/pd/model/atomic_model/base_atomic_model.py @@ -1,11 +1,12 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import logging -from typing import ( +from collections.abc import ( Callable, +) +from typing import ( NoReturn, Optional, - Union, ) import numpy as np @@ -77,8 +78,8 @@ def __init__( type_map: list[str], atom_exclude_types: list[int] = [], pair_exclude_types: list[tuple[int, int]] = [], - rcond: Optional[float] = None, - preset_out_bias: Optional[dict[str, np.ndarray]] = None, + rcond: float | None = None, + preset_out_bias: dict[str, np.ndarray] | None = None, data_stat_protect: float = 1e-2, ) -> None: paddle.nn.Layer.__init__(self) @@ -160,7 +161,7 @@ def has_default_fparam(self) -> bool: def reinit_atom_exclude( self, - exclude_types: Optional[list[int]] = None, + exclude_types: list[int] | None = None, ) -> None: if exclude_types is None: exclude_types = [] @@ -224,10 +225,10 @@ def forward_common_atomic( extended_coord: paddle.Tensor, extended_atype: paddle.Tensor, nlist: paddle.Tensor, - mapping: Optional[paddle.Tensor] = None, - fparam: Optional[paddle.Tensor] = None, - aparam: Optional[paddle.Tensor] = None, - comm_dict: Optional[list[paddle.Tensor]] = None, + mapping: paddle.Tensor | None = None, + fparam: paddle.Tensor | None = None, + aparam: paddle.Tensor | None = None, + comm_dict: list[paddle.Tensor] | None = None, ) -> dict[str, paddle.Tensor]: """Common interface for atomic inference. @@ -306,10 +307,10 @@ def forward( extended_coord: paddle.Tensor, extended_atype: paddle.Tensor, nlist: paddle.Tensor, - mapping: Optional[paddle.Tensor] = None, - fparam: Optional[paddle.Tensor] = None, - aparam: Optional[paddle.Tensor] = None, - comm_dict: Optional[list[paddle.Tensor]] = None, + mapping: paddle.Tensor | None = None, + fparam: paddle.Tensor | None = None, + aparam: paddle.Tensor | None = None, + comm_dict: list[paddle.Tensor] | None = None, ) -> dict[str, paddle.Tensor]: return self.forward_common_atomic( extended_coord, @@ -389,8 +390,8 @@ def deserialize(cls, data: dict) -> "BaseAtomicModel": def compute_or_load_stat( self, - merged: Union[Callable[[], list[dict]], list[dict]], - stat_file_path: Optional[DPPath] = None, + merged: Callable[[], list[dict]] | list[dict], + stat_file_path: DPPath | None = None, compute_or_load_out_stat: bool = True, ) -> NoReturn: """ @@ -416,8 +417,8 @@ def compute_or_load_stat( def compute_or_load_out_stat( self, - merged: Union[Callable[[], list[dict]], list[dict]], - stat_file_path: Optional[DPPath] = None, + merged: Callable[[], list[dict]] | list[dict], + stat_file_path: DPPath | None = None, ) -> None: """ Compute the output statistics (e.g. energy bias) for the fitting net from packed data. @@ -466,8 +467,8 @@ def apply_out_stat( def change_out_bias( self, - sample_merged: Union[Callable[[], list[dict]], list[dict]], - stat_file_path: Optional[DPPath] = None, + sample_merged: Callable[[], list[dict]] | list[dict], + stat_file_path: DPPath | None = None, bias_adjust_mode: str = "change-by-statistic", ) -> None: """Change the output bias according to the input data and the pretrained model. @@ -517,7 +518,7 @@ def change_out_bias( def compute_fitting_input_stat( self, - sample_merged: Union[Callable[[], list[dict]], list[dict]], + sample_merged: Callable[[], list[dict]] | list[dict], ) -> None: """Compute the input statistics (e.g. mean and stddev) for the atomic model from packed data. @@ -539,9 +540,9 @@ def _get_forward_wrapper_func(self) -> Callable[..., paddle.Tensor]: def model_forward( coord: paddle.Tensor, atype: paddle.Tensor, - box: Optional[paddle.Tensor], - fparam: Optional[paddle.Tensor] = None, - aparam: Optional[paddle.Tensor] = None, + box: paddle.Tensor | None, + fparam: paddle.Tensor | None = None, + aparam: paddle.Tensor | None = None, ) -> dict[str, paddle.Tensor]: with ( paddle.no_grad() diff --git a/deepmd/pd/model/atomic_model/dp_atomic_model.py b/deepmd/pd/model/atomic_model/dp_atomic_model.py index 7e102f6456..c5aa8b8a56 100644 --- a/deepmd/pd/model/atomic_model/dp_atomic_model.py +++ b/deepmd/pd/model/atomic_model/dp_atomic_model.py @@ -1,11 +1,12 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import functools import logging +from collections.abc import ( + Callable, +) from typing import ( Any, - Callable, Optional, - Union, ) import paddle @@ -70,7 +71,7 @@ def __init__( self.eval_fitting_last_layer_list = [] # register 'type_map' as buffer - def _string_to_array(s: Union[str, list[str]]) -> list[int]: + def _string_to_array(s: str | list[str]) -> list[int]: return [ord(c) for c in s] if type_map is not None: @@ -290,10 +291,10 @@ def forward_atomic( extended_coord: paddle.Tensor, extended_atype: paddle.Tensor, nlist: paddle.Tensor, - mapping: Optional[paddle.Tensor] = None, - fparam: Optional[paddle.Tensor] = None, - aparam: Optional[paddle.Tensor] = None, - comm_dict: Optional[dict[str, paddle.Tensor]] = None, + mapping: paddle.Tensor | None = None, + fparam: paddle.Tensor | None = None, + aparam: paddle.Tensor | None = None, + comm_dict: dict[str, paddle.Tensor] | None = None, ) -> dict[str, paddle.Tensor]: """Return atomic prediction. @@ -357,7 +358,7 @@ def get_out_bias(self) -> paddle.Tensor: def compute_or_load_stat( self, sampled_func: Callable[[], list[dict]], - stat_file_path: Optional[DPPath] = None, + stat_file_path: DPPath | None = None, compute_or_load_out_stat: bool = True, ) -> None: """ @@ -403,8 +404,8 @@ def wrapped_sampler(): def compute_fitting_input_stat( self, - sample_merged: Union[Callable[[], list[dict]], list[dict]], - stat_file_path: Optional[DPPath] = None, + sample_merged: Callable[[], list[dict]] | list[dict], + stat_file_path: DPPath | None = None, ) -> None: """Compute the input statistics (e.g. mean and stddev) for the fittings from packed data. diff --git a/deepmd/pd/model/descriptor/descriptor.py b/deepmd/pd/model/descriptor/descriptor.py index 3050b7dca3..d4ca4bc151 100644 --- a/deepmd/pd/model/descriptor/descriptor.py +++ b/deepmd/pd/model/descriptor/descriptor.py @@ -4,11 +4,11 @@ ABC, abstractmethod, ) -from typing import ( +from collections.abc import ( Callable, +) +from typing import ( NoReturn, - Optional, - Union, ) import paddle @@ -101,8 +101,8 @@ def get_env_protection(self) -> float: def compute_input_stats( self, - merged: Union[Callable[[], list[dict]], list[dict]], - path: Optional[DPPath] = None, + merged: Callable[[], list[dict]] | list[dict], + path: DPPath | None = None, ) -> NoReturn: """ Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data. @@ -177,9 +177,9 @@ def forward( nlist: paddle.Tensor, extended_coord: paddle.Tensor, extended_atype: paddle.Tensor, - extended_atype_embd: Optional[paddle.Tensor] = None, - mapping: Optional[paddle.Tensor] = None, - type_embedding: Optional[paddle.Tensor] = None, + extended_atype_embd: paddle.Tensor | None = None, + mapping: paddle.Tensor | None = None, + type_embedding: paddle.Tensor | None = None, ): """Calculate DescriptorBlock.""" pass diff --git a/deepmd/pd/model/descriptor/dpa1.py b/deepmd/pd/model/descriptor/dpa1.py index 1722140316..65baa8daa6 100644 --- a/deepmd/pd/model/descriptor/dpa1.py +++ b/deepmd/pd/model/descriptor/dpa1.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( +from collections.abc import ( Callable, - Optional, - Union, ) import paddle @@ -213,7 +211,7 @@ def __init__( self, rcut: float, rcut_smth: float, - sel: Union[list[int], int], + sel: list[int] | int, ntypes: int, neuron: list = [25, 50, 100], axis_neuron: int = 16, @@ -235,17 +233,17 @@ def __init__( concat_output_tebd: bool = True, trainable: bool = True, trainable_ln: bool = True, - ln_eps: Optional[float] = 1e-5, + ln_eps: float | None = 1e-5, smooth_type_embedding: bool = True, type_one_side: bool = False, - stripped_type_embedding: Optional[bool] = None, - seed: Optional[Union[int, list[int]]] = None, + stripped_type_embedding: bool | None = None, + seed: int | list[int] | None = None, use_econf_tebd: bool = False, use_tebd_bias: bool = False, - type_map: Optional[list[str]] = None, + type_map: list[str] | None = None, # not implemented spin=None, - type: Optional[str] = None, + type: str | None = None, ) -> None: super().__init__() # Ensure compatibility with the deprecated stripped_type_embedding option. @@ -436,8 +434,8 @@ def dim_emb(self): def compute_input_stats( self, - merged: Union[Callable[[], list[dict]], list[dict]], - path: Optional[DPPath] = None, + merged: Callable[[], list[dict]] | list[dict], + path: DPPath | None = None, ): """ Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data. @@ -620,8 +618,8 @@ def forward( extended_coord: paddle.Tensor, extended_atype: paddle.Tensor, nlist: paddle.Tensor, - mapping: Optional[paddle.Tensor] = None, - comm_dict: Optional[list[paddle.Tensor]] = None, + mapping: paddle.Tensor | None = None, + comm_dict: list[paddle.Tensor] | None = None, ): """Compute the descriptor. @@ -689,9 +687,9 @@ def forward( def update_sel( cls, train_data: DeepmdDataSystem, - type_map: Optional[list[str]], + type_map: list[str] | None, local_jdata: dict, - ) -> tuple[dict, Optional[float]]: + ) -> tuple[dict, float | None]: """Update the selection and perform neighbor statistics. Parameters diff --git a/deepmd/pd/model/descriptor/dpa2.py b/deepmd/pd/model/descriptor/dpa2.py index c8eb4ca117..645ad2f2fe 100644 --- a/deepmd/pd/model/descriptor/dpa2.py +++ b/deepmd/pd/model/descriptor/dpa2.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( +from collections.abc import ( Callable, - Optional, - Union, ) import paddle @@ -80,9 +78,9 @@ def __init__( self, ntypes: int, # args for repinit - repinit: Union[RepinitArgs, dict], + repinit: RepinitArgs | dict, # args for repformer - repformer: Union[RepformerArgs, dict], + repformer: RepformerArgs | dict, # kwargs for descriptor concat_output_tebd: bool = True, precision: str = "float64", @@ -90,11 +88,11 @@ def __init__( exclude_types: list[tuple[int, int]] = [], env_protection: float = 0.0, trainable: bool = True, - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, add_tebd_to_repinit_out: bool = False, use_econf_tebd: bool = False, use_tebd_bias: bool = False, - type_map: Optional[list[str]] = None, + type_map: list[str] | None = None, ) -> None: r"""The DPA-2 descriptor[1]_. @@ -503,8 +501,8 @@ def dim_emb(self): def compute_input_stats( self, - merged: Union[Callable[[], list[dict]], list[dict]], - path: Optional[DPPath] = None, + merged: Callable[[], list[dict]] | list[dict], + path: DPPath | None = None, ) -> None: """ Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data. @@ -727,8 +725,8 @@ def forward( extended_coord: paddle.Tensor, extended_atype: paddle.Tensor, nlist: paddle.Tensor, - mapping: Optional[paddle.Tensor] = None, - comm_dict: Optional[list[paddle.Tensor]] = None, + mapping: paddle.Tensor | None = None, + comm_dict: list[paddle.Tensor] | None = None, ): """Compute the descriptor. @@ -851,9 +849,9 @@ def forward( def update_sel( cls, train_data: DeepmdDataSystem, - type_map: Optional[list[str]], + type_map: list[str] | None, local_jdata: dict, - ) -> tuple[dict, Optional[float]]: + ) -> tuple[dict, float | None]: """Update the selection and perform neighbor statistics. Parameters diff --git a/deepmd/pd/model/descriptor/dpa3.py b/deepmd/pd/model/descriptor/dpa3.py index 80c5b1c000..205829860f 100644 --- a/deepmd/pd/model/descriptor/dpa3.py +++ b/deepmd/pd/model/descriptor/dpa3.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( +from collections.abc import ( Callable, - Optional, - Union, ) import paddle @@ -106,7 +104,7 @@ def __init__( self, ntypes: int, # args for repflow - repflow: Union[RepFlowArgs, dict], + repflow: RepFlowArgs | dict, # kwargs for descriptor concat_output_tebd: bool = False, activation_function: str = "silu", @@ -114,11 +112,11 @@ def __init__( exclude_types: list[tuple[int, int]] = [], env_protection: float = 0.0, trainable: bool = True, - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, use_econf_tebd: bool = False, use_tebd_bias: bool = False, use_loc_mapping: bool = True, - type_map: Optional[list[str]] = None, + type_map: list[str] | None = None, ) -> None: super().__init__() @@ -371,8 +369,8 @@ def dim_emb(self): def compute_input_stats( self, - merged: Union[Callable[[], list[dict]], list[dict]], - path: Optional[DPPath] = None, + merged: Callable[[], list[dict]] | list[dict], + path: DPPath | None = None, ) -> None: """ Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data. @@ -488,8 +486,8 @@ def forward( extended_coord: paddle.Tensor, extended_atype: paddle.Tensor, nlist: paddle.Tensor, - mapping: Optional[paddle.Tensor] = None, - comm_dict: Optional[list[paddle.Tensor]] = None, + mapping: paddle.Tensor | None = None, + comm_dict: list[paddle.Tensor] | None = None, ): """Compute the descriptor. @@ -557,9 +555,9 @@ def forward( def update_sel( cls, train_data: DeepmdDataSystem, - type_map: Optional[list[str]], + type_map: list[str] | None, local_jdata: dict, - ) -> tuple[dict, Optional[float]]: + ) -> tuple[dict, float | None]: """Update the selection and perform neighbor statistics. Parameters diff --git a/deepmd/pd/model/descriptor/repflow_layer.py b/deepmd/pd/model/descriptor/repflow_layer.py index 71e4d44ce2..eca1659b67 100644 --- a/deepmd/pd/model/descriptor/repflow_layer.py +++ b/deepmd/pd/model/descriptor/repflow_layer.py @@ -1,8 +1,4 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( - Optional, - Union, -) import paddle import paddle.nn as nn @@ -63,7 +59,7 @@ def __init__( update_residual: float = 0.1, update_residual_init: str = "const", precision: str = "float64", - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, trainable: bool = True, ) -> None: super().__init__() diff --git a/deepmd/pd/model/descriptor/repflows.py b/deepmd/pd/model/descriptor/repflows.py index 2b9760bbe6..21b5984ef3 100644 --- a/deepmd/pd/model/descriptor/repflows.py +++ b/deepmd/pd/model/descriptor/repflows.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( +from collections.abc import ( Callable, - Optional, - Union, ) import paddle @@ -227,7 +225,7 @@ def __init__( sel_reduce_factor: float = 10.0, use_loc_mapping: bool = True, optim_update: bool = True, - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, trainable: bool = True, ) -> None: super().__init__() @@ -455,9 +453,9 @@ def forward( nlist: paddle.Tensor, extended_coord: paddle.Tensor, extended_atype: paddle.Tensor, - extended_atype_embd: Optional[paddle.Tensor] = None, - mapping: Optional[paddle.Tensor] = None, - comm_dict: Optional[list[paddle.Tensor]] = None, + extended_atype_embd: paddle.Tensor | None = None, + mapping: paddle.Tensor | None = None, + comm_dict: list[paddle.Tensor] | None = None, ): parallel_mode = comm_dict is not None if not parallel_mode: @@ -730,8 +728,8 @@ def forward( def compute_input_stats( self, - merged: Union[Callable[[], list[dict]], list[dict]], - path: Optional[DPPath] = None, + merged: Callable[[], list[dict]] | list[dict], + path: DPPath | None = None, ) -> None: """ Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data. diff --git a/deepmd/pd/model/descriptor/repformer_layer.py b/deepmd/pd/model/descriptor/repformer_layer.py index fc66e1d6af..efb1881d59 100644 --- a/deepmd/pd/model/descriptor/repformer_layer.py +++ b/deepmd/pd/model/descriptor/repformer_layer.py @@ -1,8 +1,4 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( - Optional, - Union, -) import paddle import paddle.nn as nn @@ -44,7 +40,7 @@ def get_residual( _mode: str = "norm", trainable: bool = True, precision: str = "float64", - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, ) -> paddle.Tensor: r""" Get residual tensor for one update vector. @@ -162,7 +158,7 @@ def __init__( smooth: bool = True, attnw_shift: float = 20.0, precision: str = "float64", - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, trainable: bool = True, ) -> None: """Return neighbor-wise multi-head self-attention maps, with gate mechanism.""" @@ -289,7 +285,7 @@ def __init__( input_dim: int, head_num: int, precision: str = "float64", - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, trainable: bool = True, ) -> None: super().__init__() @@ -379,7 +375,7 @@ def __init__( input_dim: int, head_num: int, precision: str = "float64", - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, trainable: bool = True, ) -> None: super().__init__() @@ -458,7 +454,7 @@ def __init__( smooth: bool = True, attnw_shift: float = 20.0, precision: str = "float64", - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, trainable: bool = True, ) -> None: super().__init__() @@ -622,11 +618,11 @@ def __init__( smooth: bool = True, precision: str = "float64", trainable_ln: bool = True, - ln_eps: Optional[float] = 1e-5, + ln_eps: float | None = 1e-5, use_sqrt_nnei: bool = True, g1_out_conv: bool = True, g1_out_mlp: bool = True, - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, trainable: bool = True, ) -> None: super().__init__() diff --git a/deepmd/pd/model/descriptor/repformers.py b/deepmd/pd/model/descriptor/repformers.py index f0d18a0908..24f92f1bee 100644 --- a/deepmd/pd/model/descriptor/repformers.py +++ b/deepmd/pd/model/descriptor/repformers.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( +from collections.abc import ( Callable, - Optional, - Union, ) import paddle @@ -113,8 +111,8 @@ def __init__( env_protection: float = 0.0, precision: str = "float64", trainable_ln: bool = True, - ln_eps: Optional[float] = 1e-5, - seed: Optional[Union[int, list[int]]] = None, + ln_eps: float | None = 1e-5, + seed: int | list[int] | None = None, use_sqrt_nnei: bool = True, g1_out_conv: bool = True, g1_out_mlp: bool = True, @@ -420,10 +418,10 @@ def forward( nlist: paddle.Tensor, extended_coord: paddle.Tensor, extended_atype: paddle.Tensor, - extended_atype_embd: Optional[paddle.Tensor] = None, - mapping: Optional[paddle.Tensor] = None, - type_embedding: Optional[paddle.Tensor] = None, - comm_dict: Optional[list[paddle.Tensor]] = None, + extended_atype_embd: paddle.Tensor | None = None, + mapping: paddle.Tensor | None = None, + type_embedding: paddle.Tensor | None = None, + comm_dict: list[paddle.Tensor] | None = None, ): if (comm_dict is None or len(comm_dict) == 0) and paddle.in_dynamic_mode(): assert mapping is not None @@ -599,8 +597,8 @@ def forward( def compute_input_stats( self, - merged: Union[Callable[[], list[dict]], list[dict]], - path: Optional[DPPath] = None, + merged: Callable[[], list[dict]] | list[dict], + path: DPPath | None = None, ) -> None: """ Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data. diff --git a/deepmd/pd/model/descriptor/se_a.py b/deepmd/pd/model/descriptor/se_a.py index 109c7ba3c4..17ef5d67c4 100644 --- a/deepmd/pd/model/descriptor/se_a.py +++ b/deepmd/pd/model/descriptor/se_a.py @@ -1,10 +1,10 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import itertools -from typing import ( +from collections.abc import ( Callable, +) +from typing import ( ClassVar, - Optional, - Union, ) import numpy as np @@ -84,9 +84,9 @@ def __init__( env_protection: float = 0.0, type_one_side: bool = True, trainable: bool = True, - seed: Optional[Union[int, list[int]]] = None, - ntypes: Optional[int] = None, # to be compat with input - type_map: Optional[list[str]] = None, + seed: int | list[int] | None = None, + ntypes: int | None = None, # to be compat with input + type_map: list[str] | None = None, # not implemented spin=None, ) -> None: @@ -226,8 +226,8 @@ def change_type_map( def compute_input_stats( self, - merged: Union[Callable[[], list[dict]], list[dict]], - path: Optional[DPPath] = None, + merged: Callable[[], list[dict]] | list[dict], + path: DPPath | None = None, ): """ Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data. @@ -284,8 +284,8 @@ def forward( coord_ext: paddle.Tensor, atype_ext: paddle.Tensor, nlist: paddle.Tensor, - mapping: Optional[paddle.Tensor] = None, - comm_dict: Optional[list[paddle.Tensor]] = None, + mapping: paddle.Tensor | None = None, + comm_dict: list[paddle.Tensor] | None = None, ): """Compute the descriptor. @@ -401,9 +401,9 @@ def t_cvt(xx): def update_sel( cls, train_data: DeepmdDataSystem, - type_map: Optional[list[str]], + type_map: list[str] | None, local_jdata: dict, - ) -> tuple[dict, Optional[float]]: + ) -> tuple[dict, float | None]: """Update the selection and perform neighbor statistics. Parameters @@ -449,7 +449,7 @@ def __init__( env_protection: float = 0.0, type_one_side: bool = True, trainable: bool = True, - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, **kwargs, ) -> None: """Construct an embedding net of type `se_a`. @@ -623,8 +623,8 @@ def __getitem__(self, key): def compute_input_stats( self, - merged: Union[Callable[[], list[dict]], list[dict]], - path: Optional[DPPath] = None, + merged: Callable[[], list[dict]] | list[dict], + path: DPPath | None = None, ) -> None: """ Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data. @@ -684,7 +684,7 @@ def reinit_exclude( def enable_compression( self, table_data: dict[str, paddle.Tensor], - table_config: list[Union[int, float]], + table_config: list[int | float], lower: dict[str, int], upper: dict[str, int], ) -> None: @@ -722,9 +722,9 @@ def forward( nlist: paddle.Tensor, extended_coord: paddle.Tensor, extended_atype: paddle.Tensor, - extended_atype_embd: Optional[paddle.Tensor] = None, - mapping: Optional[paddle.Tensor] = None, - type_embedding: Optional[paddle.Tensor] = None, + extended_atype_embd: paddle.Tensor | None = None, + mapping: paddle.Tensor | None = None, + type_embedding: paddle.Tensor | None = None, ): """Calculate decoded embedding for each atom. diff --git a/deepmd/pd/model/descriptor/se_atten.py b/deepmd/pd/model/descriptor/se_atten.py index ceae16f409..2c93c35ef8 100644 --- a/deepmd/pd/model/descriptor/se_atten.py +++ b/deepmd/pd/model/descriptor/se_atten.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( +from collections.abc import ( Callable, - Optional, - Union, ) import paddle @@ -56,7 +54,7 @@ def __init__( self, rcut: float, rcut_smth: float, - sel: Union[list[int], int], + sel: list[int] | int, ntypes: int, neuron: list = [25, 50, 100], axis_neuron: int = 16, @@ -78,9 +76,9 @@ def __init__( exclude_types: list[tuple[int, int]] = [], env_protection: float = 0.0, trainable_ln: bool = True, - ln_eps: Optional[float] = 1e-5, - seed: Optional[Union[int, list[int]]] = None, - type: Optional[str] = None, + ln_eps: float | None = 1e-5, + seed: int | list[int] | None = None, + type: str | None = None, trainable: bool = True, ) -> None: r"""Construct an embedding net of type `se_atten`. @@ -369,8 +367,8 @@ def dim_emb(self): def compute_input_stats( self, - merged: Union[Callable[[], list[dict]], list[dict]], - path: Optional[DPPath] = None, + merged: Callable[[], list[dict]] | list[dict], + path: DPPath | None = None, ) -> None: """ Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data. @@ -459,9 +457,9 @@ def forward( nlist: paddle.Tensor, extended_coord: paddle.Tensor, extended_atype: paddle.Tensor, - extended_atype_embd: Optional[paddle.Tensor] = None, - mapping: Optional[paddle.Tensor] = None, - type_embedding: Optional[paddle.Tensor] = None, + extended_atype_embd: paddle.Tensor | None = None, + mapping: paddle.Tensor | None = None, + type_embedding: paddle.Tensor | None = None, ): """Compute the descriptor. @@ -676,12 +674,12 @@ def __init__( do_mask: bool = False, scaling_factor: float = 1.0, normalize: bool = True, - temperature: Optional[float] = None, + temperature: float | None = None, trainable_ln: bool = True, ln_eps: float = 1e-5, smooth: bool = True, precision: str = DEFAULT_PRECISION, - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, trainable: bool = True, ) -> None: """Construct a neighbor-wise attention net.""" @@ -727,8 +725,8 @@ def forward( self, input_G, nei_mask, - input_r: Optional[paddle.Tensor] = None, - sw: Optional[paddle.Tensor] = None, + input_r: paddle.Tensor | None = None, + sw: paddle.Tensor | None = None, ): """Compute the multi-layer gated self-attention. @@ -820,12 +818,12 @@ def __init__( do_mask: bool = False, scaling_factor: float = 1.0, normalize: bool = True, - temperature: Optional[float] = None, + temperature: float | None = None, smooth: bool = True, trainable_ln: bool = True, ln_eps: float = 1e-5, precision: str = DEFAULT_PRECISION, - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, trainable: bool = True, ) -> None: """Construct a neighbor-wise attention layer.""" @@ -868,8 +866,8 @@ def forward( self, x, nei_mask, - input_r: Optional[paddle.Tensor] = None, - sw: Optional[paddle.Tensor] = None, + input_r: paddle.Tensor | None = None, + sw: paddle.Tensor | None = None, ): residual = x x, _ = self.attention_layer(x, nei_mask, input_r=input_r, sw=sw) @@ -930,11 +928,11 @@ def __init__( do_mask: bool = False, scaling_factor: float = 1.0, normalize: bool = True, - temperature: Optional[float] = None, + temperature: float | None = None, bias: bool = True, smooth: bool = True, precision: str = DEFAULT_PRECISION, - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, trainable: bool = True, ) -> None: """Construct a multi-head neighbor-wise attention net.""" @@ -986,8 +984,8 @@ def forward( self, query, nei_mask, - input_r: Optional[paddle.Tensor] = None, - sw: Optional[paddle.Tensor] = None, + input_r: paddle.Tensor | None = None, + sw: paddle.Tensor | None = None, attnw_shift: float = 20.0, ): """Compute the multi-head gated self-attention. diff --git a/deepmd/pd/model/descriptor/se_atten_v2.py b/deepmd/pd/model/descriptor/se_atten_v2.py index 6c90846f9a..8855fb3037 100644 --- a/deepmd/pd/model/descriptor/se_atten_v2.py +++ b/deepmd/pd/model/descriptor/se_atten_v2.py @@ -1,8 +1,4 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( - Optional, - Union, -) import paddle @@ -40,7 +36,7 @@ def __init__( self, rcut: float, rcut_smth: float, - sel: Union[list[int], int], + sel: list[int] | int, ntypes: int, neuron: list = [25, 50, 100], axis_neuron: int = 16, @@ -61,16 +57,16 @@ def __init__( concat_output_tebd: bool = True, trainable: bool = True, trainable_ln: bool = True, - ln_eps: Optional[float] = 1e-5, + ln_eps: float | None = 1e-5, type_one_side: bool = False, - stripped_type_embedding: Optional[bool] = None, - seed: Optional[Union[int, list[int]]] = None, + stripped_type_embedding: bool | None = None, + seed: int | list[int] | None = None, use_econf_tebd: bool = False, use_tebd_bias: bool = False, - type_map: Optional[list[str]] = None, + type_map: list[str] | None = None, # not implemented spin=None, - type: Optional[str] = None, + type: str | None = None, ) -> None: r"""Construct smooth version of embedding net of type `se_atten_v2`. diff --git a/deepmd/pd/model/descriptor/se_t_tebd.py b/deepmd/pd/model/descriptor/se_t_tebd.py index e9d4053612..b0f409f2bb 100644 --- a/deepmd/pd/model/descriptor/se_t_tebd.py +++ b/deepmd/pd/model/descriptor/se_t_tebd.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( +from collections.abc import ( Callable, - Optional, - Union, ) import paddle @@ -124,7 +122,7 @@ def __init__( self, rcut: float, rcut_smth: float, - sel: Union[list[int], int], + sel: list[int] | int, ntypes: int, neuron: list = [2, 4, 8], tebd_dim: int = 8, @@ -136,8 +134,8 @@ def __init__( exclude_types: list[tuple[int, int]] = [], precision: str = "float64", trainable: bool = True, - seed: Optional[Union[int, list[int]]] = None, - type_map: Optional[list[str]] = None, + seed: int | list[int] | None = None, + type_map: list[str] | None = None, concat_output_tebd: bool = True, use_econf_tebd: bool = False, use_tebd_bias=False, @@ -296,8 +294,8 @@ def dim_emb(self): def compute_input_stats( self, - merged: Union[Callable[[], list[dict]], list[dict]], - path: Optional[DPPath] = None, + merged: Callable[[], list[dict]] | list[dict], + path: DPPath | None = None, ): """ Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data. @@ -431,8 +429,8 @@ def forward( extended_coord: paddle.Tensor, extended_atype: paddle.Tensor, nlist: paddle.Tensor, - mapping: Optional[paddle.Tensor] = None, - comm_dict: Optional[list[paddle.Tensor]] = None, + mapping: paddle.Tensor | None = None, + comm_dict: list[paddle.Tensor] | None = None, ): """Compute the descriptor. @@ -500,9 +498,9 @@ def forward( def update_sel( cls, train_data: DeepmdDataSystem, - type_map: Optional[list[str]], + type_map: list[str] | None, local_jdata: dict, - ) -> tuple[dict, Optional[float]]: + ) -> tuple[dict, float | None]: """Update the selection and perform neighbor statistics. Parameters @@ -535,7 +533,7 @@ def __init__( self, rcut: float, rcut_smth: float, - sel: Union[list[int], int], + sel: list[int] | int, ntypes: int, neuron: list = [25, 50, 100], tebd_dim: int = 8, @@ -547,7 +545,7 @@ def __init__( exclude_types: list[tuple[int, int]] = [], env_protection: float = 0.0, smooth: bool = True, - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, trainable: bool = True, ) -> None: super().__init__() @@ -719,8 +717,8 @@ def dim_emb(self): def compute_input_stats( self, - merged: Union[Callable[[], list[dict]], list[dict]], - path: Optional[DPPath] = None, + merged: Callable[[], list[dict]] | list[dict], + path: DPPath | None = None, ) -> None: """ Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data. @@ -782,9 +780,9 @@ def forward( nlist: paddle.Tensor, extended_coord: paddle.Tensor, extended_atype: paddle.Tensor, - extended_atype_embd: Optional[paddle.Tensor] = None, - mapping: Optional[paddle.Tensor] = None, - type_embedding: Optional[paddle.Tensor] = None, + extended_atype_embd: paddle.Tensor | None = None, + mapping: paddle.Tensor | None = None, + type_embedding: paddle.Tensor | None = None, ): """Compute the descriptor. diff --git a/deepmd/pd/model/model/dp_model.py b/deepmd/pd/model/model/dp_model.py index e014be5b68..fe107263c4 100644 --- a/deepmd/pd/model/model/dp_model.py +++ b/deepmd/pd/model/model/dp_model.py @@ -1,7 +1,4 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( - Optional, -) import paddle @@ -20,9 +17,9 @@ class DPModelCommon: def update_sel( cls, train_data: DeepmdDataSystem, - type_map: Optional[list[str]], + type_map: list[str] | None, local_jdata: dict, - ) -> tuple[dict, Optional[float]]: + ) -> tuple[dict, float | None]: """Update the selection and perform neighbor statistics. Parameters diff --git a/deepmd/pd/model/model/ener_model.py b/deepmd/pd/model/model/ener_model.py index 36cd1211ba..8111ae434d 100644 --- a/deepmd/pd/model/model/ener_model.py +++ b/deepmd/pd/model/model/ener_model.py @@ -1,7 +1,4 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( - Optional, -) import paddle @@ -68,9 +65,9 @@ def forward( self, coord, atype, - box: Optional[paddle.Tensor] = None, - fparam: Optional[paddle.Tensor] = None, - aparam: Optional[paddle.Tensor] = None, + box: paddle.Tensor | None = None, + fparam: paddle.Tensor | None = None, + aparam: paddle.Tensor | None = None, do_atomic_virial: bool = False, ) -> dict[str, paddle.Tensor]: model_ret = self.forward_common( @@ -111,11 +108,11 @@ def forward_lower( extended_coord, extended_atype, nlist, - mapping: Optional[paddle.Tensor] = None, - fparam: Optional[paddle.Tensor] = None, - aparam: Optional[paddle.Tensor] = None, + mapping: paddle.Tensor | None = None, + fparam: paddle.Tensor | None = None, + aparam: paddle.Tensor | None = None, do_atomic_virial: bool = False, - comm_dict: Optional[list[paddle.Tensor]] = None, + comm_dict: list[paddle.Tensor] | None = None, ): model_ret = self.forward_common_lower( extended_coord, diff --git a/deepmd/pd/model/model/frozen.py b/deepmd/pd/model/model/frozen.py index e8128c6bd1..f585cb1e24 100644 --- a/deepmd/pd/model/model/frozen.py +++ b/deepmd/pd/model/model/frozen.py @@ -1,8 +1,5 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import json -from typing import ( - Optional, -) import paddle @@ -101,9 +98,9 @@ def forward( self, coord, atype, - box: Optional[paddle.Tensor] = None, - fparam: Optional[paddle.Tensor] = None, - aparam: Optional[paddle.Tensor] = None, + box: paddle.Tensor | None = None, + fparam: paddle.Tensor | None = None, + aparam: paddle.Tensor | None = None, do_atomic_virial: bool = False, ) -> dict[str, paddle.Tensor]: return self.model.forward( @@ -123,7 +120,7 @@ def get_model_def_script(self) -> str: # be a problem return self.model.get_model_def_script() - def get_min_nbor_dist(self) -> Optional[float]: + def get_min_nbor_dist(self) -> float | None: """Get the minimum neighbor distance.""" return self.model.get_min_nbor_dist() @@ -154,9 +151,9 @@ def get_nsel(self) -> int: def update_sel( cls, train_data: DeepmdDataSystem, - type_map: Optional[list[str]], + type_map: list[str] | None, local_jdata: dict, - ) -> tuple[dict, Optional[float]]: + ) -> tuple[dict, float | None]: """Update the selection and perform neighbor statistics. Parameters diff --git a/deepmd/pd/model/model/make_model.py b/deepmd/pd/model/model/make_model.py index dafb15e5cd..03aaf621f8 100644 --- a/deepmd/pd/model/model/make_model.py +++ b/deepmd/pd/model/model/make_model.py @@ -1,7 +1,4 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( - Optional, -) import paddle @@ -67,7 +64,7 @@ def __init__( self, *args, # underscore to prevent conflict with normal inputs - atomic_model_: Optional[T_AtomicModel] = None, + atomic_model_: T_AtomicModel | None = None, **kwargs, ) -> None: super().__init__(*args, **kwargs) @@ -129,9 +126,9 @@ def forward_common( self, coord, atype, - box: Optional[paddle.Tensor] = None, - fparam: Optional[paddle.Tensor] = None, - aparam: Optional[paddle.Tensor] = None, + box: paddle.Tensor | None = None, + fparam: paddle.Tensor | None = None, + aparam: paddle.Tensor | None = None, do_atomic_virial: bool = False, ) -> dict[str, paddle.Tensor]: """Return model prediction. @@ -236,11 +233,11 @@ def forward_common_lower( extended_coord, extended_atype, nlist, - mapping: Optional[paddle.Tensor] = None, - fparam: Optional[paddle.Tensor] = None, - aparam: Optional[paddle.Tensor] = None, + mapping: paddle.Tensor | None = None, + fparam: paddle.Tensor | None = None, + aparam: paddle.Tensor | None = None, do_atomic_virial: bool = False, - comm_dict: Optional[list[paddle.Tensor]] = None, + comm_dict: list[paddle.Tensor] | None = None, extra_nlist_sort: bool = False, ): """Return model prediction. Lower interface that takes @@ -306,14 +303,14 @@ def forward_common_lower( def input_type_cast( self, coord: paddle.Tensor, - box: Optional[paddle.Tensor] = None, - fparam: Optional[paddle.Tensor] = None, - aparam: Optional[paddle.Tensor] = None, + box: paddle.Tensor | None = None, + fparam: paddle.Tensor | None = None, + aparam: paddle.Tensor | None = None, ) -> tuple[ paddle.Tensor, - Optional[paddle.Tensor], - Optional[paddle.Tensor], - Optional[paddle.Tensor], + paddle.Tensor | None, + paddle.Tensor | None, + paddle.Tensor | None, str, ]: """Cast the input data to global float type.""" @@ -328,7 +325,7 @@ def input_type_cast( # " does not match" # f" that of the coordinate {input_prec}" # ) - _lst: list[Optional[paddle.Tensor]] = [ + _lst: list[paddle.Tensor | None] = [ vv.astype(coord.dtype) if vv is not None else None for vv in [box, fparam, aparam] ] @@ -484,7 +481,7 @@ def _format_nlist( def do_grad_r( self, - var_name: Optional[str] = None, + var_name: str | None = None, ) -> bool: """Tell if the output variable `var_name` is r_differentiable. if var_name is None, returns if any of the variable is r_differentiable. @@ -493,7 +490,7 @@ def do_grad_r( def do_grad_c( self, - var_name: Optional[str] = None, + var_name: str | None = None, ) -> bool: """Tell if the output variable `var_name` is c_differentiable. if var_name is None, returns if any of the variable is c_differentiable. @@ -594,7 +591,7 @@ def atomic_output_def(self) -> FittingOutputDef: def compute_or_load_stat( self, sampled_func, - stat_file_path: Optional[DPPath] = None, + stat_file_path: DPPath | None = None, ): """Compute or load the statistics.""" return self.atomic_model.compute_or_load_stat(sampled_func, stat_file_path) @@ -627,9 +624,9 @@ def forward( self, coord, atype, - box: Optional[paddle.Tensor] = None, - fparam: Optional[paddle.Tensor] = None, - aparam: Optional[paddle.Tensor] = None, + box: paddle.Tensor | None = None, + fparam: paddle.Tensor | None = None, + aparam: paddle.Tensor | None = None, do_atomic_virial: bool = False, ) -> dict[str, paddle.Tensor]: # directly call the forward_common method when no specific transform rule diff --git a/deepmd/pd/model/model/model.py b/deepmd/pd/model/model/model.py index 0151a9c36b..5027590a9e 100644 --- a/deepmd/pd/model/model/model.py +++ b/deepmd/pd/model/model/model.py @@ -1,7 +1,4 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( - Optional, -) import paddle @@ -23,7 +20,7 @@ def __init__(self, *args, **kwargs): def compute_or_load_stat( self, sampled_func, - stat_file_path: Optional[DPPath] = None, + stat_file_path: DPPath | None = None, ): """ Compute or load the statistics parameters of the model, @@ -46,7 +43,7 @@ def get_model_def_script(self) -> str: """Get the model definition script.""" return self.model_def_script - def get_min_nbor_dist(self) -> Optional[float]: + def get_min_nbor_dist(self) -> float | None: """Get the minimum distance between two atoms.""" return self.min_nbor_dist diff --git a/deepmd/pd/model/network/layernorm.py b/deepmd/pd/model/network/layernorm.py index 7cfffe1b93..d3d00beacd 100644 --- a/deepmd/pd/model/network/layernorm.py +++ b/deepmd/pd/model/network/layernorm.py @@ -1,8 +1,4 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( - Optional, - Union, -) import numpy as np import paddle @@ -45,7 +41,7 @@ def __init__( stddev: float = 1.0, precision: str = DEFAULT_PRECISION, trainable: bool = True, - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, ): super().__init__() self.eps = eps diff --git a/deepmd/pd/model/network/network.py b/deepmd/pd/model/network/network.py index 68053896d1..81e2dad710 100644 --- a/deepmd/pd/model/network/network.py +++ b/deepmd/pd/model/network/network.py @@ -1,8 +1,4 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( - Optional, - Union, -) import paddle import paddle.nn as nn @@ -41,7 +37,7 @@ def __init__( bavg=0.0, stddev=1.0, precision="default", - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, use_econf_tebd=False, use_tebd_bias: bool = False, type_map=None, @@ -162,11 +158,11 @@ def __init__( activation_function: str = "tanh", precision: str = "default", trainable: bool = True, - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, padding: bool = False, use_econf_tebd: bool = False, use_tebd_bias: bool = False, - type_map: Optional[list[str]] = None, + type_map: list[str] | None = None, ) -> None: """Construct a type embedding net.""" super().__init__() diff --git a/deepmd/pd/model/network/utils.py b/deepmd/pd/model/network/utils.py index 9fae72c2cc..cfd6753638 100644 --- a/deepmd/pd/model/network/utils.py +++ b/deepmd/pd/model/network/utils.py @@ -1,7 +1,4 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( - Optional, -) import paddle @@ -10,7 +7,7 @@ def aggregate( data: paddle.Tensor, owners: paddle.Tensor, average: bool = True, - num_owner: Optional[int] = None, + num_owner: int | None = None, ) -> paddle.Tensor: """ Aggregate rows in data by specifying the owners. diff --git a/deepmd/pd/model/task/ener.py b/deepmd/pd/model/task/ener.py index 738990b2d8..3d2025bcc2 100644 --- a/deepmd/pd/model/task/ener.py +++ b/deepmd/pd/model/task/ener.py @@ -1,10 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import copy import logging -from typing import ( - Optional, - Union, -) import paddle @@ -38,7 +34,7 @@ def __init__( ntypes: int, dim_descrpt: int, neuron: list[int] = [128, 128, 128], - bias_atom_e: Optional[paddle.Tensor] = None, + bias_atom_e: paddle.Tensor | None = None, resnet_dt: bool = True, numb_fparam: int = 0, numb_aparam: int = 0, @@ -46,8 +42,8 @@ def __init__( activation_function: str = "tanh", precision: str = DEFAULT_PRECISION, mixed_types: bool = True, - seed: Optional[Union[int, list[int]]] = None, - type_map: Optional[list[str]] = None, + seed: int | list[int] | None = None, + type_map: list[str] | None = None, **kwargs, ): super().__init__( diff --git a/deepmd/pd/model/task/fitting.py b/deepmd/pd/model/task/fitting.py index cb0396e0a2..d92ad6e947 100644 --- a/deepmd/pd/model/task/fitting.py +++ b/deepmd/pd/model/task/fitting.py @@ -3,10 +3,8 @@ from abc import ( abstractmethod, ) -from typing import ( +from collections.abc import ( Callable, - Optional, - Union, ) import numpy as np @@ -77,9 +75,9 @@ def share_params(self, base_class, shared_level, resume=False) -> None: def compute_input_stats( self, - merged: Union[Callable[[], list[dict]], list[dict]], + merged: Callable[[], list[dict]] | list[dict], protection: float = 1e-2, - stat_file_path: Optional[DPPath] = None, + stat_file_path: DPPath | None = None, ) -> None: """ Compute the input statistics (e.g. mean and stddev) for the fittings from packed data. @@ -228,7 +226,7 @@ def __init__( ntypes: int, dim_descrpt: int, neuron: list[int] = [128, 128, 128], - bias_atom_e: Optional[paddle.Tensor] = None, + bias_atom_e: paddle.Tensor | None = None, resnet_dt: bool = True, numb_fparam: int = 0, numb_aparam: int = 0, @@ -236,14 +234,14 @@ def __init__( activation_function: str = "tanh", precision: str = DEFAULT_PRECISION, mixed_types: bool = True, - rcond: Optional[float] = None, - seed: Optional[Union[int, list[int]]] = None, + rcond: float | None = None, + seed: int | list[int] | None = None, exclude_types: list[int] = [], - trainable: Union[bool, list[bool]] = True, - remove_vaccum_contribution: Optional[list[bool]] = None, - type_map: Optional[list[str]] = None, + trainable: bool | list[bool] = True, + remove_vaccum_contribution: list[bool] | None = None, + type_map: list[str] | None = None, use_aparam_as_mask: bool = False, - default_fparam: Optional[list[float]] = None, + default_fparam: list[float] | None = None, **kwargs, ) -> None: super().__init__() @@ -557,11 +555,11 @@ def _forward_common( self, descriptor: paddle.Tensor, atype: paddle.Tensor, - gr: Optional[paddle.Tensor] = None, - g2: Optional[paddle.Tensor] = None, - h2: Optional[paddle.Tensor] = None, - fparam: Optional[paddle.Tensor] = None, - aparam: Optional[paddle.Tensor] = None, + gr: paddle.Tensor | None = None, + g2: paddle.Tensor | None = None, + h2: paddle.Tensor | None = None, + fparam: paddle.Tensor | None = None, + aparam: paddle.Tensor | None = None, ): # cast the input to internal precsion xx = descriptor.astype(self.prec) diff --git a/deepmd/pd/model/task/invar_fitting.py b/deepmd/pd/model/task/invar_fitting.py index 176acdeb20..6dd5b29cb9 100644 --- a/deepmd/pd/model/task/invar_fitting.py +++ b/deepmd/pd/model/task/invar_fitting.py @@ -1,10 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import copy import logging -from typing import ( - Optional, - Union, -) import paddle @@ -91,7 +87,7 @@ def __init__( dim_descrpt: int, dim_out: int, neuron: list[int] = [128, 128, 128], - bias_atom_e: Optional[paddle.Tensor] = None, + bias_atom_e: paddle.Tensor | None = None, resnet_dt: bool = True, numb_fparam: int = 0, numb_aparam: int = 0, @@ -99,11 +95,11 @@ def __init__( activation_function: str = "tanh", precision: str = DEFAULT_PRECISION, mixed_types: bool = True, - rcond: Optional[float] = None, - seed: Optional[Union[int, list[int]]] = None, + rcond: float | None = None, + seed: int | list[int] | None = None, exclude_types: list[int] = [], - atom_ener: Optional[list[Optional[paddle.Tensor]]] = None, - type_map: Optional[list[str]] = None, + atom_ener: list[paddle.Tensor | None] | None = None, + type_map: list[str] | None = None, use_aparam_as_mask: bool = False, **kwargs, ): @@ -167,11 +163,11 @@ def forward( self, descriptor: paddle.Tensor, atype: paddle.Tensor, - gr: Optional[paddle.Tensor] = None, - g2: Optional[paddle.Tensor] = None, - h2: Optional[paddle.Tensor] = None, - fparam: Optional[paddle.Tensor] = None, - aparam: Optional[paddle.Tensor] = None, + gr: paddle.Tensor | None = None, + g2: paddle.Tensor | None = None, + h2: paddle.Tensor | None = None, + fparam: paddle.Tensor | None = None, + aparam: paddle.Tensor | None = None, ): """Based on embedding net output, alculate total energy. diff --git a/deepmd/pd/train/wrapper.py b/deepmd/pd/train/wrapper.py index bd28b17c88..6c8db691c4 100644 --- a/deepmd/pd/train/wrapper.py +++ b/deepmd/pd/train/wrapper.py @@ -7,13 +7,10 @@ from collections import ( OrderedDict, ) -from typing import ( - Union, -) import paddle -_StateDict = Union[dict[str, paddle.Tensor], OrderedDict[str, paddle.Tensor]] +_StateDict = dict[str, paddle.Tensor] | OrderedDict[str, paddle.Tensor] log = logging.getLogger(__name__) diff --git a/deepmd/pd/utils/dataset.py b/deepmd/pd/utils/dataset.py index 1f0533d8fc..fa9106044c 100644 --- a/deepmd/pd/utils/dataset.py +++ b/deepmd/pd/utils/dataset.py @@ -1,10 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( - Optional, -) - from paddle.io import ( Dataset, ) @@ -16,7 +12,7 @@ class DeepmdDataSetForLoader(Dataset): - def __init__(self, system: str, type_map: Optional[list[str]] = None): + def __init__(self, system: str, type_map: list[str] | None = None): """Construct DeePMD-style dataset containing frames cross different systems. Args: diff --git a/deepmd/pd/utils/env_mat_stat.py b/deepmd/pd/utils/env_mat_stat.py index 4ea4850ba7..0e41243924 100644 --- a/deepmd/pd/utils/env_mat_stat.py +++ b/deepmd/pd/utils/env_mat_stat.py @@ -4,7 +4,6 @@ ) from typing import ( TYPE_CHECKING, - Union, ) import numpy as np @@ -77,7 +76,7 @@ def __init__(self, descriptor: "DescriptorBlock"): ) # se_r=1, se_a=4 def iter( - self, data: list[dict[str, Union[paddle.Tensor, list[tuple[int, int]]]]] + self, data: list[dict[str, paddle.Tensor | list[tuple[int, int]]]] ) -> Iterator[dict[str, StatItem]]: """Get the iterator of the environment matrix. diff --git a/deepmd/pd/utils/neighbor_stat.py b/deepmd/pd/utils/neighbor_stat.py index 9bc0079475..f569999bfc 100644 --- a/deepmd/pd/utils/neighbor_stat.py +++ b/deepmd/pd/utils/neighbor_stat.py @@ -2,9 +2,6 @@ from collections.abc import ( Iterator, ) -from typing import ( - Optional, -) import numpy as np import paddle @@ -52,7 +49,7 @@ def forward( self, coord: paddle.Tensor, atype: paddle.Tensor, - cell: Optional[paddle.Tensor], + cell: paddle.Tensor | None, ) -> tuple[paddle.Tensor, paddle.Tensor]: """Calculate the neareest neighbor distance between atoms, maximum nbor size of atoms and the output data range of the environment matrix. @@ -173,7 +170,7 @@ def _execute( self, coord: np.ndarray, atype: np.ndarray, - cell: Optional[np.ndarray], + cell: np.ndarray | None, ): """Execute the operation. diff --git a/deepmd/pd/utils/nlist.py b/deepmd/pd/utils/nlist.py index 9157fba61a..cd041a345c 100644 --- a/deepmd/pd/utils/nlist.py +++ b/deepmd/pd/utils/nlist.py @@ -1,8 +1,4 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( - Optional, - Union, -) import paddle @@ -22,7 +18,7 @@ def extend_input_and_build_neighbor_list( rcut: float, sel: list[int], mixed_types: bool = False, - box: Optional[paddle.Tensor] = None, + box: paddle.Tensor | None = None, ): nframes, nloc = atype.shape[:2] if box is not None: @@ -54,7 +50,7 @@ def build_neighbor_list( atype: paddle.Tensor, nloc: int, rcut: float, - sel: Union[int, list[int]], + sel: int | list[int], distinguish_types: bool = True, ) -> paddle.Tensor: """Build neighbor list for a single frame. keeps nsel neighbors. @@ -181,7 +177,7 @@ def build_directional_neighbor_list( coord_neig: paddle.Tensor, atype_neig: paddle.Tensor, rcut: float, - sel: Union[int, list[int]], + sel: int | list[int], distinguish_types: bool = True, ) -> paddle.Tensor: """Build directional neighbor list. @@ -417,9 +413,9 @@ def build_multiple_neighbor_list( def extend_coord_with_ghosts( coord: paddle.Tensor, atype: paddle.Tensor, - cell: Optional[paddle.Tensor], + cell: paddle.Tensor | None, rcut: float, - cell_cpu: Optional[paddle.Tensor] = None, + cell_cpu: paddle.Tensor | None = None, ) -> tuple[paddle.Tensor, paddle.Tensor, paddle.Tensor]: """Extend the coordinates of the atoms by appending peridoc images. The number of images is large enough to ensure all the neighbors diff --git a/deepmd/pd/utils/stat.py b/deepmd/pd/utils/stat.py index 4132d0a5f7..3642d309b8 100644 --- a/deepmd/pd/utils/stat.py +++ b/deepmd/pd/utils/stat.py @@ -3,11 +3,11 @@ from collections import ( defaultdict, ) +from collections.abc import ( + Callable, +) from typing import ( Any, - Callable, - Optional, - Union, ) import numpy as np @@ -97,7 +97,7 @@ def make_stat_input( def _restore_from_file( stat_file_path: DPPath, keys: list[str] = ["energy"], -) -> Optional[dict]: +) -> dict | None: if stat_file_path is None: return None, None stat_files = [stat_file_path / f"bias_atom_{kk}" for kk in keys] @@ -159,7 +159,7 @@ def _post_process_stat( def _compute_model_predict( - sampled: Union[Callable[[], list[dict]], list[dict]], + sampled: Callable[[], list[dict]] | list[dict], keys: list[str], model_forward: Callable[..., paddle.Tensor], ) -> dict[str, list[paddle.Tensor]]: @@ -199,8 +199,8 @@ def model_forward_auto_batch_size(*args, **kwargs): def _make_preset_out_bias( ntypes: int, - ibias: list[Optional[np.ndarray]], -) -> Optional[np.ndarray]: + ibias: list[np.ndarray | None], +) -> np.ndarray | None: """Make preset out bias. output: @@ -223,9 +223,9 @@ def _make_preset_out_bias( def _fill_stat_with_global( - atomic_stat: Union[np.ndarray, None], + atomic_stat: np.ndarray | None, global_stat: np.ndarray, -) -> Union[np.ndarray, None]: +) -> np.ndarray | None: """This function is used to fill atomic stat with global stat. Parameters @@ -249,13 +249,13 @@ def _fill_stat_with_global( def compute_output_stats( - merged: Union[Callable[[], list[dict]], list[dict]], + merged: Callable[[], list[dict]] | list[dict], ntypes: int, - keys: Union[str, list[str]] = ["energy"], - stat_file_path: Optional[DPPath] = None, - rcond: Optional[float] = None, - preset_bias: Optional[dict[str, list[Optional[np.ndarray]]]] = None, - model_forward: Optional[Callable[..., paddle.Tensor]] = None, + keys: str | list[str] = ["energy"], + stat_file_path: DPPath | None = None, + rcond: float | None = None, + preset_bias: dict[str, list[np.ndarray | None]] | None = None, + model_forward: Callable[..., paddle.Tensor] | None = None, stats_distinguish_types: bool = True, intensive: bool = False, ) -> dict[str, Any]: @@ -420,9 +420,9 @@ def compute_output_stats_global( sampled: list[dict], ntypes: int, keys: list[str], - rcond: Optional[float] = None, - preset_bias: Optional[dict[str, list[Optional[paddle.Tensor]]]] = None, - model_pred: Optional[dict[str, np.ndarray]] = None, + rcond: float | None = None, + preset_bias: dict[str, list[paddle.Tensor | None]] | None = None, + model_pred: dict[str, np.ndarray] | None = None, stats_distinguish_types: bool = True, intensive: bool = False, ) -> tuple[dict[str, np.ndarray], dict[str, np.ndarray]]: @@ -554,7 +554,7 @@ def compute_output_stats_atomic( sampled: list[dict], ntypes: int, keys: list[str], - model_pred: Optional[dict[str, np.ndarray]] = None, + model_pred: dict[str, np.ndarray] | None = None, ) -> tuple[dict[str, np.ndarray], dict[str, np.ndarray]]: # get label dict from sample; for each key, only picking the system with atomic labels. outputs = { diff --git a/deepmd/pt/entrypoints/compress.py b/deepmd/pt/entrypoints/compress.py index 3a4efaeea5..83f2999971 100644 --- a/deepmd/pt/entrypoints/compress.py +++ b/deepmd/pt/entrypoints/compress.py @@ -1,9 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import json import logging -from typing import ( - Optional, -) import torch @@ -35,7 +32,7 @@ def enable_compression( stride: float = 0.01, extrapolate: int = 5, check_frequency: int = -1, - training_script: Optional[str] = None, + training_script: str | None = None, ) -> None: saved_model = torch.jit.load(input_file, map_location="cpu") model_def_script = json.loads(saved_model.model_def_script) diff --git a/deepmd/pt/entrypoints/main.py b/deepmd/pt/entrypoints/main.py index 06a7603cc0..d349d519ba 100644 --- a/deepmd/pt/entrypoints/main.py +++ b/deepmd/pt/entrypoints/main.py @@ -9,8 +9,6 @@ ) from typing import ( Any, - Optional, - Union, ) import h5py @@ -97,13 +95,13 @@ def get_trainer( config: dict[str, Any], - init_model: Optional[str] = None, - restart_model: Optional[str] = None, - finetune_model: Optional[str] = None, + init_model: str | None = None, + restart_model: str | None = None, + finetune_model: str | None = None, force_load: bool = False, - init_frz_model: Optional[str] = None, - shared_links: Optional[dict[str, Any]] = None, - finetune_links: Optional[dict[str, Any]] = None, + init_frz_model: str | None = None, + shared_links: dict[str, Any] | None = None, + finetune_links: dict[str, Any] | None = None, ) -> training.Trainer: multi_task = "model_dict" in config.get("model", {}) @@ -111,8 +109,8 @@ def prepare_trainer_input_single( model_params_single: dict[str, Any], data_dict_single: dict[str, Any], rank: int = 0, - seed: Optional[int] = None, - ) -> tuple[DpLoaderSet, Optional[DpLoaderSet], Optional[DPPath]]: + seed: int | None = None, + ) -> tuple[DpLoaderSet, DpLoaderSet | None, DPPath | None]: training_dataset_params = data_dict_single["training_data"] validation_dataset_params = data_dict_single.get("validation_data", None) validation_systems = ( @@ -245,10 +243,10 @@ def get_backend_info(self) -> dict: def train( input_file: str, - init_model: Optional[str], - restart: Optional[str], - finetune: Optional[str], - init_frz_model: Optional[str], + init_model: str | None, + restart: str | None, + finetune: str | None, + init_frz_model: str | None, model_branch: str, skip_neighbor_stat: bool = False, use_pretrain_script: bool = False, @@ -372,7 +370,7 @@ def train( def freeze( model: str, output: str = "frozen_model.pth", - head: Optional[str] = None, + head: str | None = None, ) -> None: model = inference.Tester(model, head=head).model model.eval() @@ -389,12 +387,12 @@ def freeze( def change_bias( input_file: str, mode: str = "change", - bias_value: Optional[list] = None, - datafile: Optional[str] = None, + bias_value: list | None = None, + datafile: str | None = None, system: str = ".", numb_batch: int = 0, - model_branch: Optional[str] = None, - output: Optional[str] = None, + model_branch: str | None = None, + output: str | None = None, ) -> None: if input_file.endswith(".pt"): old_state_dict = torch.load( @@ -518,7 +516,7 @@ def change_bias( @record -def main(args: Optional[Union[list[str], argparse.Namespace]] = None) -> None: +def main(args: list[str] | argparse.Namespace | None = None) -> None: if not isinstance(args, argparse.Namespace): FLAGS = parse_args(args=args) else: diff --git a/deepmd/pt/infer/deep_eval.py b/deepmd/pt/infer/deep_eval.py index d4800bb4ca..02cdb702ef 100644 --- a/deepmd/pt/infer/deep_eval.py +++ b/deepmd/pt/infer/deep_eval.py @@ -1,12 +1,13 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import json import logging +from collections.abc import ( + Callable, +) from typing import ( TYPE_CHECKING, Any, - Callable, Optional, - Union, ) import numpy as np @@ -108,9 +109,9 @@ def __init__( model_file: str, output_def: ModelOutputDef, *args: Any, - auto_batch_size: Union[bool, int, AutoBatchSize] = True, + auto_batch_size: bool | int | AutoBatchSize = True, neighbor_list: Optional["ase.neighborlist.NewPrimitiveNeighborList"] = None, - head: Optional[Union[str, int]] = None, + head: str | int | None = None, no_jit: bool = False, **kwargs: Any, ) -> None: @@ -306,11 +307,11 @@ def get_model_branch(self) -> tuple[dict[str, str], dict[str, dict[str, Any]]]: def eval( self, coords: np.ndarray, - cells: Optional[np.ndarray], + cells: np.ndarray | None, atom_types: np.ndarray, atomic: bool = False, - fparam: Optional[np.ndarray] = None, - aparam: Optional[np.ndarray] = None, + fparam: np.ndarray | None = None, + aparam: np.ndarray | None = None, **kwargs: Any, ) -> dict[str, np.ndarray]: """Evaluate the energy, force and virial by using this DP. @@ -467,10 +468,10 @@ def _get_natoms_and_nframes( def _eval_model( self, coords: np.ndarray, - cells: Optional[np.ndarray], + cells: np.ndarray | None, atom_types: np.ndarray, - fparam: Optional[np.ndarray], - aparam: Optional[np.ndarray], + fparam: np.ndarray | None, + aparam: np.ndarray | None, request_defs: list[OutputVariableDef], ) -> tuple[np.ndarray, ...]: model = self.dp.to(DEVICE) @@ -544,11 +545,11 @@ def _eval_model( def _eval_model_spin( self, coords: np.ndarray, - cells: Optional[np.ndarray], + cells: np.ndarray | None, atom_types: np.ndarray, spins: np.ndarray, - fparam: Optional[np.ndarray], - aparam: Optional[np.ndarray], + fparam: np.ndarray | None, + aparam: np.ndarray | None, request_defs: list[OutputVariableDef], ) -> tuple[np.ndarray, ...]: model = self.dp.to(DEVICE) @@ -740,10 +741,10 @@ def get_model(self) -> "BaseModel": def eval_descriptor( self, coords: np.ndarray, - cells: Optional[np.ndarray], + cells: np.ndarray | None, atom_types: np.ndarray, - fparam: Optional[np.ndarray] = None, - aparam: Optional[np.ndarray] = None, + fparam: np.ndarray | None = None, + aparam: np.ndarray | None = None, **kwargs: Any, ) -> np.ndarray: """Evaluate descriptors by using this DP. @@ -795,10 +796,10 @@ def eval_descriptor( def eval_fitting_last_layer( self, coords: np.ndarray, - cells: Optional[np.ndarray], + cells: np.ndarray | None, atom_types: np.ndarray, - fparam: Optional[np.ndarray] = None, - aparam: Optional[np.ndarray] = None, + fparam: np.ndarray | None = None, + aparam: np.ndarray | None = None, **kwargs: Any, ) -> np.ndarray: """Evaluate fitting before last layer by using this DP. diff --git a/deepmd/pt/infer/inference.py b/deepmd/pt/infer/inference.py index ac11d160aa..4c49abeef8 100644 --- a/deepmd/pt/infer/inference.py +++ b/deepmd/pt/infer/inference.py @@ -3,10 +3,6 @@ from copy import ( deepcopy, ) -from typing import ( - Optional, - Union, -) import torch @@ -29,8 +25,8 @@ class Tester: def __init__( self, - model_ckpt: Union[str, torch.nn.Module], - head: Optional[str] = None, + model_ckpt: str | torch.nn.Module, + head: str | None = None, ) -> None: """Construct a DeePMD tester. diff --git a/deepmd/pt/loss/ener.py b/deepmd/pt/loss/ener.py index cccdc8949e..1778806343 100644 --- a/deepmd/pt/loss/ener.py +++ b/deepmd/pt/loss/ener.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( Any, - Optional, ) import torch @@ -49,7 +48,7 @@ def __init__( limit_pref_ae: float = 0.0, start_pref_pf: float = 0.0, limit_pref_pf: float = 0.0, - relative_f: Optional[float] = None, + relative_f: float | None = None, enable_atom_ener_coeff: bool = False, start_pref_gf: float = 0.0, limit_pref_gf: float = 0.0, diff --git a/deepmd/pt/loss/loss.py b/deepmd/pt/loss/loss.py index 13cad6f59b..51eb7a9319 100644 --- a/deepmd/pt/loss/loss.py +++ b/deepmd/pt/loss/loss.py @@ -6,7 +6,6 @@ from typing import ( Any, NoReturn, - Union, ) import torch @@ -30,7 +29,7 @@ def forward( model: torch.nn.Module, label: dict[str, torch.Tensor], natoms: int, - learning_rate: Union[float, torch.Tensor], + learning_rate: float | torch.Tensor, ) -> NoReturn: """Return loss .""" raise NotImplementedError diff --git a/deepmd/pt/loss/property.py b/deepmd/pt/loss/property.py index 1cd842650d..189bcb2a4a 100644 --- a/deepmd/pt/loss/property.py +++ b/deepmd/pt/loss/property.py @@ -2,7 +2,6 @@ import logging from typing import ( Any, - Union, ) import torch @@ -29,8 +28,8 @@ def __init__( loss_func: str = "smooth_mae", metric: list[str] = ["mae"], beta: float = 1.00, - out_bias: Union[list, None] = None, - out_std: Union[list, None] = None, + out_bias: list | None = None, + out_std: list | None = None, intensive: bool = False, **kwargs: Any, ) -> None: diff --git a/deepmd/pt/model/atomic_model/base_atomic_model.py b/deepmd/pt/model/atomic_model/base_atomic_model.py index 6b786a796f..0ccf539757 100644 --- a/deepmd/pt/model/atomic_model/base_atomic_model.py +++ b/deepmd/pt/model/atomic_model/base_atomic_model.py @@ -1,11 +1,12 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import logging -from typing import ( +from collections.abc import ( Callable, +) +from typing import ( NoReturn, Optional, - Union, ) import numpy as np @@ -77,8 +78,8 @@ def __init__( type_map: list[str], atom_exclude_types: list[int] = [], pair_exclude_types: list[tuple[int, int]] = [], - rcond: Optional[float] = None, - preset_out_bias: Optional[dict[str, np.ndarray]] = None, + rcond: float | None = None, + preset_out_bias: dict[str, np.ndarray] | None = None, data_stat_protect: float = 1e-2, ) -> None: torch.nn.Module.__init__(self) @@ -203,10 +204,10 @@ def forward_common_atomic( extended_coord: torch.Tensor, extended_atype: torch.Tensor, nlist: torch.Tensor, - mapping: Optional[torch.Tensor] = None, - fparam: Optional[torch.Tensor] = None, - aparam: Optional[torch.Tensor] = None, - comm_dict: Optional[dict[str, torch.Tensor]] = None, + mapping: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + comm_dict: dict[str, torch.Tensor] | None = None, ) -> dict[str, torch.Tensor]: """Common interface for atomic inference. @@ -284,10 +285,10 @@ def forward( extended_coord: torch.Tensor, extended_atype: torch.Tensor, nlist: torch.Tensor, - mapping: Optional[torch.Tensor] = None, - fparam: Optional[torch.Tensor] = None, - aparam: Optional[torch.Tensor] = None, - comm_dict: Optional[dict[str, torch.Tensor]] = None, + mapping: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + comm_dict: dict[str, torch.Tensor] | None = None, ) -> dict[str, torch.Tensor]: return self.forward_common_atomic( extended_coord, @@ -367,8 +368,8 @@ def deserialize(cls, data: dict) -> "BaseAtomicModel": def compute_or_load_stat( self, - merged: Union[Callable[[], list[dict]], list[dict]], - stat_file_path: Optional[DPPath] = None, + merged: Callable[[], list[dict]] | list[dict], + stat_file_path: DPPath | None = None, compute_or_load_out_stat: bool = True, ) -> NoReturn: """ @@ -394,8 +395,8 @@ def compute_or_load_stat( def compute_or_load_out_stat( self, - merged: Union[Callable[[], list[dict]], list[dict]], - stat_file_path: Optional[DPPath] = None, + merged: Callable[[], list[dict]] | list[dict], + stat_file_path: DPPath | None = None, ) -> None: """ Compute the output statistics (e.g. energy bias) for the fitting net from packed data. @@ -444,8 +445,8 @@ def apply_out_stat( def change_out_bias( self, - sample_merged: Union[Callable[[], list[dict]], list[dict]], - stat_file_path: Optional[DPPath] = None, + sample_merged: Callable[[], list[dict]] | list[dict], + stat_file_path: DPPath | None = None, bias_adjust_mode: str = "change-by-statistic", ) -> None: """Change the output bias according to the input data and the pretrained model. @@ -495,7 +496,7 @@ def change_out_bias( def compute_fitting_input_stat( self, - sample_merged: Union[Callable[[], list[dict]], list[dict]], + sample_merged: Callable[[], list[dict]] | list[dict], ) -> None: """Compute the input statistics (e.g. mean and stddev) for the atomic model from packed data. @@ -517,9 +518,9 @@ def _get_forward_wrapper_func(self) -> Callable[..., torch.Tensor]: def model_forward( coord: torch.Tensor, atype: torch.Tensor, - box: Optional[torch.Tensor], - fparam: Optional[torch.Tensor] = None, - aparam: Optional[torch.Tensor] = None, + box: torch.Tensor | None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: with ( torch.no_grad() diff --git a/deepmd/pt/model/atomic_model/dp_atomic_model.py b/deepmd/pt/model/atomic_model/dp_atomic_model.py index fdff461efa..4ebef0106c 100644 --- a/deepmd/pt/model/atomic_model/dp_atomic_model.py +++ b/deepmd/pt/model/atomic_model/dp_atomic_model.py @@ -1,11 +1,12 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import functools import logging +from collections.abc import ( + Callable, +) from typing import ( Any, - Callable, Optional, - Union, ) import torch @@ -222,10 +223,10 @@ def forward_atomic( extended_coord: torch.Tensor, extended_atype: torch.Tensor, nlist: torch.Tensor, - mapping: Optional[torch.Tensor] = None, - fparam: Optional[torch.Tensor] = None, - aparam: Optional[torch.Tensor] = None, - comm_dict: Optional[dict[str, torch.Tensor]] = None, + mapping: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + comm_dict: dict[str, torch.Tensor] | None = None, ) -> dict[str, torch.Tensor]: """Return atomic prediction. @@ -289,7 +290,7 @@ def get_out_bias(self) -> torch.Tensor: def compute_or_load_stat( self, sampled_func: Callable[[], list[dict]], - stat_file_path: Optional[DPPath] = None, + stat_file_path: DPPath | None = None, compute_or_load_out_stat: bool = True, ) -> None: """ @@ -344,8 +345,8 @@ def wrapped_sampler() -> list[dict]: def compute_fitting_input_stat( self, - sample_merged: Union[Callable[[], list[dict]], list[dict]], - stat_file_path: Optional[DPPath] = None, + sample_merged: Callable[[], list[dict]] | list[dict], + stat_file_path: DPPath | None = None, ) -> None: """Compute the input statistics (e.g. mean and stddev) for the fittings from packed data. @@ -375,7 +376,7 @@ def has_default_fparam(self) -> bool: """Check if the model has default frame parameters.""" return self.fitting_net.has_default_fparam() - def get_default_fparam(self) -> Optional[torch.Tensor]: + def get_default_fparam(self) -> torch.Tensor | None: return self.fitting_net.get_default_fparam() def get_dim_aparam(self) -> int: diff --git a/deepmd/pt/model/atomic_model/linear_atomic_model.py b/deepmd/pt/model/atomic_model/linear_atomic_model.py index b510448ec3..96b3baf6ec 100644 --- a/deepmd/pt/model/atomic_model/linear_atomic_model.py +++ b/deepmd/pt/model/atomic_model/linear_atomic_model.py @@ -1,10 +1,11 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import functools +from collections.abc import ( + Callable, +) from typing import ( Any, - Callable, Optional, - Union, ) import torch @@ -57,7 +58,7 @@ def __init__( self, models: list[BaseAtomicModel], type_map: list[str], - weights: Optional[Union[str, list[float]]] = "mean", + weights: str | list[float] | None = "mean", **kwargs: Any, ) -> None: super().__init__(type_map, **kwargs) @@ -232,10 +233,10 @@ def forward_atomic( extended_coord: torch.Tensor, extended_atype: torch.Tensor, nlist: torch.Tensor, - mapping: Optional[torch.Tensor] = None, - fparam: Optional[torch.Tensor] = None, - aparam: Optional[torch.Tensor] = None, - comm_dict: Optional[dict[str, torch.Tensor]] = None, + mapping: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + comm_dict: dict[str, torch.Tensor] | None = None, ) -> dict[str, torch.Tensor]: """Return atomic prediction. @@ -476,7 +477,7 @@ def is_aparam_nall(self) -> bool: def compute_or_load_stat( self, sampled_func: Callable[[], list[dict[str, Any]]], - stat_file_path: Optional[DPPath] = None, + stat_file_path: DPPath | None = None, compute_or_load_out_stat: bool = True, ) -> None: """ @@ -551,7 +552,7 @@ def __init__( sw_rmin: float, sw_rmax: float, type_map: list[str], - smin_alpha: Optional[float] = 0.1, + smin_alpha: float | None = 0.1, **kwargs: Any, ) -> None: models = [dp_model, zbl_model] diff --git a/deepmd/pt/model/atomic_model/pairtab_atomic_model.py b/deepmd/pt/model/atomic_model/pairtab_atomic_model.py index b022e6bfc9..6933fdc19a 100644 --- a/deepmd/pt/model/atomic_model/pairtab_atomic_model.py +++ b/deepmd/pt/model/atomic_model/pairtab_atomic_model.py @@ -1,9 +1,10 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +from collections.abc import ( + Callable, +) from typing import ( Any, - Callable, Optional, - Union, ) import torch @@ -67,7 +68,7 @@ def __init__( self, tab_file: str, rcut: float, - sel: Union[int, list[int]], + sel: int | list[int], type_map: list[str], **kwargs: Any, ) -> None: @@ -227,8 +228,8 @@ def deserialize(cls, data: dict[str, Any]) -> "PairTabAtomicModel": def compute_or_load_stat( self, - sampled_func: Union[Callable[[], list[dict]], list[dict]], - stat_file_path: Optional[DPPath] = None, + sampled_func: Callable[[], list[dict]] | list[dict], + stat_file_path: DPPath | None = None, compute_or_load_out_stat: bool = True, ) -> None: """ @@ -258,11 +259,11 @@ def forward_atomic( extended_coord: torch.Tensor, extended_atype: torch.Tensor, nlist: torch.Tensor, - mapping: Optional[torch.Tensor] = None, - fparam: Optional[torch.Tensor] = None, - aparam: Optional[torch.Tensor] = None, + mapping: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, do_atomic_virial: bool = False, - comm_dict: Optional[dict[str, torch.Tensor]] = None, + comm_dict: dict[str, torch.Tensor] | None = None, ) -> dict[str, torch.Tensor]: nframes, nloc, nnei = nlist.shape extended_coord = extended_coord.view(nframes, -1, 3) diff --git a/deepmd/pt/model/descriptor/descriptor.py b/deepmd/pt/model/descriptor/descriptor.py index c1a3529ae0..2a6e857a6f 100644 --- a/deepmd/pt/model/descriptor/descriptor.py +++ b/deepmd/pt/model/descriptor/descriptor.py @@ -4,12 +4,12 @@ ABC, abstractmethod, ) +from collections.abc import ( + Callable, +) from typing import ( Any, - Callable, NoReturn, - Optional, - Union, ) import torch @@ -102,8 +102,8 @@ def get_env_protection(self) -> float: def compute_input_stats( self, - merged: Union[Callable[[], list[dict]], list[dict]], - path: Optional[DPPath] = None, + merged: Callable[[], list[dict]] | list[dict], + path: DPPath | None = None, ) -> NoReturn: """ Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data. @@ -178,15 +178,15 @@ def forward( nlist: torch.Tensor, extended_coord: torch.Tensor, extended_atype: torch.Tensor, - extended_atype_embd: Optional[torch.Tensor] = None, - mapping: Optional[torch.Tensor] = None, - type_embedding: Optional[torch.Tensor] = None, + extended_atype_embd: torch.Tensor | None = None, + mapping: torch.Tensor | None = None, + type_embedding: torch.Tensor | None = None, ) -> tuple[ torch.Tensor, - Optional[torch.Tensor], - Optional[torch.Tensor], - Optional[torch.Tensor], - Optional[torch.Tensor], + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, ]: """Calculate DescriptorBlock.""" pass @@ -211,7 +211,7 @@ def make_default_type_embedding( def extend_descrpt_stat( des: DescriptorBlock, type_map: list[str], - des_with_stat: Optional[DescriptorBlock] = None, + des_with_stat: DescriptorBlock | None = None, ) -> None: r""" Extend the statistics of a descriptor block with types from newly provided `type_map`. diff --git a/deepmd/pt/model/descriptor/dpa1.py b/deepmd/pt/model/descriptor/dpa1.py index 7f600ccc2e..359bf2f084 100644 --- a/deepmd/pt/model/descriptor/dpa1.py +++ b/deepmd/pt/model/descriptor/dpa1.py @@ -1,10 +1,10 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import warnings +from collections.abc import ( + Callable, +) from typing import ( Any, - Callable, - Optional, - Union, ) import torch @@ -221,7 +221,7 @@ def __init__( self, rcut: float, rcut_smth: float, - sel: Union[list[int], int], + sel: list[int] | int, ntypes: int, neuron: list = [25, 50, 100], axis_neuron: int = 16, @@ -239,21 +239,21 @@ def __init__( env_protection: float = 0.0, scaling_factor: int = 1.0, normalize: bool = True, - temperature: Optional[float] = None, + temperature: float | None = None, concat_output_tebd: bool = True, trainable: bool = True, trainable_ln: bool = True, - ln_eps: Optional[float] = 1e-5, + ln_eps: float | None = 1e-5, smooth_type_embedding: bool = True, type_one_side: bool = False, - stripped_type_embedding: Optional[bool] = None, - seed: Optional[Union[int, list[int]]] = None, + stripped_type_embedding: bool | None = None, + seed: int | list[int] | None = None, use_econf_tebd: bool = False, use_tebd_bias: bool = False, - type_map: Optional[list[str]] = None, + type_map: list[str] | None = None, # not implemented - spin: Optional[Any] = None, - type: Optional[str] = None, + spin: Any | None = None, + type: str | None = None, ) -> None: super().__init__() # Ensure compatibility with the deprecated stripped_type_embedding option. @@ -418,8 +418,8 @@ def dim_emb(self) -> int: def compute_input_stats( self, - merged: Union[Callable[[], list[dict]], list[dict]], - path: Optional[DPPath] = None, + merged: Callable[[], list[dict]] | list[dict], + path: DPPath | None = None, ) -> None: """ Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data. @@ -453,7 +453,7 @@ def get_stat_mean_and_stddev(self) -> tuple[torch.Tensor, torch.Tensor]: return self.se_atten.mean, self.se_atten.stddev def change_type_map( - self, type_map: list[str], model_with_new_type_stat: Optional[Any] = None + self, type_map: list[str], model_with_new_type_stat: Any | None = None ) -> None: """Change the type related params to new ones, according to `type_map` and the original one in the model. If there are new types in `type_map`, statistics will be updated accordingly to `model_with_new_type_stat` for these new types. @@ -664,14 +664,14 @@ def forward( extended_coord: torch.Tensor, extended_atype: torch.Tensor, nlist: torch.Tensor, - mapping: Optional[torch.Tensor] = None, - comm_dict: Optional[dict[str, torch.Tensor]] = None, + mapping: torch.Tensor | None = None, + comm_dict: dict[str, torch.Tensor] | None = None, ) -> tuple[ torch.Tensor, - Optional[torch.Tensor], - Optional[torch.Tensor], - Optional[torch.Tensor], - Optional[torch.Tensor], + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, ]: """Compute the descriptor. @@ -741,9 +741,9 @@ def forward( def update_sel( cls, train_data: DeepmdDataSystem, - type_map: Optional[list[str]], + type_map: list[str] | None, local_jdata: dict, - ) -> tuple[dict, Optional[float]]: + ) -> tuple[dict, float | None]: """Update the selection and perform neighbor statistics. Parameters diff --git a/deepmd/pt/model/descriptor/dpa2.py b/deepmd/pt/model/descriptor/dpa2.py index 583f18f2be..30d1987b97 100644 --- a/deepmd/pt/model/descriptor/dpa2.py +++ b/deepmd/pt/model/descriptor/dpa2.py @@ -1,10 +1,10 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import warnings +from collections.abc import ( + Callable, +) from typing import ( Any, - Callable, - Optional, - Union, ) import torch @@ -86,9 +86,9 @@ def __init__( self, ntypes: int, # args for repinit - repinit: Union[RepinitArgs, dict], + repinit: RepinitArgs | dict, # args for repformer - repformer: Union[RepformerArgs, dict], + repformer: RepformerArgs | dict, # kwargs for descriptor concat_output_tebd: bool = True, precision: str = "float64", @@ -96,11 +96,11 @@ def __init__( exclude_types: list[tuple[int, int]] = [], env_protection: float = 0.0, trainable: bool = True, - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, add_tebd_to_repinit_out: bool = False, use_econf_tebd: bool = False, use_tebd_bias: bool = False, - type_map: Optional[list[str]] = None, + type_map: list[str] | None = None, ) -> None: r"""The DPA-2 descriptor[1]_. @@ -426,7 +426,7 @@ def share_params( raise NotImplementedError def change_type_map( - self, type_map: list[str], model_with_new_type_stat: Optional[Any] = None + self, type_map: list[str], model_with_new_type_stat: Any | None = None ) -> None: """Change the type related params to new ones, according to `type_map` and the original one in the model. If there are new types in `type_map`, statistics will be updated accordingly to `model_with_new_type_stat` for these new types. @@ -491,8 +491,8 @@ def dim_emb(self) -> int: def compute_input_stats( self, - merged: Union[Callable[[], list[dict]], list[dict]], - path: Optional[DPPath] = None, + merged: Callable[[], list[dict]] | list[dict], + path: DPPath | None = None, ) -> None: """ Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data. @@ -713,14 +713,14 @@ def forward( extended_coord: torch.Tensor, extended_atype: torch.Tensor, nlist: torch.Tensor, - mapping: Optional[torch.Tensor] = None, - comm_dict: Optional[dict[str, torch.Tensor]] = None, + mapping: torch.Tensor | None = None, + comm_dict: dict[str, torch.Tensor] | None = None, ) -> tuple[ torch.Tensor, - Optional[torch.Tensor], - Optional[torch.Tensor], - Optional[torch.Tensor], - Optional[torch.Tensor], + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, ]: """Compute the descriptor. @@ -842,9 +842,9 @@ def forward( def update_sel( cls, train_data: DeepmdDataSystem, - type_map: Optional[list[str]], + type_map: list[str] | None, local_jdata: dict, - ) -> tuple[dict, Optional[float]]: + ) -> tuple[dict, float | None]: """Update the selection and perform neighbor statistics. Parameters diff --git a/deepmd/pt/model/descriptor/dpa3.py b/deepmd/pt/model/descriptor/dpa3.py index 2de7851a51..136527123e 100644 --- a/deepmd/pt/model/descriptor/dpa3.py +++ b/deepmd/pt/model/descriptor/dpa3.py @@ -1,9 +1,9 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +from collections.abc import ( + Callable, +) from typing import ( Any, - Callable, - Optional, - Union, ) import torch @@ -107,7 +107,7 @@ def __init__( self, ntypes: int, # args for repflow - repflow: Union[RepFlowArgs, dict], + repflow: RepFlowArgs | dict, # kwargs for descriptor concat_output_tebd: bool = False, activation_function: str = "silu", @@ -115,11 +115,11 @@ def __init__( exclude_types: list[tuple[int, int]] = [], env_protection: float = 0.0, trainable: bool = True, - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, use_econf_tebd: bool = False, use_tebd_bias: bool = False, use_loc_mapping: bool = True, - type_map: Optional[list[str]] = None, + type_map: list[str] | None = None, ) -> None: super().__init__() @@ -299,7 +299,7 @@ def share_params( raise NotImplementedError def change_type_map( - self, type_map: list[str], model_with_new_type_stat: Optional[Any] = None + self, type_map: list[str], model_with_new_type_stat: Any | None = None ) -> None: """Change the type related params to new ones, according to `type_map` and the original one in the model. If there are new types in `type_map`, statistics will be updated accordingly to `model_with_new_type_stat` for these new types. @@ -338,8 +338,8 @@ def dim_emb(self) -> int: def compute_input_stats( self, - merged: Union[Callable[[], list[dict]], list[dict]], - path: Optional[DPPath] = None, + merged: Callable[[], list[dict]] | list[dict], + path: DPPath | None = None, ) -> None: """ Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data. @@ -453,14 +453,14 @@ def forward( extended_coord: torch.Tensor, extended_atype: torch.Tensor, nlist: torch.Tensor, - mapping: Optional[torch.Tensor] = None, - comm_dict: Optional[dict[str, torch.Tensor]] = None, + mapping: torch.Tensor | None = None, + comm_dict: dict[str, torch.Tensor] | None = None, ) -> tuple[ torch.Tensor, - Optional[torch.Tensor], - Optional[torch.Tensor], - Optional[torch.Tensor], - Optional[torch.Tensor], + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, ]: """Compute the descriptor. @@ -532,9 +532,9 @@ def forward( def update_sel( cls, train_data: DeepmdDataSystem, - type_map: Optional[list[str]], + type_map: list[str] | None, local_jdata: dict, - ) -> tuple[dict, Optional[float]]: + ) -> tuple[dict, float | None]: """Update the selection and perform neighbor statistics. Parameters diff --git a/deepmd/pt/model/descriptor/hybrid.py b/deepmd/pt/model/descriptor/hybrid.py index 545fba7019..0f001bc4c8 100644 --- a/deepmd/pt/model/descriptor/hybrid.py +++ b/deepmd/pt/model/descriptor/hybrid.py @@ -3,7 +3,6 @@ from typing import ( Any, Optional, - Union, ) import numpy as np @@ -44,7 +43,7 @@ class DescrptHybrid(BaseDescriptor, torch.nn.Module): def __init__( self, - list: list[Union[BaseDescriptor, dict[str, Any]]], + list: list[BaseDescriptor | dict[str, Any]], **kwargs: Any, ) -> None: super().__init__() @@ -200,7 +199,7 @@ def change_type_map( ) def compute_input_stats( - self, merged: list[dict], path: Optional[DPPath] = None + self, merged: list[dict], path: DPPath | None = None ) -> None: """Update mean and stddev for descriptor elements.""" for descrpt in self.descrpt_list: @@ -208,8 +207,8 @@ def compute_input_stats( def set_stat_mean_and_stddev( self, - mean: list[Union[torch.Tensor, list[torch.Tensor]]], - stddev: list[Union[torch.Tensor, list[torch.Tensor]]], + mean: list[torch.Tensor | list[torch.Tensor]], + stddev: list[torch.Tensor | list[torch.Tensor]], ) -> None: """Update mean and stddev for descriptor.""" for ii, descrpt in enumerate(self.descrpt_list): @@ -218,8 +217,8 @@ def set_stat_mean_and_stddev( def get_stat_mean_and_stddev( self, ) -> tuple[ - list[Union[torch.Tensor, list[torch.Tensor]]], - list[Union[torch.Tensor, list[torch.Tensor]]], + list[torch.Tensor | list[torch.Tensor]], + list[torch.Tensor | list[torch.Tensor]], ]: """Get mean and stddev for descriptor.""" mean_list = [] @@ -267,14 +266,14 @@ def forward( coord_ext: torch.Tensor, atype_ext: torch.Tensor, nlist: torch.Tensor, - mapping: Optional[torch.Tensor] = None, - comm_dict: Optional[dict[str, torch.Tensor]] = None, + mapping: torch.Tensor | None = None, + comm_dict: dict[str, torch.Tensor] | None = None, ) -> tuple[ torch.Tensor, - Optional[torch.Tensor], - Optional[torch.Tensor], - Optional[torch.Tensor], - Optional[torch.Tensor], + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, ]: """Compute the descriptor. @@ -309,9 +308,9 @@ def forward( """ out_descriptor = [] out_gr = [] - out_g2: Optional[torch.Tensor] = None - out_h2: Optional[torch.Tensor] = None - out_sw: Optional[torch.Tensor] = None + out_g2: torch.Tensor | None = None + out_h2: torch.Tensor | None = None + out_sw: torch.Tensor | None = None if self.sel_no_mixed_types is not None: nl_distinguish_types = nlist_distinguish_types( nlist, @@ -344,9 +343,9 @@ def forward( def update_sel( cls, train_data: DeepmdDataSystem, - type_map: Optional[list[str]], + type_map: list[str] | None, local_jdata: dict, - ) -> tuple[dict, Optional[float]]: + ) -> tuple[dict, float | None]: """Update the selection and perform neighbor statistics. Parameters diff --git a/deepmd/pt/model/descriptor/repflow_layer.py b/deepmd/pt/model/descriptor/repflow_layer.py index 62145958c8..338f48b060 100644 --- a/deepmd/pt/model/descriptor/repflow_layer.py +++ b/deepmd/pt/model/descriptor/repflow_layer.py @@ -1,8 +1,4 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( - Optional, - Union, -) import torch import torch.nn as nn @@ -63,7 +59,7 @@ def __init__( update_residual: float = 0.1, update_residual_init: str = "const", precision: str = "float64", - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, trainable: bool = True, ) -> None: super().__init__() diff --git a/deepmd/pt/model/descriptor/repflows.py b/deepmd/pt/model/descriptor/repflows.py index 69b5e3b593..433897860f 100644 --- a/deepmd/pt/model/descriptor/repflows.py +++ b/deepmd/pt/model/descriptor/repflows.py @@ -1,9 +1,9 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +from collections.abc import ( + Callable, +) from typing import ( Any, - Callable, - Optional, - Union, ) import torch @@ -221,7 +221,7 @@ def __init__( sel_reduce_factor: float = 10.0, use_loc_mapping: bool = True, optim_update: bool = True, - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, trainable: bool = True, ) -> None: super().__init__() @@ -436,15 +436,15 @@ def forward( nlist: torch.Tensor, extended_coord: torch.Tensor, extended_atype: torch.Tensor, - extended_atype_embd: Optional[torch.Tensor] = None, - mapping: Optional[torch.Tensor] = None, - comm_dict: Optional[dict[str, torch.Tensor]] = None, + extended_atype_embd: torch.Tensor | None = None, + mapping: torch.Tensor | None = None, + comm_dict: dict[str, torch.Tensor] | None = None, ) -> tuple[ torch.Tensor, - Optional[torch.Tensor], - Optional[torch.Tensor], - Optional[torch.Tensor], - Optional[torch.Tensor], + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, ]: parallel_mode = comm_dict is not None if not parallel_mode: @@ -679,8 +679,8 @@ def forward( def compute_input_stats( self, - merged: Union[Callable[[], list[dict]], list[dict]], - path: Optional[DPPath] = None, + merged: Callable[[], list[dict]] | list[dict], + path: DPPath | None = None, ) -> None: """ Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data. diff --git a/deepmd/pt/model/descriptor/repformer_layer.py b/deepmd/pt/model/descriptor/repformer_layer.py index 32012af92d..fd72dbc0a3 100644 --- a/deepmd/pt/model/descriptor/repformer_layer.py +++ b/deepmd/pt/model/descriptor/repformer_layer.py @@ -1,8 +1,4 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( - Optional, - Union, -) import torch import torch.nn as nn @@ -43,7 +39,7 @@ def get_residual( _mode: str = "norm", trainable: bool = True, precision: str = "float64", - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, ) -> torch.Tensor: r""" Get residual tensor for one update vector. @@ -159,7 +155,7 @@ def __init__( smooth: bool = True, attnw_shift: float = 20.0, precision: str = "float64", - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, trainable: bool = True, ) -> None: """Return neighbor-wise multi-head self-attention maps, with gate mechanism.""" @@ -286,7 +282,7 @@ def __init__( input_dim: int, head_num: int, precision: str = "float64", - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, trainable: bool = True, ) -> None: super().__init__() @@ -374,7 +370,7 @@ def __init__( input_dim: int, head_num: int, precision: str = "float64", - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, trainable: bool = True, ) -> None: super().__init__() @@ -453,7 +449,7 @@ def __init__( smooth: bool = True, attnw_shift: float = 20.0, precision: str = "float64", - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, trainable: bool = True, ) -> None: super().__init__() @@ -612,11 +608,11 @@ def __init__( smooth: bool = True, precision: str = "float64", trainable_ln: bool = True, - ln_eps: Optional[float] = 1e-5, + ln_eps: float | None = 1e-5, use_sqrt_nnei: bool = True, g1_out_conv: bool = True, g1_out_mlp: bool = True, - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, trainable: bool = True, ) -> None: super().__init__() diff --git a/deepmd/pt/model/descriptor/repformers.py b/deepmd/pt/model/descriptor/repformers.py index 2c383640f1..69f2cc4eaa 100644 --- a/deepmd/pt/model/descriptor/repformers.py +++ b/deepmd/pt/model/descriptor/repformers.py @@ -1,9 +1,9 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +from collections.abc import ( + Callable, +) from typing import ( Any, - Callable, - Optional, - Union, ) import torch @@ -107,8 +107,8 @@ def __init__( env_protection: float = 0.0, precision: str = "float64", trainable_ln: bool = True, - ln_eps: Optional[float] = 1e-5, - seed: Optional[Union[int, list[int]]] = None, + ln_eps: float | None = 1e-5, + seed: int | list[int] | None = None, use_sqrt_nnei: bool = True, g1_out_conv: bool = True, g1_out_mlp: bool = True, @@ -396,16 +396,16 @@ def forward( nlist: torch.Tensor, extended_coord: torch.Tensor, extended_atype: torch.Tensor, - extended_atype_embd: Optional[torch.Tensor] = None, - mapping: Optional[torch.Tensor] = None, - type_embedding: Optional[torch.Tensor] = None, - comm_dict: Optional[dict[str, torch.Tensor]] = None, + extended_atype_embd: torch.Tensor | None = None, + mapping: torch.Tensor | None = None, + type_embedding: torch.Tensor | None = None, + comm_dict: dict[str, torch.Tensor] | None = None, ) -> tuple[ torch.Tensor, - Optional[torch.Tensor], - Optional[torch.Tensor], - Optional[torch.Tensor], - Optional[torch.Tensor], + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, ]: if comm_dict is None: assert mapping is not None @@ -546,8 +546,8 @@ def forward( def compute_input_stats( self, - merged: Union[Callable[[], list[dict]], list[dict]], - path: Optional[DPPath] = None, + merged: Callable[[], list[dict]] | list[dict], + path: DPPath | None = None, ) -> None: """ Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data. diff --git a/deepmd/pt/model/descriptor/se_a.py b/deepmd/pt/model/descriptor/se_a.py index 17fa6a830e..59c165ddb0 100644 --- a/deepmd/pt/model/descriptor/se_a.py +++ b/deepmd/pt/model/descriptor/se_a.py @@ -1,11 +1,11 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import itertools +from collections.abc import ( + Callable, +) from typing import ( Any, - Callable, ClassVar, - Optional, - Union, ) import numpy as np @@ -96,7 +96,7 @@ def __init__( self, rcut: float, rcut_smth: float, - sel: Union[list[int], int], + sel: list[int] | int, neuron: list[int] = [25, 50, 100], axis_neuron: int = 16, set_davg_zero: bool = False, @@ -107,11 +107,11 @@ def __init__( env_protection: float = 0.0, type_one_side: bool = True, trainable: bool = True, - seed: Optional[Union[int, list[int]]] = None, - ntypes: Optional[int] = None, # to be compat with input - type_map: Optional[list[str]] = None, + seed: int | list[int] | None = None, + ntypes: int | None = None, # to be compat with input + type_map: list[str] | None = None, # not implemented - spin: Optional[Any] = None, + spin: Any | None = None, ) -> None: del ntypes if spin is not None: @@ -213,7 +213,7 @@ def dim_out(self) -> int: return self.sea.dim_out def change_type_map( - self, type_map: list[str], model_with_new_type_stat: Optional[Any] = None + self, type_map: list[str], model_with_new_type_stat: Any | None = None ) -> None: """Change the type related params to new ones, according to `type_map` and the original one in the model. If there are new types in `type_map`, statistics will be updated accordingly to `model_with_new_type_stat` for these new types. @@ -226,8 +226,8 @@ def change_type_map( def compute_input_stats( self, - merged: Union[Callable[[], list[dict]], list[dict]], - path: Optional[DPPath] = None, + merged: Callable[[], list[dict]] | list[dict], + path: DPPath | None = None, ) -> None: """ Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data. @@ -306,14 +306,14 @@ def forward( coord_ext: torch.Tensor, atype_ext: torch.Tensor, nlist: torch.Tensor, - mapping: Optional[torch.Tensor] = None, - comm_dict: Optional[dict[str, torch.Tensor]] = None, + mapping: torch.Tensor | None = None, + comm_dict: dict[str, torch.Tensor] | None = None, ) -> tuple[ torch.Tensor, - Optional[torch.Tensor], - Optional[torch.Tensor], - Optional[torch.Tensor], - Optional[torch.Tensor], + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, ]: """Compute the descriptor. @@ -431,9 +431,9 @@ def t_cvt(xx: Any) -> torch.Tensor: def update_sel( cls, train_data: DeepmdDataSystem, - type_map: Optional[list[str]], + type_map: list[str] | None, local_jdata: dict, - ) -> tuple[dict, Optional[float]]: + ) -> tuple[dict, float | None]: """Update the selection and perform neighbor statistics. Parameters @@ -468,7 +468,7 @@ def __init__( self, rcut: float, rcut_smth: float, - sel: Union[int, list[int]], + sel: int | list[int], neuron: list[int] = [25, 50, 100], axis_neuron: int = 16, set_davg_zero: bool = False, @@ -479,7 +479,7 @@ def __init__( env_protection: float = 0.0, type_one_side: bool = True, trainable: bool = True, - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, **kwargs: Any, ) -> None: """Construct an embedding net of type `se_a`. @@ -640,8 +640,8 @@ def __getitem__(self, key: str) -> torch.Tensor: def compute_input_stats( self, - merged: Union[Callable[[], list[dict]], list[dict]], - path: Optional[DPPath] = None, + merged: Callable[[], list[dict]] | list[dict], + path: DPPath | None = None, ) -> None: """ Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data. @@ -699,7 +699,7 @@ def reinit_exclude( def enable_compression( self, table_data: dict[str, torch.Tensor], - table_config: list[Union[int, float]], + table_config: list[int | float], lower: dict[str, int], upper: dict[str, int], ) -> None: @@ -737,15 +737,15 @@ def forward( nlist: torch.Tensor, extended_coord: torch.Tensor, extended_atype: torch.Tensor, - extended_atype_embd: Optional[torch.Tensor] = None, - mapping: Optional[torch.Tensor] = None, - type_embedding: Optional[torch.Tensor] = None, + extended_atype_embd: torch.Tensor | None = None, + mapping: torch.Tensor | None = None, + type_embedding: torch.Tensor | None = None, ) -> tuple[ torch.Tensor, - Optional[torch.Tensor], - Optional[torch.Tensor], - Optional[torch.Tensor], - Optional[torch.Tensor], + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, ]: """Calculate decoded embedding for each atom. diff --git a/deepmd/pt/model/descriptor/se_atten.py b/deepmd/pt/model/descriptor/se_atten.py index 7c1de6146a..32c8830b1e 100644 --- a/deepmd/pt/model/descriptor/se_atten.py +++ b/deepmd/pt/model/descriptor/se_atten.py @@ -1,9 +1,9 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +from collections.abc import ( + Callable, +) from typing import ( Any, - Callable, - Optional, - Union, ) import torch @@ -79,7 +79,7 @@ def __init__( self, rcut: float, rcut_smth: float, - sel: Union[list[int], int], + sel: list[int] | int, ntypes: int, neuron: list = [25, 50, 100], axis_neuron: int = 16, @@ -95,15 +95,15 @@ def __init__( resnet_dt: bool = False, scaling_factor: float = 1.0, normalize: bool = True, - temperature: Optional[float] = None, + temperature: float | None = None, smooth: bool = True, type_one_side: bool = False, exclude_types: list[tuple[int, int]] = [], env_protection: float = 0.0, trainable_ln: bool = True, - ln_eps: Optional[float] = 1e-5, - seed: Optional[Union[int, list[int]]] = None, - type: Optional[str] = None, + ln_eps: float | None = 1e-5, + seed: int | list[int] | None = None, + type: str | None = None, trainable: bool = True, ) -> None: r"""Construct an embedding net of type `se_atten`. @@ -375,8 +375,8 @@ def dim_emb(self) -> int: def compute_input_stats( self, - merged: Union[Callable[[], list[dict]], list[dict]], - path: Optional[DPPath] = None, + merged: Callable[[], list[dict]] | list[dict], + path: DPPath | None = None, ) -> None: """ Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data. @@ -512,15 +512,15 @@ def forward( nlist: torch.Tensor, extended_coord: torch.Tensor, extended_atype: torch.Tensor, - extended_atype_embd: Optional[torch.Tensor] = None, - mapping: Optional[torch.Tensor] = None, - type_embedding: Optional[torch.Tensor] = None, + extended_atype_embd: torch.Tensor | None = None, + mapping: torch.Tensor | None = None, + type_embedding: torch.Tensor | None = None, ) -> tuple[ torch.Tensor, - Optional[torch.Tensor], - Optional[torch.Tensor], - Optional[torch.Tensor], - Optional[torch.Tensor], + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, ]: """Compute the descriptor. @@ -749,12 +749,12 @@ def __init__( do_mask: bool = False, scaling_factor: float = 1.0, normalize: bool = True, - temperature: Optional[float] = None, + temperature: float | None = None, trainable_ln: bool = True, ln_eps: float = 1e-5, smooth: bool = True, precision: str = DEFAULT_PRECISION, - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, trainable: bool = True, ) -> None: """Construct a neighbor-wise attention net.""" @@ -800,8 +800,8 @@ def forward( self, input_G: torch.Tensor, nei_mask: torch.Tensor, - input_r: Optional[torch.Tensor] = None, - sw: Optional[torch.Tensor] = None, + input_r: torch.Tensor | None = None, + sw: torch.Tensor | None = None, ) -> torch.Tensor: """Compute the multi-layer gated self-attention. @@ -894,12 +894,12 @@ def __init__( do_mask: bool = False, scaling_factor: float = 1.0, normalize: bool = True, - temperature: Optional[float] = None, + temperature: float | None = None, smooth: bool = True, trainable_ln: bool = True, ln_eps: float = 1e-5, precision: str = DEFAULT_PRECISION, - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, trainable: bool = True, ) -> None: """Construct a neighbor-wise attention layer.""" @@ -942,8 +942,8 @@ def forward( self, x: torch.Tensor, nei_mask: torch.Tensor, - input_r: Optional[torch.Tensor] = None, - sw: Optional[torch.Tensor] = None, + input_r: torch.Tensor | None = None, + sw: torch.Tensor | None = None, ) -> torch.Tensor: residual = x x, _ = self.attention_layer(x, nei_mask, input_r=input_r, sw=sw) @@ -1004,11 +1004,11 @@ def __init__( do_mask: bool = False, scaling_factor: float = 1.0, normalize: bool = True, - temperature: Optional[float] = None, + temperature: float | None = None, bias: bool = True, smooth: bool = True, precision: str = DEFAULT_PRECISION, - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, trainable: bool = True, ) -> None: """Construct a multi-head neighbor-wise attention net.""" @@ -1060,8 +1060,8 @@ def forward( self, query: torch.Tensor, nei_mask: torch.Tensor, - input_r: Optional[torch.Tensor] = None, - sw: Optional[torch.Tensor] = None, + input_r: torch.Tensor | None = None, + sw: torch.Tensor | None = None, attnw_shift: float = 20.0, ) -> tuple[torch.Tensor, torch.Tensor]: """Compute the multi-head gated self-attention. diff --git a/deepmd/pt/model/descriptor/se_atten_v2.py b/deepmd/pt/model/descriptor/se_atten_v2.py index 5377d919b0..5fa031b8bf 100644 --- a/deepmd/pt/model/descriptor/se_atten_v2.py +++ b/deepmd/pt/model/descriptor/se_atten_v2.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( Any, - Optional, - Union, ) import torch @@ -41,7 +39,7 @@ def __init__( self, rcut: float, rcut_smth: float, - sel: Union[list[int], int], + sel: list[int] | int, ntypes: int, neuron: list = [25, 50, 100], axis_neuron: int = 16, @@ -58,20 +56,20 @@ def __init__( env_protection: float = 0.0, scaling_factor: int = 1.0, normalize: bool = True, - temperature: Optional[float] = None, + temperature: float | None = None, concat_output_tebd: bool = True, trainable: bool = True, trainable_ln: bool = True, - ln_eps: Optional[float] = 1e-5, + ln_eps: float | None = 1e-5, type_one_side: bool = False, - stripped_type_embedding: Optional[bool] = None, - seed: Optional[Union[int, list[int]]] = None, + stripped_type_embedding: bool | None = None, + seed: int | list[int] | None = None, use_econf_tebd: bool = False, use_tebd_bias: bool = False, - type_map: Optional[list[str]] = None, + type_map: list[str] | None = None, # not implemented - spin: Optional[Any] = None, - type: Optional[str] = None, + spin: Any | None = None, + type: str | None = None, ) -> None: r"""Construct smooth version of embedding net of type `se_atten_v2`. diff --git a/deepmd/pt/model/descriptor/se_r.py b/deepmd/pt/model/descriptor/se_r.py index 294323a48c..1d5a8fc1a8 100644 --- a/deepmd/pt/model/descriptor/se_r.py +++ b/deepmd/pt/model/descriptor/se_r.py @@ -1,9 +1,9 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +from collections.abc import ( + Callable, +) from typing import ( Any, - Callable, - Optional, - Union, ) import numpy as np @@ -84,7 +84,7 @@ def __init__( self, rcut: float, rcut_smth: float, - sel: Union[list[int], int], + sel: list[int] | int, neuron: list[int] = [25, 50, 100], set_davg_zero: bool = False, activation_function: str = "tanh", @@ -93,8 +93,8 @@ def __init__( exclude_types: list[tuple[int, int]] = [], env_protection: float = 0.0, trainable: bool = True, - seed: Optional[Union[int, list[int]]] = None, - type_map: Optional[list[str]] = None, + seed: int | list[int] | None = None, + type_map: list[str] | None = None, **kwargs: Any, ) -> None: super().__init__() @@ -271,7 +271,7 @@ def share_params( raise NotImplementedError def change_type_map( - self, type_map: list[str], model_with_new_type_stat: Optional[Any] = None + self, type_map: list[str], model_with_new_type_stat: Any | None = None ) -> None: """Change the type related params to new ones, according to `type_map` and the original one in the model. If there are new types in `type_map`, statistics will be updated accordingly to `model_with_new_type_stat` for these new types. @@ -284,8 +284,8 @@ def change_type_map( def compute_input_stats( self, - merged: Union[Callable[[], list[dict]], list[dict]], - path: Optional[DPPath] = None, + merged: Callable[[], list[dict]] | list[dict], + path: DPPath | None = None, ) -> None: """ Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data. @@ -425,14 +425,14 @@ def forward( coord_ext: torch.Tensor, atype_ext: torch.Tensor, nlist: torch.Tensor, - mapping: Optional[torch.Tensor] = None, - comm_dict: Optional[dict[str, torch.Tensor]] = None, + mapping: torch.Tensor | None = None, + comm_dict: dict[str, torch.Tensor] | None = None, ) -> tuple[ torch.Tensor, - Optional[torch.Tensor], - Optional[torch.Tensor], - Optional[torch.Tensor], - Optional[torch.Tensor], + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, ]: """Compute the descriptor. @@ -596,9 +596,9 @@ def t_cvt(xx: Any) -> torch.Tensor: def update_sel( cls, train_data: DeepmdDataSystem, - type_map: Optional[list[str]], + type_map: list[str] | None, local_jdata: dict, - ) -> tuple[dict, Optional[float]]: + ) -> tuple[dict, float | None]: """Update the selection and perform neighbor statistics. Parameters diff --git a/deepmd/pt/model/descriptor/se_t.py b/deepmd/pt/model/descriptor/se_t.py index c489d0be06..9771d8fe6f 100644 --- a/deepmd/pt/model/descriptor/se_t.py +++ b/deepmd/pt/model/descriptor/se_t.py @@ -1,11 +1,11 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import itertools +from collections.abc import ( + Callable, +) from typing import ( Any, - Callable, ClassVar, - Optional, - Union, ) import numpy as np @@ -143,11 +143,11 @@ def __init__( exclude_types: list[tuple[int, int]] = [], precision: str = "float64", trainable: bool = True, - seed: Optional[Union[int, list[int]]] = None, - type_map: Optional[list[str]] = None, - ntypes: Optional[int] = None, # to be compat with input + seed: int | list[int] | None = None, + type_map: list[str] | None = None, + ntypes: int | None = None, # to be compat with input # not implemented - spin: Optional[dict] = None, + spin: dict | None = None, ) -> None: del ntypes if spin is not None: @@ -247,7 +247,7 @@ def dim_out(self) -> int: return self.seat.dim_out def change_type_map( - self, type_map: list[str], model_with_new_type_stat: Optional[Any] = None + self, type_map: list[str], model_with_new_type_stat: Any | None = None ) -> None: """Change the type related params to new ones, according to `type_map` and the original one in the model. If there are new types in `type_map`, statistics will be updated accordingly to `model_with_new_type_stat` for these new types. @@ -260,8 +260,8 @@ def change_type_map( def compute_input_stats( self, - merged: Union[Callable[[], list[dict]], list[dict]], - path: Optional[DPPath] = None, + merged: Callable[[], list[dict]] | list[dict], + path: DPPath | None = None, ) -> None: """ Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data. @@ -341,14 +341,14 @@ def forward( coord_ext: torch.Tensor, atype_ext: torch.Tensor, nlist: torch.Tensor, - mapping: Optional[torch.Tensor] = None, - comm_dict: Optional[dict[str, torch.Tensor]] = None, + mapping: torch.Tensor | None = None, + comm_dict: dict[str, torch.Tensor] | None = None, ) -> tuple[ torch.Tensor, - Optional[torch.Tensor], - Optional[torch.Tensor], - Optional[torch.Tensor], - Optional[torch.Tensor], + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, ]: """Compute the descriptor. @@ -460,9 +460,9 @@ def t_cvt(xx: Any) -> torch.Tensor: def update_sel( cls, train_data: DeepmdDataSystem, - type_map: Optional[list[str]], + type_map: list[str] | None, local_jdata: dict, - ) -> tuple[dict, Optional[float]]: + ) -> tuple[dict, float | None]: """Update the selection and perform neighbor statistics. Parameters @@ -506,7 +506,7 @@ def __init__( exclude_types: list[tuple[int, int]] = [], precision: str = "float64", trainable: bool = True, - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, ) -> None: r"""Construct an embedding net of type `se_e3`. @@ -684,8 +684,8 @@ def __getitem__(self, key: str) -> Any: def compute_input_stats( self, - merged: Union[Callable[[], list[dict]], list[dict]], - path: Optional[DPPath] = None, + merged: Callable[[], list[dict]] | list[dict], + path: DPPath | None = None, ) -> None: """ Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data. @@ -774,15 +774,15 @@ def forward( nlist: torch.Tensor, extended_coord: torch.Tensor, extended_atype: torch.Tensor, - extended_atype_embd: Optional[torch.Tensor] = None, - mapping: Optional[torch.Tensor] = None, - type_embedding: Optional[torch.Tensor] = None, + extended_atype_embd: torch.Tensor | None = None, + mapping: torch.Tensor | None = None, + type_embedding: torch.Tensor | None = None, ) -> tuple[ torch.Tensor, - Optional[torch.Tensor], - Optional[torch.Tensor], - Optional[torch.Tensor], - Optional[torch.Tensor], + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, ]: """Compute the descriptor. diff --git a/deepmd/pt/model/descriptor/se_t_tebd.py b/deepmd/pt/model/descriptor/se_t_tebd.py index 8225d8e4af..d639cc94bc 100644 --- a/deepmd/pt/model/descriptor/se_t_tebd.py +++ b/deepmd/pt/model/descriptor/se_t_tebd.py @@ -1,9 +1,9 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +from collections.abc import ( + Callable, +) from typing import ( Any, - Callable, - Optional, - Union, ) import torch @@ -132,7 +132,7 @@ def __init__( self, rcut: float, rcut_smth: float, - sel: Union[list[int], int], + sel: list[int] | int, ntypes: int, neuron: list = [2, 4, 8], tebd_dim: int = 8, @@ -144,8 +144,8 @@ def __init__( exclude_types: list[tuple[int, int]] = [], precision: str = "float64", trainable: bool = True, - seed: Optional[Union[int, list[int]]] = None, - type_map: Optional[list[str]] = None, + seed: int | list[int] | None = None, + type_map: list[str] | None = None, concat_output_tebd: bool = True, use_econf_tebd: bool = False, use_tebd_bias: bool = False, @@ -286,8 +286,8 @@ def dim_emb(self) -> int: def compute_input_stats( self, - merged: Union[Callable[[], list[dict]], list[dict]], - path: Optional[DPPath] = None, + merged: Callable[[], list[dict]] | list[dict], + path: DPPath | None = None, ) -> None: """ Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data. @@ -321,7 +321,7 @@ def get_stat_mean_and_stddev(self) -> tuple[torch.Tensor, torch.Tensor]: return self.se_ttebd.mean, self.se_ttebd.stddev def change_type_map( - self, type_map: list[str], model_with_new_type_stat: Optional[Any] = None + self, type_map: list[str], model_with_new_type_stat: Any | None = None ) -> None: """Change the type related params to new ones, according to `type_map` and the original one in the model. If there are new types in `type_map`, statistics will be updated accordingly to `model_with_new_type_stat` for these new types. @@ -421,14 +421,14 @@ def forward( extended_coord: torch.Tensor, extended_atype: torch.Tensor, nlist: torch.Tensor, - mapping: Optional[torch.Tensor] = None, - comm_dict: Optional[dict[str, torch.Tensor]] = None, + mapping: torch.Tensor | None = None, + comm_dict: dict[str, torch.Tensor] | None = None, ) -> tuple[ torch.Tensor, - Optional[torch.Tensor], - Optional[torch.Tensor], - Optional[torch.Tensor], - Optional[torch.Tensor], + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, ]: """Compute the descriptor. @@ -496,9 +496,9 @@ def forward( def update_sel( cls, train_data: DeepmdDataSystem, - type_map: Optional[list[str]], + type_map: list[str] | None, local_jdata: dict, - ) -> tuple[dict, Optional[float]]: + ) -> tuple[dict, float | None]: """Update the selection and perform neighbor statistics. Parameters @@ -611,7 +611,7 @@ def __init__( self, rcut: float, rcut_smth: float, - sel: Union[list[int], int], + sel: list[int] | int, ntypes: int, neuron: list = [25, 50, 100], tebd_dim: int = 8, @@ -623,7 +623,7 @@ def __init__( exclude_types: list[tuple[int, int]] = [], env_protection: float = 0.0, smooth: bool = True, - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, trainable: bool = True, ) -> None: super().__init__() @@ -790,8 +790,8 @@ def dim_emb(self) -> int: def compute_input_stats( self, - merged: Union[Callable[[], list[dict]], list[dict]], - path: Optional[DPPath] = None, + merged: Callable[[], list[dict]] | list[dict], + path: DPPath | None = None, ) -> None: """ Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data. @@ -851,15 +851,15 @@ def forward( nlist: torch.Tensor, extended_coord: torch.Tensor, extended_atype: torch.Tensor, - extended_atype_embd: Optional[torch.Tensor] = None, - mapping: Optional[torch.Tensor] = None, - type_embedding: Optional[torch.Tensor] = None, + extended_atype_embd: torch.Tensor | None = None, + mapping: torch.Tensor | None = None, + type_embedding: torch.Tensor | None = None, ) -> tuple[ torch.Tensor, - Optional[torch.Tensor], - Optional[torch.Tensor], - Optional[torch.Tensor], - Optional[torch.Tensor], + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, ]: """Compute the descriptor. diff --git a/deepmd/pt/model/model/__init__.py b/deepmd/pt/model/model/__init__.py index 1be46e084a..b8f7b171d4 100644 --- a/deepmd/pt/model/model/__init__.py +++ b/deepmd/pt/model/model/__init__.py @@ -15,7 +15,6 @@ import json from typing import ( Any, - Optional, ) import numpy as np @@ -210,7 +209,7 @@ def get_zbl_model(model_params: dict) -> DPZBLModel: return model -def _can_be_converted_to_float(value: Any) -> Optional[bool]: +def _can_be_converted_to_float(value: Any) -> bool | None: try: float(value) return True @@ -220,8 +219,8 @@ def _can_be_converted_to_float(value: Any) -> Optional[bool]: def _convert_preset_out_bias_to_array( - preset_out_bias: Optional[dict], type_map: list[str] -) -> Optional[dict]: + preset_out_bias: dict | None, type_map: list[str] +) -> dict | None: if preset_out_bias is not None: for kk in preset_out_bias: if len(preset_out_bias[kk]) != len(type_map): diff --git a/deepmd/pt/model/model/dipole_model.py b/deepmd/pt/model/model/dipole_model.py index de089e7de7..c6813ce079 100644 --- a/deepmd/pt/model/model/dipole_model.py +++ b/deepmd/pt/model/model/dipole_model.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( Any, - Optional, ) import torch @@ -57,9 +56,9 @@ def forward( self, coord: torch.Tensor, atype: torch.Tensor, - box: Optional[torch.Tensor] = None, - fparam: Optional[torch.Tensor] = None, - aparam: Optional[torch.Tensor] = None, + box: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, do_atomic_virial: bool = False, ) -> dict[str, torch.Tensor]: model_ret = self.forward_common( @@ -95,11 +94,11 @@ def forward_lower( extended_coord: torch.Tensor, extended_atype: torch.Tensor, nlist: torch.Tensor, - mapping: Optional[torch.Tensor] = None, - fparam: Optional[torch.Tensor] = None, - aparam: Optional[torch.Tensor] = None, + mapping: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, do_atomic_virial: bool = False, - comm_dict: Optional[dict[str, torch.Tensor]] = None, + comm_dict: dict[str, torch.Tensor] | None = None, ) -> dict[str, torch.Tensor]: model_ret = self.forward_common_lower( extended_coord, diff --git a/deepmd/pt/model/model/dos_model.py b/deepmd/pt/model/model/dos_model.py index a68735984f..75f89c141a 100644 --- a/deepmd/pt/model/model/dos_model.py +++ b/deepmd/pt/model/model/dos_model.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( Any, - Optional, ) import torch @@ -49,9 +48,9 @@ def forward( self, coord: torch.Tensor, atype: torch.Tensor, - box: Optional[torch.Tensor] = None, - fparam: Optional[torch.Tensor] = None, - aparam: Optional[torch.Tensor] = None, + box: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, do_atomic_virial: bool = False, ) -> dict[str, torch.Tensor]: model_ret = self.forward_common( @@ -85,11 +84,11 @@ def forward_lower( extended_coord: torch.Tensor, extended_atype: torch.Tensor, nlist: torch.Tensor, - mapping: Optional[torch.Tensor] = None, - fparam: Optional[torch.Tensor] = None, - aparam: Optional[torch.Tensor] = None, + mapping: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, do_atomic_virial: bool = False, - comm_dict: Optional[dict[str, torch.Tensor]] = None, + comm_dict: dict[str, torch.Tensor] | None = None, ) -> dict[str, torch.Tensor]: model_ret = self.forward_common_lower( extended_coord, diff --git a/deepmd/pt/model/model/dp_linear_model.py b/deepmd/pt/model/model/dp_linear_model.py index b71c8a10c3..b43f849258 100644 --- a/deepmd/pt/model/model/dp_linear_model.py +++ b/deepmd/pt/model/model/dp_linear_model.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( Any, - Optional, ) import torch @@ -62,9 +61,9 @@ def forward( self, coord: torch.Tensor, atype: torch.Tensor, - box: Optional[torch.Tensor] = None, - fparam: Optional[torch.Tensor] = None, - aparam: Optional[torch.Tensor] = None, + box: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, do_atomic_virial: bool = False, ) -> dict[str, torch.Tensor]: model_ret = self.forward_common( @@ -97,11 +96,11 @@ def forward_lower( extended_coord: torch.Tensor, extended_atype: torch.Tensor, nlist: torch.Tensor, - mapping: Optional[torch.Tensor] = None, - fparam: Optional[torch.Tensor] = None, - aparam: Optional[torch.Tensor] = None, + mapping: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, do_atomic_virial: bool = False, - comm_dict: Optional[dict[str, torch.Tensor]] = None, + comm_dict: dict[str, torch.Tensor] | None = None, ) -> dict[str, torch.Tensor]: model_ret = self.forward_common_lower( extended_coord, @@ -135,9 +134,9 @@ def forward_lower( def update_sel( cls, train_data: DeepmdDataSystem, - type_map: Optional[list[str]], + type_map: list[str] | None, local_jdata: dict, - ) -> tuple[dict, Optional[float]]: + ) -> tuple[dict, float | None]: """Update the selection and perform neighbor statistics. Parameters diff --git a/deepmd/pt/model/model/dp_model.py b/deepmd/pt/model/model/dp_model.py index 875dc0dca0..a8b5b55584 100644 --- a/deepmd/pt/model/model/dp_model.py +++ b/deepmd/pt/model/model/dp_model.py @@ -1,7 +1,4 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( - Optional, -) import torch @@ -20,9 +17,9 @@ class DPModelCommon: def update_sel( cls, train_data: DeepmdDataSystem, - type_map: Optional[list[str]], + type_map: list[str] | None, local_jdata: dict, - ) -> tuple[dict, Optional[float]]: + ) -> tuple[dict, float | None]: """Update the selection and perform neighbor statistics. Parameters diff --git a/deepmd/pt/model/model/dp_zbl_model.py b/deepmd/pt/model/model/dp_zbl_model.py index 7f84d8abec..d533cbe125 100644 --- a/deepmd/pt/model/model/dp_zbl_model.py +++ b/deepmd/pt/model/model/dp_zbl_model.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( Any, - Optional, ) import torch @@ -59,9 +58,9 @@ def forward( self, coord: torch.Tensor, atype: torch.Tensor, - box: Optional[torch.Tensor] = None, - fparam: Optional[torch.Tensor] = None, - aparam: Optional[torch.Tensor] = None, + box: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, do_atomic_virial: bool = False, ) -> dict[str, torch.Tensor]: model_ret = self.forward_common( @@ -94,11 +93,11 @@ def forward_lower( extended_coord: torch.Tensor, extended_atype: torch.Tensor, nlist: torch.Tensor, - mapping: Optional[torch.Tensor] = None, - fparam: Optional[torch.Tensor] = None, - aparam: Optional[torch.Tensor] = None, + mapping: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, do_atomic_virial: bool = False, - comm_dict: Optional[dict[str, torch.Tensor]] = None, + comm_dict: dict[str, torch.Tensor] | None = None, ) -> dict[str, torch.Tensor]: model_ret = self.forward_common_lower( extended_coord, @@ -132,9 +131,9 @@ def forward_lower( def update_sel( cls, train_data: DeepmdDataSystem, - type_map: Optional[list[str]], + type_map: list[str] | None, local_jdata: dict, - ) -> tuple[dict, Optional[float]]: + ) -> tuple[dict, float | None]: """Update the selection and perform neighbor statistics. Parameters diff --git a/deepmd/pt/model/model/ener_model.py b/deepmd/pt/model/model/ener_model.py index dfe68d537f..8f8a3cbad7 100644 --- a/deepmd/pt/model/model/ener_model.py +++ b/deepmd/pt/model/model/ener_model.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( Any, - Optional, ) import torch @@ -95,9 +94,9 @@ def forward( self, coord: torch.Tensor, atype: torch.Tensor, - box: Optional[torch.Tensor] = None, - fparam: Optional[torch.Tensor] = None, - aparam: Optional[torch.Tensor] = None, + box: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, do_atomic_virial: bool = False, ) -> dict[str, torch.Tensor]: model_ret = self.forward_common( @@ -137,11 +136,11 @@ def forward_lower( extended_coord: torch.Tensor, extended_atype: torch.Tensor, nlist: torch.Tensor, - mapping: Optional[torch.Tensor] = None, - fparam: Optional[torch.Tensor] = None, - aparam: Optional[torch.Tensor] = None, + mapping: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, do_atomic_virial: bool = False, - comm_dict: Optional[dict[str, torch.Tensor]] = None, + comm_dict: dict[str, torch.Tensor] | None = None, ) -> dict[str, torch.Tensor]: model_ret = self.forward_common_lower( extended_coord, diff --git a/deepmd/pt/model/model/frozen.py b/deepmd/pt/model/model/frozen.py index 2a63b093db..402e78e95c 100644 --- a/deepmd/pt/model/model/frozen.py +++ b/deepmd/pt/model/model/frozen.py @@ -4,7 +4,6 @@ from typing import ( Any, NoReturn, - Optional, ) import torch @@ -119,9 +118,9 @@ def forward( self, coord: torch.Tensor, atype: torch.Tensor, - box: Optional[torch.Tensor] = None, - fparam: Optional[torch.Tensor] = None, - aparam: Optional[torch.Tensor] = None, + box: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, do_atomic_virial: bool = False, ) -> dict[str, torch.Tensor]: return self.model.forward( @@ -143,7 +142,7 @@ def get_model_def_script(self) -> str: return self.model.get_model_def_script() @torch.jit.export - def get_min_nbor_dist(self) -> Optional[float]: + def get_min_nbor_dist(self) -> float | None: """Get the minimum neighbor distance.""" return self.model.get_min_nbor_dist() @@ -176,9 +175,9 @@ def get_nsel(self) -> int: def update_sel( cls, train_data: DeepmdDataSystem, - type_map: Optional[list[str]], + type_map: list[str] | None, local_jdata: dict, - ) -> tuple[dict, Optional[float]]: + ) -> tuple[dict, float | None]: """Update the selection and perform neighbor statistics. Parameters diff --git a/deepmd/pt/model/model/make_hessian_model.py b/deepmd/pt/model/model/make_hessian_model.py index b84e63ebd7..1b1bc3feba 100644 --- a/deepmd/pt/model/model/make_hessian_model.py +++ b/deepmd/pt/model/model/make_hessian_model.py @@ -3,8 +3,6 @@ import math from typing import ( Any, - Optional, - Union, ) import torch @@ -49,7 +47,7 @@ def __init__( def requires_hessian( self, - keys: Union[str, list[str]], + keys: str | list[str], ) -> None: """Set which output variable(s) requires hessian.""" if isinstance(keys, str): @@ -66,9 +64,9 @@ def forward_common( self, coord: torch.Tensor, atype: torch.Tensor, - box: Optional[torch.Tensor] = None, - fparam: Optional[torch.Tensor] = None, - aparam: Optional[torch.Tensor] = None, + box: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, do_atomic_virial: bool = False, ) -> dict[str, torch.Tensor]: """Return model prediction. @@ -121,9 +119,9 @@ def _cal_hessian_all( self, coord: torch.Tensor, atype: torch.Tensor, - box: Optional[torch.Tensor] = None, - fparam: Optional[torch.Tensor] = None, - aparam: Optional[torch.Tensor] = None, + box: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: nf, nloc = atype.shape coord = coord.view([nf, (nloc * 3)]) @@ -166,9 +164,9 @@ def _cal_hessian_one_component( ci: int, coord: torch.Tensor, atype: torch.Tensor, - box: Optional[torch.Tensor] = None, - fparam: Optional[torch.Tensor] = None, - aparam: Optional[torch.Tensor] = None, + box: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, ) -> torch.Tensor: # coord, # (nloc x 3) # atype, # nloc @@ -189,9 +187,9 @@ def __init__( obj: CM, ci: int, atype: torch.Tensor, - box: Optional[torch.Tensor], - fparam: Optional[torch.Tensor], - aparam: Optional[torch.Tensor], + box: torch.Tensor | None, + fparam: torch.Tensor | None, + aparam: torch.Tensor | None, ) -> None: self.atype, self.box, self.fparam, self.aparam = atype, box, fparam, aparam self.ci = ci diff --git a/deepmd/pt/model/model/make_model.py b/deepmd/pt/model/model/make_model.py index 3d97c0d4db..c958a62bf6 100644 --- a/deepmd/pt/model/model/make_model.py +++ b/deepmd/pt/model/model/make_model.py @@ -1,8 +1,9 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +from collections.abc import ( + Callable, +) from typing import ( Any, - Callable, - Optional, ) import torch @@ -69,7 +70,7 @@ def __init__( self, *args: Any, # underscore to prevent conflict with normal inputs - atomic_model_: Optional[T_AtomicModel] = None, + atomic_model_: T_AtomicModel | None = None, **kwargs: Any, ) -> None: super().__init__(*args, **kwargs) @@ -133,9 +134,9 @@ def forward_common( self, coord: torch.Tensor, atype: torch.Tensor, - box: Optional[torch.Tensor] = None, - fparam: Optional[torch.Tensor] = None, - aparam: Optional[torch.Tensor] = None, + box: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, do_atomic_virial: bool = False, ) -> dict[str, torch.Tensor]: """Return model prediction. @@ -240,11 +241,11 @@ def forward_common_lower( extended_coord: torch.Tensor, extended_atype: torch.Tensor, nlist: torch.Tensor, - mapping: Optional[torch.Tensor] = None, - fparam: Optional[torch.Tensor] = None, - aparam: Optional[torch.Tensor] = None, + mapping: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, do_atomic_virial: bool = False, - comm_dict: Optional[dict[str, torch.Tensor]] = None, + comm_dict: dict[str, torch.Tensor] | None = None, extra_nlist_sort: bool = False, ) -> dict[str, torch.Tensor]: """Return model prediction. Lower interface that takes @@ -311,14 +312,14 @@ def forward_common_lower( def input_type_cast( self, coord: torch.Tensor, - box: Optional[torch.Tensor] = None, - fparam: Optional[torch.Tensor] = None, - aparam: Optional[torch.Tensor] = None, + box: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, ) -> tuple[ torch.Tensor, - Optional[torch.Tensor], - Optional[torch.Tensor], - Optional[torch.Tensor], + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, str, ]: """Cast the input data to global float type.""" @@ -333,7 +334,7 @@ def input_type_cast( # " does not match" # f" that of the coordinate {input_prec}" # ) - _lst: list[Optional[torch.Tensor]] = [ + _lst: list[torch.Tensor | None] = [ vv.to(coord.dtype) if vv is not None else None for vv in [box, fparam, aparam] ] @@ -483,7 +484,7 @@ def _format_nlist( def do_grad_r( self, - var_name: Optional[str] = None, + var_name: str | None = None, ) -> bool: """Tell if the output variable `var_name` is r_differentiable. if var_name is None, returns if any of the variable is r_differentiable. @@ -492,7 +493,7 @@ def do_grad_r( def do_grad_c( self, - var_name: Optional[str] = None, + var_name: str | None = None, ) -> bool: """Tell if the output variable `var_name` is c_differentiable. if var_name is None, returns if any of the variable is c_differentiable. @@ -500,7 +501,7 @@ def do_grad_c( return self.atomic_model.do_grad_c(var_name) def change_type_map( - self, type_map: list[str], model_with_new_type_stat: Optional[Any] = None + self, type_map: list[str], model_with_new_type_stat: Any | None = None ) -> None: """Change the type related params to new ones, according to `type_map` and the original one in the model. If there are new types in `type_map`, statistics will be updated accordingly to `model_with_new_type_stat` for these new types. @@ -532,7 +533,7 @@ def has_default_fparam(self) -> bool: """Check if the model has default frame parameters.""" return self.atomic_model.has_default_fparam() - def get_default_fparam(self) -> Optional[torch.Tensor]: + def get_default_fparam(self) -> torch.Tensor | None: return self.atomic_model.get_default_fparam() @torch.jit.export @@ -585,7 +586,7 @@ def atomic_output_def(self) -> FittingOutputDef: def compute_or_load_stat( self, sampled_func: Callable[[], Any], - stat_file_path: Optional[DPPath] = None, + stat_file_path: DPPath | None = None, ) -> None: """Compute or load the statistics.""" return self.atomic_model.compute_or_load_stat(sampled_func, stat_file_path) @@ -619,9 +620,9 @@ def forward( self, coord: torch.Tensor, atype: torch.Tensor, - box: Optional[torch.Tensor] = None, - fparam: Optional[torch.Tensor] = None, - aparam: Optional[torch.Tensor] = None, + box: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, do_atomic_virial: bool = False, ) -> dict[str, torch.Tensor]: # directly call the forward_common method when no specific transform rule diff --git a/deepmd/pt/model/model/model.py b/deepmd/pt/model/model/model.py index e3cf7bde17..d9169fdc54 100644 --- a/deepmd/pt/model/model/model.py +++ b/deepmd/pt/model/model/model.py @@ -2,7 +2,6 @@ from typing import ( Any, NoReturn, - Optional, ) import torch @@ -30,7 +29,7 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: def compute_or_load_stat( self, sampled_func: Any, - stat_file_path: Optional[DPPath] = None, + stat_file_path: DPPath | None = None, ) -> NoReturn: """ Compute or load the statistics parameters of the model, @@ -65,7 +64,7 @@ def get_model_def_script(self) -> str: return self.model_def_script @torch.jit.export - def get_min_nbor_dist(self) -> Optional[float]: + def get_min_nbor_dist(self) -> float | None: """Get the minimum distance between two atoms.""" if self.min_nbor_dist.item() == -1.0: return None diff --git a/deepmd/pt/model/model/polar_model.py b/deepmd/pt/model/model/polar_model.py index 18eac5d24c..7210117823 100644 --- a/deepmd/pt/model/model/polar_model.py +++ b/deepmd/pt/model/model/polar_model.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( Any, - Optional, ) import torch @@ -52,9 +51,9 @@ def forward( self, coord: torch.Tensor, atype: torch.Tensor, - box: Optional[torch.Tensor] = None, - fparam: Optional[torch.Tensor] = None, - aparam: Optional[torch.Tensor] = None, + box: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, do_atomic_virial: bool = False, ) -> dict[str, torch.Tensor]: model_ret = self.forward_common( @@ -82,11 +81,11 @@ def forward_lower( extended_coord: torch.Tensor, extended_atype: torch.Tensor, nlist: torch.Tensor, - mapping: Optional[torch.Tensor] = None, - fparam: Optional[torch.Tensor] = None, - aparam: Optional[torch.Tensor] = None, + mapping: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, do_atomic_virial: bool = False, - comm_dict: Optional[dict[str, torch.Tensor]] = None, + comm_dict: dict[str, torch.Tensor] | None = None, ) -> dict[str, torch.Tensor]: model_ret = self.forward_common_lower( extended_coord, diff --git a/deepmd/pt/model/model/property_model.py b/deepmd/pt/model/model/property_model.py index 0931862ae8..c24ca7fa64 100644 --- a/deepmd/pt/model/model/property_model.py +++ b/deepmd/pt/model/model/property_model.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( Any, - Optional, ) import torch @@ -52,9 +51,9 @@ def forward( self, coord: torch.Tensor, atype: torch.Tensor, - box: Optional[torch.Tensor] = None, - fparam: Optional[torch.Tensor] = None, - aparam: Optional[torch.Tensor] = None, + box: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, do_atomic_virial: bool = False, ) -> dict[str, torch.Tensor]: model_ret = self.forward_common( @@ -93,11 +92,11 @@ def forward_lower( extended_coord: torch.Tensor, extended_atype: torch.Tensor, nlist: torch.Tensor, - mapping: Optional[torch.Tensor] = None, - fparam: Optional[torch.Tensor] = None, - aparam: Optional[torch.Tensor] = None, + mapping: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, do_atomic_virial: bool = False, - comm_dict: Optional[dict[str, torch.Tensor]] = None, + comm_dict: dict[str, torch.Tensor] | None = None, ) -> dict[str, torch.Tensor]: model_ret = self.forward_common_lower( extended_coord, diff --git a/deepmd/pt/model/model/spin_model.py b/deepmd/pt/model/model/spin_model.py index bd7158fb8f..0eb7014416 100644 --- a/deepmd/pt/model/model/spin_model.py +++ b/deepmd/pt/model/model/spin_model.py @@ -1,12 +1,13 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import functools +from collections.abc import ( + Callable, +) from copy import ( deepcopy, ) from typing import ( Any, - Callable, - Optional, ) import torch @@ -70,8 +71,8 @@ def process_spin_input_lower( extended_atype: torch.Tensor, extended_spin: torch.Tensor, nlist: torch.Tensor, - mapping: Optional[torch.Tensor] = None, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: + mapping: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor | None]: """ Add `extended_spin` into `extended_coord` to generate virtual atoms, and extend `nlist` and `mapping`. Note that the final `extended_coord_updated` with shape [nframes, nall + nall, 3] has the following order: @@ -293,7 +294,7 @@ def get_model_def_script(self) -> str: return self.backbone_model.get_model_def_script() @torch.jit.export - def get_min_nbor_dist(self) -> Optional[float]: + def get_min_nbor_dist(self) -> float | None: """Get the minimum neighbor distance.""" return self.backbone_model.get_min_nbor_dist() @@ -352,7 +353,7 @@ def __getattr__(self, name: str) -> Any: def compute_or_load_stat( self, sampled_func: Callable[[], list[dict[str, Any]]], - stat_file_path: Optional[DPPath] = None, + stat_file_path: DPPath | None = None, ) -> None: """ Compute or load the statistics parameters of the model, @@ -400,9 +401,9 @@ def forward_common( coord: torch.Tensor, atype: torch.Tensor, spin: torch.Tensor, - box: Optional[torch.Tensor] = None, - fparam: Optional[torch.Tensor] = None, - aparam: Optional[torch.Tensor] = None, + box: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, do_atomic_virial: bool = False, ) -> dict[str, torch.Tensor]: nframes, nloc = atype.shape @@ -449,11 +450,11 @@ def forward_common_lower( extended_atype: torch.Tensor, extended_spin: torch.Tensor, nlist: torch.Tensor, - mapping: Optional[torch.Tensor] = None, - fparam: Optional[torch.Tensor] = None, - aparam: Optional[torch.Tensor] = None, + mapping: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, do_atomic_virial: bool = False, - comm_dict: Optional[dict[str, torch.Tensor]] = None, + comm_dict: dict[str, torch.Tensor] | None = None, extra_nlist_sort: bool = False, ) -> dict[str, torch.Tensor]: nframes, nloc = nlist.shape[:2] @@ -556,9 +557,9 @@ def forward( coord: torch.Tensor, atype: torch.Tensor, spin: torch.Tensor, - box: Optional[torch.Tensor] = None, - fparam: Optional[torch.Tensor] = None, - aparam: Optional[torch.Tensor] = None, + box: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, do_atomic_virial: bool = False, ) -> dict[str, torch.Tensor]: model_ret = self.forward_common( @@ -587,11 +588,11 @@ def forward_lower( extended_atype: torch.Tensor, extended_spin: torch.Tensor, nlist: torch.Tensor, - mapping: Optional[torch.Tensor] = None, - fparam: Optional[torch.Tensor] = None, - aparam: Optional[torch.Tensor] = None, + mapping: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, do_atomic_virial: bool = False, - comm_dict: Optional[dict[str, torch.Tensor]] = None, + comm_dict: dict[str, torch.Tensor] | None = None, ) -> dict[str, torch.Tensor]: model_ret = self.forward_common_lower( extended_coord, diff --git a/deepmd/pt/model/model/transform_output.py b/deepmd/pt/model/model/transform_output.py index cd88e4cb40..839a028fcf 100644 --- a/deepmd/pt/model/model/transform_output.py +++ b/deepmd/pt/model/model/transform_output.py @@ -1,7 +1,4 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( - Optional, -) import torch @@ -29,7 +26,7 @@ def atomic_virial_corr( ce = coord * atom_energy sumce0, sumce1, sumce2 = torch.split(torch.sum(ce, dim=1), [1, 1, 1], dim=-1) faked_grad = torch.ones_like(sumce0) - lst = torch.jit.annotate(list[Optional[torch.Tensor]], [faked_grad]) + lst = torch.jit.annotate(list[torch.Tensor | None], [faked_grad]) extended_virial_corr0 = torch.autograd.grad( [sumce0], [extended_coord], @@ -72,9 +69,9 @@ def task_deriv_one( do_virial: bool = True, do_atomic_virial: bool = False, create_graph: bool = True, -) -> tuple[torch.Tensor, Optional[torch.Tensor]]: +) -> tuple[torch.Tensor, torch.Tensor | None]: faked_grad = torch.ones_like(energy) - lst = torch.jit.annotate(list[Optional[torch.Tensor]], [faked_grad]) + lst = torch.jit.annotate(list[torch.Tensor | None], [faked_grad]) extended_force = torch.autograd.grad( [energy], [extended_coord], @@ -116,7 +113,7 @@ def take_deriv( do_virial: bool = False, do_atomic_virial: bool = False, create_graph: bool = True, -) -> tuple[torch.Tensor, Optional[torch.Tensor]]: +) -> tuple[torch.Tensor, torch.Tensor | None]: size = 1 for ii in vdef.shape: size *= ii @@ -158,7 +155,7 @@ def fit_output_to_model_output( coord_ext: torch.Tensor, do_atomic_virial: bool = False, create_graph: bool = True, - mask: Optional[torch.Tensor] = None, + mask: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: """Transform the output of the fitting network to the model output. diff --git a/deepmd/pt/model/network/init.py b/deepmd/pt/model/network/init.py index 6bdff61eea..136690b749 100644 --- a/deepmd/pt/model/network/init.py +++ b/deepmd/pt/model/network/init.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import math import warnings -from typing import Optional as _Optional import torch from torch import ( @@ -22,7 +21,7 @@ def _no_grad_uniform_( tensor: torch.Tensor, a: float, b: float, - generator: _Optional[torch.Generator] = None, + generator: torch.Generator | None = None, ) -> torch.Tensor: with torch.no_grad(): return tensor.uniform_(a, b, generator=generator) @@ -32,7 +31,7 @@ def _no_grad_normal_( tensor: torch.Tensor, mean: float, std: float, - generator: _Optional[torch.Generator] = None, + generator: torch.Generator | None = None, ) -> torch.Tensor: with torch.no_grad(): return tensor.normal_(mean, std, generator=generator) @@ -44,7 +43,7 @@ def _no_grad_trunc_normal_( std: float, a: float, b: float, - generator: _Optional[torch.Generator] = None, + generator: torch.Generator | None = None, ) -> torch.Tensor: # Method based on https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdf def norm_cdf(x: float) -> float: @@ -92,7 +91,7 @@ def _no_grad_fill_(tensor: torch.Tensor, val: float) -> torch.Tensor: return tensor.fill_(val) -def calculate_gain(nonlinearity: str, param: _Optional[float] = None) -> float: +def calculate_gain(nonlinearity: str, param: float | None = None) -> float: r"""Return the recommended gain value for the given nonlinearity function. The values are as follows: @@ -245,7 +244,7 @@ def normal_( tensor: Tensor, mean: float = 0.0, std: float = 1.0, - generator: _Optional[torch.Generator] = None, + generator: torch.Generator | None = None, ) -> Tensor: r"""Fill the input Tensor with values drawn from the normal distribution. @@ -275,7 +274,7 @@ def trunc_normal_( std: float = 1.0, a: float = -2.0, b: float = 2.0, - generator: _Optional[torch.Generator] = None, + generator: torch.Generator | None = None, ) -> Tensor: r"""Fill the input Tensor with values drawn from a truncated normal distribution. @@ -306,7 +305,7 @@ def kaiming_uniform_( a: float = 0, mode: str = "fan_in", nonlinearity: str = "leaky_relu", - generator: _Optional[torch.Generator] = None, + generator: torch.Generator | None = None, ) -> Tensor: r"""Fill the input `Tensor` with values using a Kaiming uniform distribution. @@ -364,7 +363,7 @@ def kaiming_normal_( a: float = 0, mode: str = "fan_in", nonlinearity: str = "leaky_relu", - generator: _Optional[torch.Generator] = None, + generator: torch.Generator | None = None, ) -> Tensor: r"""Fill the input `Tensor` with values using a Kaiming normal distribution. @@ -406,7 +405,7 @@ def kaiming_normal_( def xavier_uniform_( - tensor: Tensor, gain: float = 1.0, generator: _Optional[torch.Generator] = None + tensor: Tensor, gain: float = 1.0, generator: torch.Generator | None = None ) -> Tensor: r"""Fill the input `Tensor` with values using a Xavier uniform distribution. @@ -440,7 +439,7 @@ def xavier_uniform_( def xavier_normal_( tensor: Tensor, gain: float = 1.0, - generator: _Optional[torch.Generator] = None, + generator: torch.Generator | None = None, ) -> Tensor: r"""Fill the input `Tensor` with values using a Xavier normal distribution. diff --git a/deepmd/pt/model/network/layernorm.py b/deepmd/pt/model/network/layernorm.py index fdf31d0ffd..c7ffa3e205 100644 --- a/deepmd/pt/model/network/layernorm.py +++ b/deepmd/pt/model/network/layernorm.py @@ -1,8 +1,4 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( - Optional, - Union, -) import numpy as np import torch @@ -44,7 +40,7 @@ def __init__( stddev: float = 1.0, precision: str = DEFAULT_PRECISION, trainable: bool = True, - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, ) -> None: super().__init__() self.eps = eps @@ -141,7 +137,7 @@ def deserialize(cls, data: dict) -> "LayerNorm": ) prec = PRECISION_DICT[obj.precision] - def check_load_param(ss: str) -> Optional[nn.Parameter]: + def check_load_param(ss: str) -> nn.Parameter | None: return ( nn.Parameter(data=to_torch_tensor(nl[ss])) if nl[ss] is not None diff --git a/deepmd/pt/model/network/mlp.py b/deepmd/pt/model/network/mlp.py index a850c85a9b..02f7611429 100644 --- a/deepmd/pt/model/network/mlp.py +++ b/deepmd/pt/model/network/mlp.py @@ -2,8 +2,6 @@ from typing import ( Any, ClassVar, - Optional, - Union, ) import numpy as np @@ -77,13 +75,13 @@ def __init__( num_out: int, bias: bool = True, use_timestep: bool = False, - activation_function: Optional[str] = None, + activation_function: str | None = None, resnet: bool = False, bavg: float = 0.0, stddev: float = 1.0, precision: str = DEFAULT_PRECISION, init: str = "default", - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, trainable: bool = True, ) -> None: super().__init__() @@ -133,7 +131,7 @@ def __init__( def check_type_consistency(self) -> None: precision = self.precision - def check_var(var: Optional[torch.Tensor]) -> None: + def check_var(var: torch.Tensor | None) -> None: if var is not None: # assertion "float64" == "double" would fail assert PRECISION_DICT[var.dtype.name] is PRECISION_DICT[precision] @@ -152,7 +150,7 @@ def _default_normal_init( self, bavg: float = 0.0, stddev: float = 1.0, - generator: Optional[torch.Generator] = None, + generator: torch.Generator | None = None, ) -> None: normal_( self.matrix.data, @@ -165,7 +163,7 @@ def _default_normal_init( normal_(self.idt.data, mean=0.1, std=0.001, generator=generator) def _trunc_normal_init( - self, scale: float = 1.0, generator: Optional[torch.Generator] = None + self, scale: float = 1.0, generator: torch.Generator | None = None ) -> None: # Constant from scipy.stats.truncnorm.std(a=-2, b=2, loc=0., scale=1.) TRUNCATED_NORMAL_STDDEV_FACTOR = 0.87962566103423978 @@ -174,7 +172,7 @@ def _trunc_normal_init( std = (scale**0.5) / TRUNCATED_NORMAL_STDDEV_FACTOR trunc_normal_(self.matrix, mean=0.0, std=std, generator=generator) - def _glorot_uniform_init(self, generator: Optional[torch.Generator] = None) -> None: + def _glorot_uniform_init(self, generator: torch.Generator | None = None) -> None: xavier_uniform_(self.matrix, gain=1, generator=generator) def _zero_init(self, use_bias: bool = True) -> None: @@ -184,7 +182,7 @@ def _zero_init(self, use_bias: bool = True) -> None: with torch.no_grad(): self.bias.fill_(1.0) - def _normal_init(self, generator: Optional[torch.Generator] = None) -> None: + def _normal_init(self, generator: torch.Generator | None = None) -> None: kaiming_normal_(self.matrix, nonlinearity="linear", generator=generator) def forward( @@ -267,7 +265,7 @@ def deserialize(cls, data: dict) -> "MLPLayer": ) prec = PRECISION_DICT[obj.precision] - def check_load_param(ss: str) -> Optional[nn.Parameter]: + def check_load_param(ss: str) -> nn.Parameter | None: return ( nn.Parameter(data=to_torch_tensor(nl[ss])) if nl[ss] is not None diff --git a/deepmd/pt/model/network/network.py b/deepmd/pt/model/network/network.py index d95741b05c..a9662acfec 100644 --- a/deepmd/pt/model/network/network.py +++ b/deepmd/pt/model/network/network.py @@ -2,8 +2,6 @@ from typing import ( Any, Final, - Optional, - Union, ) import numpy as np @@ -47,7 +45,7 @@ def __init__( bavg: float = 0.0, stddev: float = 1.0, use_timestep: bool = False, - activate: Optional[str] = None, + activate: str | None = None, bias: bool = True, ) -> None: """Construct a linear layer. @@ -150,7 +148,7 @@ def __init__( input_dim: int, out_dim: int, activation_fn: str, - hidden: Optional[int] = None, + hidden: int | None = None, ) -> None: super().__init__() hidden = input_dim if not hidden else hidden @@ -171,7 +169,7 @@ def __init__( embed_dim: int, output_dim: int, activation_fn: str, - weight: Optional[torch.Tensor] = None, + weight: torch.Tensor | None = None, ) -> None: super().__init__() self.dense = SimpleLinear(embed_dim, embed_dim) @@ -190,7 +188,7 @@ def __init__( def forward( self, features: torch.Tensor, - masked_tokens: Optional[torch.Tensor] = None, + masked_tokens: torch.Tensor | None = None, **kwargs: Any, ) -> torch.Tensor: # Only project the masked tokens while training, @@ -273,10 +271,10 @@ def __init__( bavg: float = 0.0, stddev: float = 1.0, precision: str = "default", - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, use_econf_tebd: bool = False, use_tebd_bias: bool = False, - type_map: Optional[list[str]] = None, + type_map: list[str] | None = None, trainable: bool = True, ) -> None: """Construct a type embedding net.""" @@ -350,7 +348,7 @@ def share_params( raise NotImplementedError def change_type_map( - self, type_map: list[str], model_with_new_type_stat: Optional[Any] = None + self, type_map: list[str], model_with_new_type_stat: Any | None = None ) -> None: """Change the type related params to new ones, according to `type_map` and the original one in the model. If there are new types in `type_map`, statistics will be updated accordingly to `model_with_new_type_stat` for these new types. @@ -396,11 +394,11 @@ def __init__( activation_function: str = "tanh", precision: str = "default", trainable: bool = True, - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, padding: bool = False, use_econf_tebd: bool = False, use_tebd_bias: bool = False, - type_map: Optional[list[str]] = None, + type_map: list[str] | None = None, ) -> None: """Construct a type embedding net.""" super().__init__() @@ -457,7 +455,7 @@ def forward(self, device: torch.device) -> torch.Tensor: return embed def change_type_map( - self, type_map: list[str], model_with_new_type_stat: Optional[Any] = None + self, type_map: list[str], model_with_new_type_stat: Any | None = None ) -> None: """Change the type related params to new ones, according to `type_map` and the original one in the model. If there are new types in `type_map`, statistics will be updated accordingly to `model_with_new_type_stat` for these new types. diff --git a/deepmd/pt/model/network/utils.py b/deepmd/pt/model/network/utils.py index 7af8b7c032..0613ca3e4d 100644 --- a/deepmd/pt/model/network/utils.py +++ b/deepmd/pt/model/network/utils.py @@ -1,7 +1,4 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( - Optional, -) import torch @@ -11,7 +8,7 @@ def aggregate( data: torch.Tensor, owners: torch.Tensor, average: bool = True, - num_owner: Optional[int] = None, + num_owner: int | None = None, ) -> torch.Tensor: """ Aggregate rows in data by specifying the owners. diff --git a/deepmd/pt/model/task/denoise.py b/deepmd/pt/model/task/denoise.py index 50cae4fb12..bb155322e9 100644 --- a/deepmd/pt/model/task/denoise.py +++ b/deepmd/pt/model/task/denoise.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( Any, - Optional, ) import torch @@ -99,7 +98,7 @@ def forward( nlist_mask: torch.Tensor, features: torch.Tensor, sw: torch.Tensor, - masked_tokens: Optional[torch.Tensor] = None, + masked_tokens: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: """Calculate the updated coord. Args: diff --git a/deepmd/pt/model/task/dipole.py b/deepmd/pt/model/task/dipole.py index b6a1477f7a..6e9b0340ff 100644 --- a/deepmd/pt/model/task/dipole.py +++ b/deepmd/pt/model/task/dipole.py @@ -1,10 +1,10 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import logging +from collections.abc import ( + Callable, +) from typing import ( Any, - Callable, - Optional, - Union, ) import torch @@ -91,13 +91,13 @@ def __init__( activation_function: str = "tanh", precision: str = DEFAULT_PRECISION, mixed_types: bool = True, - rcond: Optional[float] = None, - seed: Optional[Union[int, list[int]]] = None, + rcond: float | None = None, + seed: int | list[int] | None = None, exclude_types: list[int] = [], r_differentiable: bool = True, c_differentiable: bool = True, - type_map: Optional[list[str]] = None, - default_fparam: Optional[list] = None, + type_map: list[str] | None = None, + default_fparam: list | None = None, **kwargs: Any, ) -> None: self.embedding_width = embedding_width @@ -157,8 +157,8 @@ def output_def(self) -> FittingOutputDef: def compute_output_stats( self, - merged: Union[Callable[[], list[dict]], list[dict]], - stat_file_path: Optional[DPPath] = None, + merged: Callable[[], list[dict]] | list[dict], + stat_file_path: DPPath | None = None, ) -> None: """ Compute the output statistics (e.g. energy bias) for the fitting net from packed data. @@ -182,11 +182,11 @@ def forward( self, descriptor: torch.Tensor, atype: torch.Tensor, - gr: Optional[torch.Tensor] = None, - g2: Optional[torch.Tensor] = None, - h2: Optional[torch.Tensor] = None, - fparam: Optional[torch.Tensor] = None, - aparam: Optional[torch.Tensor] = None, + gr: torch.Tensor | None = None, + g2: torch.Tensor | None = None, + h2: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: nframes, nloc, _ = descriptor.shape assert gr is not None, "Must provide the rotation matrix for dipole fitting." diff --git a/deepmd/pt/model/task/dos.py b/deepmd/pt/model/task/dos.py index afbed5f748..b11bfd7c7f 100644 --- a/deepmd/pt/model/task/dos.py +++ b/deepmd/pt/model/task/dos.py @@ -1,9 +1,5 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import logging -from typing import ( - Optional, - Union, -) import torch @@ -48,16 +44,16 @@ def __init__( numb_fparam: int = 0, numb_aparam: int = 0, dim_case_embd: int = 0, - rcond: Optional[float] = None, - bias_dos: Optional[torch.Tensor] = None, - trainable: Union[bool, list[bool]] = True, - seed: Optional[Union[int, list[int]]] = None, + rcond: float | None = None, + bias_dos: torch.Tensor | None = None, + trainable: bool | list[bool] = True, + seed: int | list[int] | None = None, activation_function: str = "tanh", precision: str = DEFAULT_PRECISION, exclude_types: list[int] = [], mixed_types: bool = True, - type_map: Optional[list[str]] = None, - default_fparam: Optional[list] = None, + type_map: list[str] | None = None, + default_fparam: list | None = None, ) -> None: if bias_dos is not None: self.bias_dos = bias_dos diff --git a/deepmd/pt/model/task/ener.py b/deepmd/pt/model/task/ener.py index af288bec10..39dabf99f8 100644 --- a/deepmd/pt/model/task/ener.py +++ b/deepmd/pt/model/task/ener.py @@ -2,8 +2,6 @@ import logging from typing import ( Any, - Optional, - Union, ) import numpy as np @@ -47,7 +45,7 @@ def __init__( ntypes: int, dim_descrpt: int, neuron: list[int] = [128, 128, 128], - bias_atom_e: Optional[torch.Tensor] = None, + bias_atom_e: torch.Tensor | None = None, resnet_dt: bool = True, numb_fparam: int = 0, numb_aparam: int = 0, @@ -55,9 +53,9 @@ def __init__( activation_function: str = "tanh", precision: str = DEFAULT_PRECISION, mixed_types: bool = True, - seed: Optional[Union[int, list[int]]] = None, - type_map: Optional[list[str]] = None, - default_fparam: Optional[list] = None, + seed: int | list[int] | None = None, + type_map: list[str] | None = None, + default_fparam: list | None = None, **kwargs: Any, ) -> None: super().__init__( @@ -108,7 +106,7 @@ def __init__( ntypes: int, dim_descrpt: int, neuron: list[int], - bias_atom_e: Optional[torch.Tensor] = None, + bias_atom_e: torch.Tensor | None = None, out_dim: int = 1, resnet_dt: bool = True, use_tebd: bool = True, @@ -190,7 +188,7 @@ def deserialize(self) -> "EnergyFittingNetDirect": raise NotImplementedError def change_type_map( - self, type_map: list[str], model_with_new_type_stat: Optional[Any] = None + self, type_map: list[str], model_with_new_type_stat: Any | None = None ) -> None: raise NotImplementedError @@ -201,11 +199,11 @@ def forward( self, inputs: torch.Tensor, atype: torch.Tensor, - gr: Optional[torch.Tensor] = None, - g2: Optional[torch.Tensor] = None, - h2: Optional[torch.Tensor] = None, - fparam: Optional[torch.Tensor] = None, - aparam: Optional[torch.Tensor] = None, + gr: torch.Tensor | None = None, + g2: torch.Tensor | None = None, + h2: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, ) -> tuple[torch.Tensor, None]: """Based on embedding net output, alculate total energy. diff --git a/deepmd/pt/model/task/fitting.py b/deepmd/pt/model/task/fitting.py index 578e1683ba..f46bc0fb7d 100644 --- a/deepmd/pt/model/task/fitting.py +++ b/deepmd/pt/model/task/fitting.py @@ -3,11 +3,12 @@ from abc import ( abstractmethod, ) +from collections.abc import ( + Callable, +) from typing import ( Any, - Callable, Optional, - Union, ) import numpy as np @@ -235,9 +236,9 @@ def restore_aparam_from_file(self, stat_file_path: DPPath) -> None: def compute_input_stats( self, - merged: Union[Callable[[], list[dict]], list[dict]], + merged: Callable[[], list[dict]] | list[dict], protection: float = 1e-2, - stat_file_path: Optional[DPPath] = None, + stat_file_path: DPPath | None = None, ) -> None: """ Compute the input statistics (e.g. mean and stddev) for the fittings from packed data. @@ -411,7 +412,7 @@ def __init__( ntypes: int, dim_descrpt: int, neuron: list[int] = [128, 128, 128], - bias_atom_e: Optional[torch.Tensor] = None, + bias_atom_e: torch.Tensor | None = None, resnet_dt: bool = True, numb_fparam: int = 0, numb_aparam: int = 0, @@ -419,14 +420,14 @@ def __init__( activation_function: str = "tanh", precision: str = DEFAULT_PRECISION, mixed_types: bool = True, - rcond: Optional[float] = None, - seed: Optional[Union[int, list[int]]] = None, + rcond: float | None = None, + seed: int | list[int] | None = None, exclude_types: list[int] = [], - trainable: Union[bool, list[bool]] = True, - remove_vaccum_contribution: Optional[list[bool]] = None, - type_map: Optional[list[str]] = None, + trainable: bool | list[bool] = True, + remove_vaccum_contribution: list[bool] | None = None, + type_map: list[str] | None = None, use_aparam_as_mask: bool = False, - default_fparam: Optional[list[float]] = None, + default_fparam: list[float] | None = None, **kwargs: Any, ) -> None: super().__init__() @@ -640,7 +641,7 @@ def has_default_fparam(self) -> bool: """Check if the fitting has default frame parameters.""" return self.default_fparam is not None - def get_default_fparam(self) -> Optional[torch.Tensor]: + def get_default_fparam(self) -> torch.Tensor | None: return self.default_fparam_tensor def get_dim_aparam(self) -> int: @@ -736,11 +737,11 @@ def _forward_common( self, descriptor: torch.Tensor, atype: torch.Tensor, - gr: Optional[torch.Tensor] = None, - g2: Optional[torch.Tensor] = None, - h2: Optional[torch.Tensor] = None, - fparam: Optional[torch.Tensor] = None, - aparam: Optional[torch.Tensor] = None, + gr: torch.Tensor | None = None, + g2: torch.Tensor | None = None, + h2: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: # cast the input to internal precsion xx = descriptor.to(self.prec) diff --git a/deepmd/pt/model/task/invar_fitting.py b/deepmd/pt/model/task/invar_fitting.py index 4ec3407901..a8953fcd2b 100644 --- a/deepmd/pt/model/task/invar_fitting.py +++ b/deepmd/pt/model/task/invar_fitting.py @@ -2,8 +2,6 @@ import logging from typing import ( Any, - Optional, - Union, ) import torch @@ -93,7 +91,7 @@ def __init__( dim_descrpt: int, dim_out: int, neuron: list[int] = [128, 128, 128], - bias_atom_e: Optional[torch.Tensor] = None, + bias_atom_e: torch.Tensor | None = None, resnet_dt: bool = True, numb_fparam: int = 0, numb_aparam: int = 0, @@ -101,13 +99,13 @@ def __init__( activation_function: str = "tanh", precision: str = DEFAULT_PRECISION, mixed_types: bool = True, - rcond: Optional[float] = None, - seed: Optional[Union[int, list[int]]] = None, + rcond: float | None = None, + seed: int | list[int] | None = None, exclude_types: list[int] = [], - atom_ener: Optional[list[Optional[torch.Tensor]]] = None, - type_map: Optional[list[str]] = None, + atom_ener: list[torch.Tensor | None] | None = None, + type_map: list[str] | None = None, use_aparam_as_mask: bool = False, - default_fparam: Optional[list[float]] = None, + default_fparam: list[float] | None = None, **kwargs: Any, ) -> None: self.dim_out = dim_out @@ -171,11 +169,11 @@ def forward( self, descriptor: torch.Tensor, atype: torch.Tensor, - gr: Optional[torch.Tensor] = None, - g2: Optional[torch.Tensor] = None, - h2: Optional[torch.Tensor] = None, - fparam: Optional[torch.Tensor] = None, - aparam: Optional[torch.Tensor] = None, + gr: torch.Tensor | None = None, + g2: torch.Tensor | None = None, + h2: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: """Based on embedding net output, alculate total energy. diff --git a/deepmd/pt/model/task/polarizability.py b/deepmd/pt/model/task/polarizability.py index bf63d9db4b..c3a7ed52a1 100644 --- a/deepmd/pt/model/task/polarizability.py +++ b/deepmd/pt/model/task/polarizability.py @@ -2,8 +2,6 @@ import logging from typing import ( Any, - Optional, - Union, ) import torch @@ -94,14 +92,14 @@ def __init__( activation_function: str = "tanh", precision: str = DEFAULT_PRECISION, mixed_types: bool = True, - rcond: Optional[float] = None, - seed: Optional[Union[int, list[int]]] = None, + rcond: float | None = None, + seed: int | list[int] | None = None, exclude_types: list[int] = [], fit_diag: bool = True, - scale: Optional[Union[list[float], float]] = None, + scale: list[float] | float | None = None, shift_diag: bool = True, - type_map: Optional[list[str]] = None, - default_fparam: Optional[list] = None, + type_map: list[str] | None = None, + default_fparam: list | None = None, **kwargs: Any, ) -> None: self.embedding_width = embedding_width @@ -168,7 +166,7 @@ def __getitem__(self, key: str) -> Any: return super().__getitem__(key) def change_type_map( - self, type_map: list[str], model_with_new_type_stat: Optional[Any] = None + self, type_map: list[str], model_with_new_type_stat: Any | None = None ) -> None: """Change the type related params to new ones, according to `type_map` and the original one in the model. If there are new types in `type_map`, statistics will be updated accordingly to `model_with_new_type_stat` for these new types. @@ -232,11 +230,11 @@ def forward( self, descriptor: torch.Tensor, atype: torch.Tensor, - gr: Optional[torch.Tensor] = None, - g2: Optional[torch.Tensor] = None, - h2: Optional[torch.Tensor] = None, - fparam: Optional[torch.Tensor] = None, - aparam: Optional[torch.Tensor] = None, + gr: torch.Tensor | None = None, + g2: torch.Tensor | None = None, + h2: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: nframes, nloc, _ = descriptor.shape assert gr is not None, ( diff --git a/deepmd/pt/model/task/property.py b/deepmd/pt/model/task/property.py index c2440b7de3..4d0b791cfa 100644 --- a/deepmd/pt/model/task/property.py +++ b/deepmd/pt/model/task/property.py @@ -2,8 +2,6 @@ import logging from typing import ( Any, - Optional, - Union, ) import torch @@ -81,7 +79,7 @@ def __init__( property_name: str, task_dim: int = 1, neuron: list[int] = [128, 128, 128], - bias_atom_p: Optional[torch.Tensor] = None, + bias_atom_p: torch.Tensor | None = None, intensive: bool = False, resnet_dt: bool = True, numb_fparam: int = 0, @@ -90,9 +88,9 @@ def __init__( activation_function: str = "tanh", precision: str = DEFAULT_PRECISION, mixed_types: bool = True, - trainable: Union[bool, list[bool]] = True, - seed: Optional[int] = None, - default_fparam: Optional[list] = None, + trainable: bool | list[bool] = True, + seed: int | None = None, + default_fparam: list | None = None, **kwargs: Any, ) -> None: self.task_dim = task_dim diff --git a/deepmd/pt/model/task/type_predict.py b/deepmd/pt/model/task/type_predict.py index 5c1b064d07..a5b2ff8884 100644 --- a/deepmd/pt/model/task/type_predict.py +++ b/deepmd/pt/model/task/type_predict.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( Any, - Optional, ) import torch @@ -40,7 +39,7 @@ def __init__( ) def forward( - self, features: torch.Tensor, masked_tokens: Optional[torch.Tensor] = None + self, features: torch.Tensor, masked_tokens: torch.Tensor | None = None ) -> torch.Tensor: """Calculate the predicted logits. Args: diff --git a/deepmd/pt/optimizer/LKF.py b/deepmd/pt/optimizer/LKF.py index aeb1120bff..c79e875f3e 100644 --- a/deepmd/pt/optimizer/LKF.py +++ b/deepmd/pt/optimizer/LKF.py @@ -3,7 +3,6 @@ import math from typing import ( Any, - Optional, ) import torch @@ -319,7 +318,7 @@ def step(self, error: torch.Tensor) -> None: self.__update(H, error, weights) - def get_device_id(self, index: int) -> Optional[int]: + def get_device_id(self, index: int) -> int | None: for i, (start, end) in enumerate(self.dindex): if start <= index < end: return i diff --git a/deepmd/pt/train/training.py b/deepmd/pt/train/training.py index d099b8b20b..46cbda1894 100644 --- a/deepmd/pt/train/training.py +++ b/deepmd/pt/train/training.py @@ -3,6 +3,7 @@ import logging import time from collections.abc import ( + Callable, Generator, Iterable, ) @@ -14,8 +15,6 @@ ) from typing import ( Any, - Callable, - Optional, ) import numpy as np @@ -97,15 +96,15 @@ def __init__( self, config: dict[str, Any], training_data: DpLoaderSet, - stat_file_path: Optional[str] = None, - validation_data: Optional[DpLoaderSet] = None, - init_model: Optional[str] = None, - restart_model: Optional[str] = None, - finetune_model: Optional[str] = None, + stat_file_path: str | None = None, + validation_data: DpLoaderSet | None = None, + init_model: str | None = None, + restart_model: str | None = None, + finetune_model: str | None = None, force_load: bool = False, - shared_links: Optional[dict[str, str]] = None, - finetune_links: Optional[dict[str, str]] = None, - init_frz_model: Optional[str] = None, + shared_links: dict[str, str] | None = None, + finetune_links: dict[str, str] | None = None, + init_frz_model: str | None = None, ) -> None: """Construct a DeePMD trainer. @@ -185,13 +184,13 @@ def cycle_iterator(iterable: Iterable) -> Generator[Any, None, None]: def get_data_loader( _training_data: DpLoaderSet, - _validation_data: Optional[DpLoaderSet], + _validation_data: DpLoaderSet | None, _training_params: dict[str, Any], ) -> tuple[ DataLoader, Generator[Any, None, None], - Optional[DataLoader], - Optional[Generator[Any, None, None]], + DataLoader | None, + Generator[Any, None, None] | None, int, ]: def get_dataloader_and_iter( @@ -246,8 +245,8 @@ def single_model_stat( _model: Any, _data_stat_nbatch: int, _training_data: DpLoaderSet, - _validation_data: Optional[DpLoaderSet], - _stat_file_path: Optional[str], + _validation_data: DpLoaderSet | None, + _stat_file_path: str | None, _data_requirement: list[DataRequirementItem], finetune_has_new_type: bool = False, ) -> Callable[[], Any]: @@ -1443,7 +1442,7 @@ def get_single_model( def get_model_for_wrapper( _model_params: dict[str, Any], resuming: bool = False, - _loss_params: Optional[dict[str, Any]] = None, + _loss_params: dict[str, Any] | None = None, ) -> Any: if "model_dict" not in _model_params: if _loss_params is not None and whether_hessian(_loss_params): diff --git a/deepmd/pt/train/wrapper.py b/deepmd/pt/train/wrapper.py index c65787958a..2669e3d832 100644 --- a/deepmd/pt/train/wrapper.py +++ b/deepmd/pt/train/wrapper.py @@ -2,8 +2,6 @@ import logging from typing import ( Any, - Optional, - Union, ) import torch @@ -18,10 +16,10 @@ class ModelWrapper(torch.nn.Module): def __init__( self, - model: Union[torch.nn.Module, dict], - loss: Union[torch.nn.Module, dict] = None, - model_params: Optional[dict[str, Any]] = None, - shared_links: Optional[dict[str, Any]] = None, + model: torch.nn.Module | dict, + loss: torch.nn.Module | dict = None, + model_params: dict[str, Any] | None = None, + shared_links: dict[str, Any] | None = None, ) -> None: """Construct a DeePMD model wrapper. @@ -155,15 +153,15 @@ def forward( self, coord: torch.Tensor, atype: torch.Tensor, - spin: Optional[torch.Tensor] = None, - box: Optional[torch.Tensor] = None, - cur_lr: Optional[torch.Tensor] = None, - label: Optional[torch.Tensor] = None, - task_key: Optional[torch.Tensor] = None, + spin: torch.Tensor | None = None, + box: torch.Tensor | None = None, + cur_lr: torch.Tensor | None = None, + label: torch.Tensor | None = None, + task_key: torch.Tensor | None = None, inference_only: bool = False, do_atomic_virial: bool = False, - fparam: Optional[torch.Tensor] = None, - aparam: Optional[torch.Tensor] = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, ) -> tuple[Any, Any, Any]: if not self.multi_task: task_key = "Default" diff --git a/deepmd/pt/utils/dataloader.py b/deepmd/pt/utils/dataloader.py index c434341ab9..c991e59daa 100644 --- a/deepmd/pt/utils/dataloader.py +++ b/deepmd/pt/utils/dataloader.py @@ -6,8 +6,6 @@ ) from typing import ( Any, - Optional, - Union, ) import h5py @@ -50,7 +48,7 @@ torch.multiprocessing.set_sharing_strategy("file_system") -def setup_seed(seed: Union[int, list[int], tuple[int, ...]]) -> None: +def setup_seed(seed: int | list[int] | tuple[int, ...]) -> None: if isinstance(seed, (list, tuple)): mixed_seed = mix_entropy(seed) else: @@ -80,10 +78,10 @@ class DpLoaderSet(Dataset): def __init__( self, - systems: Union[str, list[str]], + systems: str | list[str], batch_size: int, - type_map: Optional[list[str]], - seed: Optional[int] = None, + type_map: list[str] | None, + seed: int | None = None, shuffle: bool = True, ) -> None: if seed is not None: diff --git a/deepmd/pt/utils/dataset.py b/deepmd/pt/utils/dataset.py index 2cbe47cc3e..481fa04497 100644 --- a/deepmd/pt/utils/dataset.py +++ b/deepmd/pt/utils/dataset.py @@ -3,7 +3,6 @@ from typing import ( Any, - Optional, ) from torch.utils.data import ( @@ -17,7 +16,7 @@ class DeepmdDataSetForLoader(Dataset): - def __init__(self, system: str, type_map: Optional[list[str]] = None) -> None: + def __init__(self, system: str, type_map: list[str] | None = None) -> None: """Construct DeePMD-style dataset containing frames cross different systems. Args: diff --git a/deepmd/pt/utils/env_mat_stat.py b/deepmd/pt/utils/env_mat_stat.py index 1f89c09621..01822c7f3f 100644 --- a/deepmd/pt/utils/env_mat_stat.py +++ b/deepmd/pt/utils/env_mat_stat.py @@ -4,7 +4,6 @@ ) from typing import ( TYPE_CHECKING, - Union, ) import numpy as np @@ -77,7 +76,7 @@ def __init__(self, descriptor: "DescriptorBlock") -> None: ) # se_r=1, se_a=4 def iter( - self, data: list[dict[str, Union[torch.Tensor, list[tuple[int, int]]]]] + self, data: list[dict[str, torch.Tensor | list[tuple[int, int]]]] ) -> Iterator[dict[str, StatItem]]: """Get the iterator of the environment matrix. diff --git a/deepmd/pt/utils/multi_task.py b/deepmd/pt/utils/multi_task.py index 87b020c17b..f5e7799df7 100644 --- a/deepmd/pt/utils/multi_task.py +++ b/deepmd/pt/utils/multi_task.py @@ -4,7 +4,6 @@ ) from typing import ( Any, - Optional, ) from deepmd.pt.model.descriptor import ( @@ -107,7 +106,7 @@ def replace_one_item( key_type: str, key_in_dict: str, suffix: str = "", - index: Optional[int] = None, + index: int | None = None, ) -> None: shared_type = key_type shared_key = key_in_dict diff --git a/deepmd/pt/utils/neighbor_stat.py b/deepmd/pt/utils/neighbor_stat.py index b0e9eca141..292a27080b 100644 --- a/deepmd/pt/utils/neighbor_stat.py +++ b/deepmd/pt/utils/neighbor_stat.py @@ -2,9 +2,6 @@ from collections.abc import ( Iterator, ) -from typing import ( - Optional, -) import numpy as np import torch @@ -52,7 +49,7 @@ def forward( self, coord: torch.Tensor, atype: torch.Tensor, - cell: Optional[torch.Tensor], + cell: torch.Tensor | None, ) -> tuple[torch.Tensor, torch.Tensor]: """Calculate the neareest neighbor distance between atoms, maximum nbor size of atoms and the output data range of the environment matrix. @@ -170,7 +167,7 @@ def _execute( self, coord: np.ndarray, atype: np.ndarray, - cell: Optional[np.ndarray], + cell: np.ndarray | None, ) -> tuple[np.ndarray, np.ndarray]: """Execute the operation. diff --git a/deepmd/pt/utils/nlist.py b/deepmd/pt/utils/nlist.py index 8023645f8c..7f74e65f26 100644 --- a/deepmd/pt/utils/nlist.py +++ b/deepmd/pt/utils/nlist.py @@ -1,8 +1,4 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( - Optional, - Union, -) import torch @@ -21,7 +17,7 @@ def extend_input_and_build_neighbor_list( rcut: float, sel: list[int], mixed_types: bool = False, - box: Optional[torch.Tensor] = None, + box: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: nframes, nloc = atype.shape[:2] if box is not None: @@ -53,7 +49,7 @@ def build_neighbor_list( atype: torch.Tensor, nloc: int, rcut: float, - sel: Union[int, list[int]], + sel: int | list[int], distinguish_types: bool = True, ) -> torch.Tensor: """Build neighbor list for a single frame. keeps nsel neighbors. @@ -192,7 +188,7 @@ def build_directional_neighbor_list( coord_neig: torch.Tensor, atype_neig: torch.Tensor, rcut: float, - sel: Union[int, list[int]], + sel: int | list[int], distinguish_types: bool = True, ) -> torch.Tensor: """Build directional neighbor list. @@ -411,9 +407,9 @@ def build_multiple_neighbor_list( def extend_coord_with_ghosts( coord: torch.Tensor, atype: torch.Tensor, - cell: Optional[torch.Tensor], + cell: torch.Tensor | None, rcut: float, - cell_cpu: Optional[torch.Tensor] = None, + cell_cpu: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """Extend the coordinates of the atoms by appending peridoc images. The number of images is large enough to ensure all the neighbors diff --git a/deepmd/pt/utils/stat.py b/deepmd/pt/utils/stat.py index 7312d95a06..cf82461a7e 100644 --- a/deepmd/pt/utils/stat.py +++ b/deepmd/pt/utils/stat.py @@ -3,11 +3,11 @@ from collections import ( defaultdict, ) +from collections.abc import ( + Callable, +) from typing import ( Any, - Callable, - Optional, - Union, ) import numpy as np @@ -97,7 +97,7 @@ def make_stat_input( def _restore_from_file( stat_file_path: DPPath, keys: list[str] = ["energy"], -) -> Optional[dict]: +) -> dict | None: if stat_file_path is None: return None, None stat_files = [stat_file_path / f"bias_atom_{kk}" for kk in keys] @@ -159,7 +159,7 @@ def _post_process_stat( def _compute_model_predict( - sampled: Union[Callable[[], list[dict]], list[dict]], + sampled: Callable[[], list[dict]] | list[dict], keys: list[str], model_forward: Callable[..., torch.Tensor], ) -> dict[str, list[torch.Tensor]]: @@ -199,8 +199,8 @@ def model_forward_auto_batch_size(*args: Any, **kwargs: Any) -> Any: def _make_preset_out_bias( ntypes: int, - ibias: list[Optional[np.ndarray]], -) -> Optional[np.ndarray]: + ibias: list[np.ndarray | None], +) -> np.ndarray | None: """Make preset out bias. output: @@ -223,9 +223,9 @@ def _make_preset_out_bias( def _fill_stat_with_global( - atomic_stat: Union[np.ndarray, None], + atomic_stat: np.ndarray | None, global_stat: np.ndarray, -) -> Union[np.ndarray, None]: +) -> np.ndarray | None: """This function is used to fill atomic stat with global stat. Parameters @@ -249,13 +249,13 @@ def _fill_stat_with_global( def compute_output_stats( - merged: Union[Callable[[], list[dict]], list[dict]], + merged: Callable[[], list[dict]] | list[dict], ntypes: int, - keys: Union[str, list[str]] = ["energy"], - stat_file_path: Optional[DPPath] = None, - rcond: Optional[float] = None, - preset_bias: Optional[dict[str, list[Optional[np.ndarray]]]] = None, - model_forward: Optional[Callable[..., torch.Tensor]] = None, + keys: str | list[str] = ["energy"], + stat_file_path: DPPath | None = None, + rcond: float | None = None, + preset_bias: dict[str, list[np.ndarray | None]] | None = None, + model_forward: Callable[..., torch.Tensor] | None = None, stats_distinguish_types: bool = True, intensive: bool = False, ) -> dict[str, Any]: @@ -420,9 +420,9 @@ def compute_output_stats_global( sampled: list[dict], ntypes: int, keys: list[str], - rcond: Optional[float] = None, - preset_bias: Optional[dict[str, list[Optional[np.ndarray]]]] = None, - model_pred: Optional[dict[str, np.ndarray]] = None, + rcond: float | None = None, + preset_bias: dict[str, list[np.ndarray | None]] | None = None, + model_pred: dict[str, np.ndarray] | None = None, stats_distinguish_types: bool = True, intensive: bool = False, ) -> tuple[dict[str, np.ndarray], dict[str, np.ndarray]]: @@ -551,7 +551,7 @@ def compute_output_stats_atomic( sampled: list[dict], ntypes: int, keys: list[str], - model_pred: Optional[dict[str, np.ndarray]] = None, + model_pred: dict[str, np.ndarray] | None = None, ) -> tuple[dict[str, np.ndarray], dict[str, np.ndarray]]: # get label dict from sample; for each key, only picking the system with atomic labels. outputs = { diff --git a/deepmd/pt/utils/utils.py b/deepmd/pt/utils/utils.py index d06e2c1640..ab066bdf93 100644 --- a/deepmd/pt/utils/utils.py +++ b/deepmd/pt/utils/utils.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( Any, - Optional, - Union, overload, ) @@ -173,7 +171,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: class ActivationFn(torch.nn.Module): - def __init__(self, activation: Optional[str]) -> None: + def __init__(self, activation: str | None) -> None: super().__init__() self.activation: str = activation if activation is not None else "linear" if self.activation.lower().startswith( @@ -229,8 +227,8 @@ def to_numpy_array(xx: None) -> None: ... def to_numpy_array( - xx: Optional[torch.Tensor], -) -> Optional[np.ndarray]: + xx: torch.Tensor | None, +) -> np.ndarray | None: if xx is None: return None assert xx is not None @@ -256,8 +254,8 @@ def to_torch_tensor(xx: None) -> None: ... def to_torch_tensor( - xx: Optional[np.ndarray], -) -> Optional[torch.Tensor]: + xx: np.ndarray | None, +) -> torch.Tensor | None: if xx is None: return None assert xx is not None @@ -326,8 +324,8 @@ def mix_entropy(entropy_array: list[int]) -> int: def get_generator( - seed: Optional[Union[int, list[int]]] = None, -) -> Optional[torch.Generator]: + seed: int | list[int] | None = None, +) -> torch.Generator | None: if seed is not None: if isinstance(seed, list): seed = mix_entropy(seed) diff --git a/deepmd/tf/cluster/__init__.py b/deepmd/tf/cluster/__init__.py index 0f8916038d..a394b30f12 100644 --- a/deepmd/tf/cluster/__init__.py +++ b/deepmd/tf/cluster/__init__.py @@ -1,16 +1,12 @@ # SPDX-License-Identifier: LGPL-3.0-or-later """Module that reads node resources, auto detects if running local or on SLURM.""" -from typing import ( - Optional, -) - from .local import get_resource as get_local_res __all__ = ["get_resource"] -def get_resource() -> tuple[str, list[str], Optional[list[int]]]: +def get_resource() -> tuple[str, list[str], list[int] | None]: """Get local or slurm resources: nodename, nodelist, and gpus. Returns diff --git a/deepmd/tf/cluster/local.py b/deepmd/tf/cluster/local.py index 25fb1cc645..0d1af07d68 100644 --- a/deepmd/tf/cluster/local.py +++ b/deepmd/tf/cluster/local.py @@ -3,9 +3,6 @@ import subprocess as sp import sys -from typing import ( - Optional, -) from deepmd.tf.env import ( tf, @@ -49,7 +46,7 @@ def get_gpus(): return list(range(num_gpus)) if num_gpus > 0 else None -def get_resource() -> tuple[str, list[str], Optional[list[int]]]: +def get_resource() -> tuple[str, list[str], list[int] | None]: """Get local resources: nodename, nodelist, and gpus. Returns diff --git a/deepmd/tf/common.py b/deepmd/tf/common.py index 985c36c686..bf2681594f 100644 --- a/deepmd/tf/common.py +++ b/deepmd/tf/common.py @@ -2,13 +2,15 @@ """Collection of functions and classes used throughout the whole package.""" import warnings +from collections.abc import ( + Callable, +) from functools import ( wraps, ) from typing import ( TYPE_CHECKING, Any, - Callable, Union, ) diff --git a/deepmd/tf/descriptor/descriptor.py b/deepmd/tf/descriptor/descriptor.py index bd1af8c72e..e0b61ebb7e 100644 --- a/deepmd/tf/descriptor/descriptor.py +++ b/deepmd/tf/descriptor/descriptor.py @@ -4,7 +4,6 @@ ) from typing import ( Any, - Optional, ) import numpy as np @@ -173,7 +172,7 @@ def build( box_: tf.Tensor, mesh: tf.Tensor, input_dict: dict[str, Any], - reuse: Optional[bool] = None, + reuse: bool | None = None, suffix: str = "", ) -> tf.Tensor: """Build the computational graph for the descriptor. @@ -253,7 +252,7 @@ def enable_compression( f"Descriptor {type(self).__name__} doesn't support compression!" ) - def enable_mixed_precision(self, mixed_prec: Optional[dict] = None) -> None: + def enable_mixed_precision(self, mixed_prec: dict | None = None) -> None: """Receive the mixed precision setting. Parameters @@ -466,9 +465,9 @@ def explicit_ntypes(self) -> bool: def update_sel( cls, train_data: DeepmdDataSystem, - type_map: Optional[list[str]], + type_map: list[str] | None, local_jdata: dict, - ) -> tuple[dict, Optional[float]]: + ) -> tuple[dict, float | None]: """Update the selection and perform neighbor statistics. Parameters diff --git a/deepmd/tf/descriptor/hybrid.py b/deepmd/tf/descriptor/hybrid.py index 57c21f0ee6..1611344c9a 100644 --- a/deepmd/tf/descriptor/hybrid.py +++ b/deepmd/tf/descriptor/hybrid.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( Any, - Optional, - Union, ) import numpy as np @@ -45,9 +43,9 @@ class DescrptHybrid(Descriptor): def __init__( self, - list: list[Union[Descriptor, dict[str, Any]]], - ntypes: Optional[int] = None, - spin: Optional[Spin] = None, + list: list[Descriptor | dict[str, Any]], + ntypes: int | None = None, + spin: Spin | None = None, **kwargs, ) -> None: """Constructor.""" @@ -143,7 +141,7 @@ def compute_input_stats( mesh: list, input_dict: dict, mixed_type: bool = False, - real_natoms_vec: Optional[list] = None, + real_natoms_vec: list | None = None, **kwargs, ) -> None: """Compute the statisitcs (avg and std) of the training data. The input will be normalized by the statistics. @@ -213,7 +211,7 @@ def build( box_: tf.Tensor, mesh: tf.Tensor, input_dict: dict, - reuse: Optional[bool] = None, + reuse: bool | None = None, suffix: str = "", ) -> tf.Tensor: """Build the computational graph for the descriptor. @@ -351,7 +349,7 @@ def enable_compression( suffix=f"{suffix}_{idx}", ) - def enable_mixed_precision(self, mixed_prec: Optional[dict] = None) -> None: + def enable_mixed_precision(self, mixed_prec: dict | None = None) -> None: """Receive the mixed precision setting. Parameters @@ -426,9 +424,9 @@ def explicit_ntypes(self) -> bool: def update_sel( cls, train_data: DeepmdDataSystem, - type_map: Optional[list[str]], + type_map: list[str] | None, local_jdata: dict, - ) -> tuple[dict, Optional[float]]: + ) -> tuple[dict, float | None]: """Update the selection and perform neighbor statistics. Parameters diff --git a/deepmd/tf/descriptor/loc_frame.py b/deepmd/tf/descriptor/loc_frame.py index 9b338a5d25..87db875038 100644 --- a/deepmd/tf/descriptor/loc_frame.py +++ b/deepmd/tf/descriptor/loc_frame.py @@ -1,7 +1,4 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( - Optional, -) import numpy as np @@ -218,7 +215,7 @@ def build( box_: tf.Tensor, mesh: tf.Tensor, input_dict: dict, - reuse: Optional[bool] = None, + reuse: bool | None = None, suffix: str = "", ) -> tf.Tensor: """Build the computational graph for the descriptor. @@ -435,9 +432,9 @@ def init_variables( def update_sel( cls, train_data: DeepmdDataSystem, - type_map: Optional[list[str]], + type_map: list[str] | None, local_jdata: dict, - ) -> tuple[dict, Optional[float]]: + ) -> tuple[dict, float | None]: """Update the selection and perform neighbor statistics. Parameters diff --git a/deepmd/tf/descriptor/se.py b/deepmd/tf/descriptor/se.py index 5b04c5ba00..e62885eb62 100644 --- a/deepmd/tf/descriptor/se.py +++ b/deepmd/tf/descriptor/se.py @@ -1,8 +1,5 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import re -from typing import ( - Optional, -) from deepmd.dpmodel.utils.network import ( EmbeddingNet, @@ -154,9 +151,9 @@ def precision(self) -> tf.DType: def update_sel( cls, train_data: DeepmdDataSystem, - type_map: Optional[list[str]], + type_map: list[str] | None, local_jdata: dict, - ) -> tuple[dict, Optional[float]]: + ) -> tuple[dict, float | None]: """Update the selection and perform neighbor statistics. Parameters diff --git a/deepmd/tf/descriptor/se_a.py b/deepmd/tf/descriptor/se_a.py index ba21925b13..3d2ab39c9f 100644 --- a/deepmd/tf/descriptor/se_a.py +++ b/deepmd/tf/descriptor/se_a.py @@ -1,7 +1,4 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( - Optional, -) import numpy as np @@ -172,16 +169,16 @@ def __init__( axis_neuron: int = 8, resnet_dt: bool = False, trainable: bool = True, - seed: Optional[int] = None, + seed: int | None = None, type_one_side: bool = True, exclude_types: list[list[int]] = [], set_davg_zero: bool = False, activation_function: str = "tanh", precision: str = "default", uniform_seed: bool = False, - spin: Optional[Spin] = None, + spin: Spin | None = None, tebd_input_mode: str = "concat", - type_map: Optional[list[str]] = None, # to be compat with input + type_map: list[str] | None = None, # to be compat with input env_protection: float = 0.0, # not implement!! **kwargs, ) -> None: @@ -547,7 +544,7 @@ def enable_compression( self.davg = get_tensor_by_name_from_graph(graph, f"descrpt_attr{suffix}/t_avg") self.dstd = get_tensor_by_name_from_graph(graph, f"descrpt_attr{suffix}/t_std") - def enable_mixed_precision(self, mixed_prec: Optional[dict] = None) -> None: + def enable_mixed_precision(self, mixed_prec: dict | None = None) -> None: """Receive the mixed precision setting. Parameters @@ -566,7 +563,7 @@ def build( box_: tf.Tensor, mesh: tf.Tensor, input_dict: dict, - reuse: Optional[bool] = None, + reuse: bool | None = None, suffix: str = "", ) -> tf.Tensor: """Build the computational graph for the descriptor. diff --git a/deepmd/tf/descriptor/se_a_ebd.py b/deepmd/tf/descriptor/se_a_ebd.py index ae76308e69..189a013c13 100644 --- a/deepmd/tf/descriptor/se_a_ebd.py +++ b/deepmd/tf/descriptor/se_a_ebd.py @@ -1,7 +1,4 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( - Optional, -) import numpy as np @@ -78,7 +75,7 @@ def __init__( axis_neuron: int = 8, resnet_dt: bool = False, trainable: bool = True, - seed: Optional[int] = None, + seed: int | None = None, type_one_side: bool = True, type_nchanl: int = 2, type_nlayer: int = 1, @@ -118,7 +115,7 @@ def build( box_: tf.Tensor, mesh: tf.Tensor, input_dict: dict, - reuse: Optional[bool] = None, + reuse: bool | None = None, suffix: str = "", ) -> tf.Tensor: """Build the computational graph for the descriptor. diff --git a/deepmd/tf/descriptor/se_a_ebd_v2.py b/deepmd/tf/descriptor/se_a_ebd_v2.py index 035fc6509c..65e28e99d9 100644 --- a/deepmd/tf/descriptor/se_a_ebd_v2.py +++ b/deepmd/tf/descriptor/se_a_ebd_v2.py @@ -1,8 +1,5 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import logging -from typing import ( - Optional, -) from deepmd.tf.utils.spin import ( Spin, @@ -35,14 +32,14 @@ def __init__( axis_neuron: int = 8, resnet_dt: bool = False, trainable: bool = True, - seed: Optional[int] = None, + seed: int | None = None, type_one_side: bool = True, exclude_types: list[list[int]] = [], set_davg_zero: bool = False, activation_function: str = "tanh", precision: str = "default", uniform_seed: bool = False, - spin: Optional[Spin] = None, + spin: Spin | None = None, **kwargs, ) -> None: DescrptSeA.__init__( diff --git a/deepmd/tf/descriptor/se_a_ef.py b/deepmd/tf/descriptor/se_a_ef.py index d5476cbe89..37aa830431 100644 --- a/deepmd/tf/descriptor/se_a_ef.py +++ b/deepmd/tf/descriptor/se_a_ef.py @@ -1,7 +1,4 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( - Optional, -) import numpy as np @@ -77,7 +74,7 @@ def __init__( axis_neuron: int = 8, resnet_dt: bool = False, trainable: bool = True, - seed: Optional[int] = None, + seed: int | None = None, type_one_side: bool = True, exclude_types: list[list[int]] = [], set_davg_zero: bool = False, @@ -207,7 +204,7 @@ def build( box_: tf.Tensor, mesh: tf.Tensor, input_dict: dict, - reuse: Optional[bool] = None, + reuse: bool | None = None, suffix: str = "", ) -> tf.Tensor: """Build the computational graph for the descriptor. @@ -308,7 +305,7 @@ def __init__( axis_neuron: int = 8, resnet_dt: bool = False, trainable: bool = True, - seed: Optional[int] = None, + seed: int | None = None, type_one_side: bool = True, exclude_types: list[list[int]] = [], set_davg_zero: bool = False, diff --git a/deepmd/tf/descriptor/se_a_mask.py b/deepmd/tf/descriptor/se_a_mask.py index 5667122809..a0c5adaf76 100644 --- a/deepmd/tf/descriptor/se_a_mask.py +++ b/deepmd/tf/descriptor/se_a_mask.py @@ -2,7 +2,6 @@ import warnings from typing import ( Any, - Optional, ) import numpy as np @@ -124,7 +123,7 @@ def __init__( trainable: bool = True, type_one_side: bool = False, exclude_types: list[list[int]] = [], - seed: Optional[int] = None, + seed: int | None = None, activation_function: str = "tanh", precision: str = "default", uniform_seed: bool = False, @@ -269,7 +268,7 @@ def build( box_: tf.Tensor, mesh: tf.Tensor, input_dict: dict[str, Any], - reuse: Optional[bool] = None, + reuse: bool | None = None, suffix: str = "", ) -> tf.Tensor: """Build the computational graph for the descriptor. @@ -427,9 +426,9 @@ def prod_force_virial( def update_sel( cls, train_data: DeepmdDataSystem, - type_map: Optional[list[str]], + type_map: list[str] | None, local_jdata: dict, - ) -> tuple[dict, Optional[float]]: + ) -> tuple[dict, float | None]: """Update the selection and perform neighbor statistics. Parameters diff --git a/deepmd/tf/descriptor/se_atten.py b/deepmd/tf/descriptor/se_atten.py index 002e7bd3d3..3978709715 100644 --- a/deepmd/tf/descriptor/se_atten.py +++ b/deepmd/tf/descriptor/se_atten.py @@ -4,8 +4,6 @@ import warnings from typing import ( Any, - Optional, - Union, ) import numpy as np @@ -172,13 +170,13 @@ def __init__( self, rcut: float, rcut_smth: float, - sel: Union[list[int], int], + sel: list[int] | int, ntypes: int, neuron: list[int] = [25, 50, 100], axis_neuron: int = 8, resnet_dt: bool = False, trainable: bool = True, - seed: Optional[int] = None, + seed: int | None = None, type_one_side: bool = True, set_davg_zero: bool = True, exclude_types: list[list[int]] = [], @@ -196,11 +194,11 @@ def __init__( normalize=True, temperature=None, trainable_ln: bool = True, - ln_eps: Optional[float] = 1e-3, + ln_eps: float | None = 1e-3, concat_output_tebd: bool = True, env_protection: float = 0.0, # not implement!! - stripped_type_embedding: Optional[bool] = None, - type_map: Optional[list[str]] = None, # to be compat with input + stripped_type_embedding: bool | None = None, + type_map: list[str] | None = None, # to be compat with input **kwargs, ) -> None: # Ensure compatibility with the deprecated stripped_type_embedding option. @@ -343,7 +341,7 @@ def compute_input_stats( mesh: list, input_dict: dict, mixed_type: bool = False, - real_natoms_vec: Optional[list] = None, + real_natoms_vec: list | None = None, **kwargs, ) -> None: """Compute the statisitcs (avg and std) of the training data. The input will be normalized by the statistics. @@ -522,7 +520,7 @@ def build( box_: tf.Tensor, mesh: tf.Tensor, input_dict: dict, - reuse: Optional[bool] = None, + reuse: bool | None = None, suffix: str = "", ) -> tf.Tensor: """Build the computational graph for the descriptor. @@ -1513,9 +1511,9 @@ def explicit_ntypes(self) -> bool: def update_sel( cls, train_data: DeepmdDataSystem, - type_map: Optional[list[str]], + type_map: list[str] | None, local_jdata: dict, - ) -> tuple[dict, Optional[float]]: + ) -> tuple[dict, float | None]: """Update the selection and perform neighbor statistics. Parameters @@ -2184,7 +2182,7 @@ def __init__( self, rcut: float, rcut_smth: float, - sel: Union[list[int], int], + sel: list[int] | int, ntypes: int, neuron: list[int] = [25, 50, 100], axis_neuron: int = 8, @@ -2204,17 +2202,17 @@ def __init__( precision: str = "default", scaling_factor=1.0, normalize: bool = True, - temperature: Optional[float] = None, + temperature: float | None = None, trainable_ln: bool = True, - ln_eps: Optional[float] = 1e-3, + ln_eps: float | None = 1e-3, smooth_type_embedding: bool = True, concat_output_tebd: bool = True, use_econf_tebd: bool = False, use_tebd_bias: bool = False, - type_map: Optional[list[str]] = None, - spin: Optional[Any] = None, + type_map: list[str] | None = None, + spin: Any | None = None, # consistent with argcheck, not used though - seed: Optional[int] = None, + seed: int | None = None, uniform_seed: bool = False, ) -> None: if not normalize: @@ -2302,7 +2300,7 @@ def build( box_: tf.Tensor, mesh: tf.Tensor, input_dict: dict, - reuse: Optional[bool] = None, + reuse: bool | None = None, suffix: str = "", ) -> tf.Tensor: type_embedding = self.type_embedding.build(self.ntypes, suffix=suffix) diff --git a/deepmd/tf/descriptor/se_atten_v2.py b/deepmd/tf/descriptor/se_atten_v2.py index 69efe004c4..269fa1c083 100644 --- a/deepmd/tf/descriptor/se_atten_v2.py +++ b/deepmd/tf/descriptor/se_atten_v2.py @@ -1,8 +1,5 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import logging -from typing import ( - Optional, -) from deepmd.tf.utils.type_embed import ( TypeEmbedNet, @@ -77,7 +74,7 @@ def __init__( axis_neuron: int = 8, resnet_dt: bool = False, trainable: bool = True, - seed: Optional[int] = None, + seed: int | None = None, type_one_side: bool = True, set_davg_zero: bool = False, exclude_types: list[list[int]] = [], diff --git a/deepmd/tf/descriptor/se_r.py b/deepmd/tf/descriptor/se_r.py index 9613a9fa9b..c38a13d35a 100644 --- a/deepmd/tf/descriptor/se_r.py +++ b/deepmd/tf/descriptor/se_r.py @@ -1,7 +1,4 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( - Optional, -) import numpy as np @@ -95,15 +92,15 @@ def __init__( neuron: list[int] = [24, 48, 96], resnet_dt: bool = False, trainable: bool = True, - seed: Optional[int] = None, + seed: int | None = None, type_one_side: bool = True, exclude_types: list[list[int]] = [], set_davg_zero: bool = False, activation_function: str = "tanh", precision: str = "default", uniform_seed: bool = False, - spin: Optional[Spin] = None, - type_map: Optional[list[str]] = None, # to be compat with input + spin: Spin | None = None, + type_map: list[str] | None = None, # to be compat with input env_protection: float = 0.0, # not implement!! **kwargs, ) -> None: @@ -382,7 +379,7 @@ def build( box_: tf.Tensor, mesh: tf.Tensor, input_dict: dict, - reuse: Optional[bool] = None, + reuse: bool | None = None, suffix: str = "", ) -> tf.Tensor: """Build the computational graph for the descriptor. diff --git a/deepmd/tf/descriptor/se_t.py b/deepmd/tf/descriptor/se_t.py index ec6a1122d6..b03746a9c6 100644 --- a/deepmd/tf/descriptor/se_t.py +++ b/deepmd/tf/descriptor/se_t.py @@ -1,8 +1,5 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import re -from typing import ( - Optional, -) import numpy as np @@ -99,13 +96,13 @@ def __init__( neuron: list[int] = [24, 48, 96], resnet_dt: bool = False, trainable: bool = True, - seed: Optional[int] = None, + seed: int | None = None, exclude_types: list[list[int]] = [], set_davg_zero: bool = False, activation_function: str = "tanh", precision: str = "default", uniform_seed: bool = False, - type_map: Optional[list[str]] = None, # to be compat with input + type_map: list[str] | None = None, # to be compat with input env_protection: float = 0.0, # not implement!! **kwargs, ) -> None: @@ -387,7 +384,7 @@ def build( box_: tf.Tensor, mesh: tf.Tensor, input_dict: dict, - reuse: Optional[bool] = None, + reuse: bool | None = None, suffix: str = "", ) -> tf.Tensor: """Build the computational graph for the descriptor. diff --git a/deepmd/tf/entrypoints/change_bias.py b/deepmd/tf/entrypoints/change_bias.py index efb4f9ae35..adf65c0e2b 100644 --- a/deepmd/tf/entrypoints/change_bias.py +++ b/deepmd/tf/entrypoints/change_bias.py @@ -8,9 +8,6 @@ from pathlib import ( Path, ) -from typing import ( - Optional, -) import numpy as np @@ -54,12 +51,12 @@ def change_bias( INPUT: str, mode: str = "change", - bias_value: Optional[list] = None, - datafile: Optional[str] = None, + bias_value: list | None = None, + datafile: str | None = None, system: str = ".", numb_batch: int = 0, - model_branch: Optional[str] = None, - output: Optional[str] = None, + model_branch: str | None = None, + output: str | None = None, log_level: int = 0, **kwargs, ) -> None: @@ -131,12 +128,12 @@ def change_bias( def _change_bias_checkpoint_file( checkpoint_prefix: str, mode: str, - bias_value: Optional[list], - datafile: Optional[str], + bias_value: list | None, + datafile: str | None, system: str, numb_batch: int, - model_branch: Optional[str], - output: Optional[str], + model_branch: str | None, + output: str | None, log_level: int, ) -> None: """Change bias for individual checkpoint files.""" @@ -255,12 +252,12 @@ def _change_bias_checkpoint_file( def _change_bias_frozen_model( frozen_model_path: str, mode: str, - bias_value: Optional[list], - datafile: Optional[str], + bias_value: list | None, + datafile: str | None, system: str, numb_batch: int, - model_branch: Optional[str], - output: Optional[str], + model_branch: str | None, + output: str | None, log_level: int, ) -> None: """Change bias for frozen model (.pb file).""" @@ -285,7 +282,7 @@ def _change_bias_frozen_model( def _load_data_systems( - datafile: Optional[str], system: str, trainer: DPTrainer + datafile: str | None, system: str, trainer: DPTrainer ) -> DeepmdDataSystem: """Load data systems for bias calculation.""" if datafile is not None: diff --git a/deepmd/tf/entrypoints/compress.py b/deepmd/tf/entrypoints/compress.py index e8ceec7e9c..96bbd2c416 100644 --- a/deepmd/tf/entrypoints/compress.py +++ b/deepmd/tf/entrypoints/compress.py @@ -4,9 +4,6 @@ import json import logging import os -from typing import ( - Optional, -) from deepmd.tf.common import ( j_loader, @@ -58,7 +55,7 @@ def compress( checkpoint_folder: str, training_script: str, mpi_log: str, - log_path: Optional[str], + log_path: str | None, log_level: int, **kwargs, ) -> None: diff --git a/deepmd/tf/entrypoints/freeze.py b/deepmd/tf/entrypoints/freeze.py index 83d2ff33f4..366dba9c31 100755 --- a/deepmd/tf/entrypoints/freeze.py +++ b/deepmd/tf/entrypoints/freeze.py @@ -14,10 +14,6 @@ from pathlib import ( Path, ) -from typing import ( - Optional, - Union, -) import google.protobuf.message @@ -76,9 +72,9 @@ def _transfer_fitting_net_trainable_variables(sess, old_graph_def, raw_graph_def def _make_node_names( model_type: str, - modifier_type: Optional[str] = None, + modifier_type: str | None = None, out_suffix: str = "", - node_names: Optional[Union[str, list]] = None, + node_names: str | list | None = None, ) -> list[str]: """Get node names based on model type. @@ -297,8 +293,8 @@ def freeze( *, checkpoint_folder: str, output: str, - node_names: Optional[str] = None, - nvnmd_weight: Optional[str] = None, + node_names: str | None = None, + nvnmd_weight: str | None = None, **kwargs, ) -> None: """Freeze the graph in supplied folder. diff --git a/deepmd/tf/entrypoints/main.py b/deepmd/tf/entrypoints/main.py index ac2edc8ddd..8d83eac826 100644 --- a/deepmd/tf/entrypoints/main.py +++ b/deepmd/tf/entrypoints/main.py @@ -5,10 +5,6 @@ from pathlib import ( Path, ) -from typing import ( - Optional, - Union, -) from deepmd.backend.suffix import ( format_model_suffix, @@ -39,7 +35,7 @@ __all__ = ["get_ll", "main", "main_parser", "parse_args"] -def main(args: Optional[Union[list[str], argparse.Namespace]] = None) -> None: +def main(args: list[str] | argparse.Namespace | None = None) -> None: """DeePMD-Kit entry point. Parameters diff --git a/deepmd/tf/entrypoints/train.py b/deepmd/tf/entrypoints/train.py index 5bcca9a4e3..3ab55e190c 100755 --- a/deepmd/tf/entrypoints/train.py +++ b/deepmd/tf/entrypoints/train.py @@ -10,7 +10,6 @@ import time from typing import ( Any, - Optional, ) from deepmd.common import ( @@ -55,16 +54,16 @@ def train( *, INPUT: str, - init_model: Optional[str], - restart: Optional[str], + init_model: str | None, + restart: str | None, output: str, init_frz_model: str, mpi_log: str, log_level: int, - log_path: Optional[str], + log_path: str | None, is_compress: bool = False, skip_neighbor_stat: bool = False, - finetune: Optional[str] = None, + finetune: str | None = None, use_pretrain_script: bool = False, **kwargs, ) -> None: @@ -276,7 +275,7 @@ def _do_work( def get_modifier(modi_data=None): - modifier: Optional[BaseModifier] + modifier: BaseModifier | None if modi_data is not None: modifier_params = copy.deepcopy(modi_data) modifier_type = modifier_params.pop("type") diff --git a/deepmd/tf/entrypoints/transfer.py b/deepmd/tf/entrypoints/transfer.py index 0ad022776a..e441322764 100644 --- a/deepmd/tf/entrypoints/transfer.py +++ b/deepmd/tf/entrypoints/transfer.py @@ -6,9 +6,6 @@ from collections.abc import ( Sequence, ) -from typing import ( - Optional, -) import numpy as np @@ -32,7 +29,7 @@ def convert_number(number: int) -> float: def convert_matrix( - matrix: np.ndarray, shape: Sequence[int], dtype: Optional[type] = None + matrix: np.ndarray, shape: Sequence[int], dtype: type | None = None ) -> np.ndarray: """Convert matrix of integers to self defined binary format. @@ -188,7 +185,7 @@ def __init__(self, node) -> None: self.node = node def from_array( - self, tensor: np.ndarray, dtype: type, shape: Optional[Sequence[int]] = None + self, tensor: np.ndarray, dtype: type, shape: Sequence[int] | None = None ) -> None: if shape is None: shape = tensor.shape diff --git a/deepmd/tf/fit/dipole.py b/deepmd/tf/fit/dipole.py index 09da354116..961198b8e7 100644 --- a/deepmd/tf/fit/dipole.py +++ b/deepmd/tf/fit/dipole.py @@ -1,7 +1,4 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( - Optional, -) import numpy as np @@ -97,15 +94,15 @@ def __init__( numb_fparam: int = 0, numb_aparam: int = 0, dim_case_embd: int = 0, - sel_type: Optional[list[int]] = None, - seed: Optional[int] = None, + sel_type: list[int] | None = None, + seed: int | None = None, activation_function: str = "tanh", precision: str = "default", uniform_seed: bool = False, mixed_types: bool = False, - type_map: Optional[list[str]] = None, # to be compat with input - default_fparam: Optional[list[float]] = None, # to be compat with input - trainable: Optional[list[bool]] = None, + type_map: list[str] | None = None, # to be compat with input + default_fparam: list[float] | None = None, # to be compat with input + trainable: list[bool] | None = None, **kwargs, ) -> None: """Constructor.""" @@ -240,8 +237,8 @@ def build( input_d: tf.Tensor, rot_mat: tf.Tensor, natoms: tf.Tensor, - input_dict: Optional[dict] = None, - reuse: Optional[bool] = None, + input_dict: dict | None = None, + reuse: bool | None = None, suffix: str = "", ) -> tf.Tensor: """Build the computational graph for fitting net. @@ -373,7 +370,7 @@ def init_variables( graph_def, suffix=suffix ) - def enable_mixed_precision(self, mixed_prec: Optional[dict] = None) -> None: + def enable_mixed_precision(self, mixed_prec: dict | None = None) -> None: """Receive the mixed precision setting. Parameters diff --git a/deepmd/tf/fit/dos.py b/deepmd/tf/fit/dos.py index 81d80a4f4d..250d803d8f 100644 --- a/deepmd/tf/fit/dos.py +++ b/deepmd/tf/fit/dos.py @@ -1,8 +1,5 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import logging -from typing import ( - Optional, -) import numpy as np @@ -118,17 +115,17 @@ def __init__( numb_aparam: int = 0, dim_case_embd: int = 0, numb_dos: int = 300, - rcond: Optional[float] = None, - trainable: Optional[list[bool]] = None, - seed: Optional[int] = None, + rcond: float | None = None, + trainable: list[bool] | None = None, + seed: int | None = None, activation_function: str = "tanh", precision: str = "default", uniform_seed: bool = False, - layer_name: Optional[list[Optional[str]]] = None, + layer_name: list[str | None] | None = None, use_aparam_as_mask: bool = False, mixed_types: bool = False, - type_map: Optional[list[str]] = None, # to be compat with input - default_fparam: Optional[list[float]] = None, # to be compat with input + type_map: list[str] | None = None, # to be compat with input + default_fparam: list[float] | None = None, # to be compat with input **kwargs, ) -> None: """Constructor.""" @@ -406,8 +403,8 @@ def build( self, inputs: tf.Tensor, natoms: tf.Tensor, - input_dict: Optional[dict] = None, - reuse: Optional[bool] = None, + input_dict: dict | None = None, + reuse: bool | None = None, suffix: str = "", ) -> tf.Tensor: """Build the computational graph for fitting net. @@ -640,7 +637,7 @@ def init_variables( # for compatibility, old models has no t_bias_dos pass - def enable_mixed_precision(self, mixed_prec: Optional[dict] = None) -> None: + def enable_mixed_precision(self, mixed_prec: dict | None = None) -> None: """Receive the mixed precision setting. Parameters diff --git a/deepmd/tf/fit/ener.py b/deepmd/tf/fit/ener.py index 547c0eefb1..2b8b1b906e 100644 --- a/deepmd/tf/fit/ener.py +++ b/deepmd/tf/fit/ener.py @@ -1,8 +1,5 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import logging -from typing import ( - Optional, -) import numpy as np @@ -164,20 +161,20 @@ def __init__( numb_fparam: int = 0, numb_aparam: int = 0, dim_case_embd: int = 0, - rcond: Optional[float] = None, + rcond: float | None = None, tot_ener_zero: bool = False, - trainable: Optional[list[bool]] = None, - seed: Optional[int] = None, + trainable: list[bool] | None = None, + seed: int | None = None, atom_ener: list[float] = [], activation_function: str = "tanh", precision: str = "default", uniform_seed: bool = False, - layer_name: Optional[list[Optional[str]]] = None, + layer_name: list[str | None] | None = None, use_aparam_as_mask: bool = False, - spin: Optional[Spin] = None, + spin: Spin | None = None, mixed_types: bool = False, - type_map: Optional[list[str]] = None, # to be compat with input - default_fparam: Optional[list[float]] = None, # to be compat with input + type_map: list[str] | None = None, # to be compat with input + default_fparam: list[float] | None = None, # to be compat with input **kwargs, ) -> None: """Constructor.""" @@ -478,8 +475,8 @@ def build( self, inputs: tf.Tensor, natoms: tf.Tensor, - input_dict: Optional[dict] = None, - reuse: Optional[bool] = None, + input_dict: dict | None = None, + reuse: bool | None = None, suffix: str = "", ) -> tf.Tensor: """Build the computational graph for fitting net. @@ -841,7 +838,7 @@ def change_energy_bias( ntest=ntest, ) - def enable_mixed_precision(self, mixed_prec: Optional[dict] = None) -> None: + def enable_mixed_precision(self, mixed_prec: dict | None = None) -> None: """Receive the mixed precision setting. Parameters diff --git a/deepmd/tf/fit/fitting.py b/deepmd/tf/fit/fitting.py index 0e109fea60..b33559f12f 100644 --- a/deepmd/tf/fit/fitting.py +++ b/deepmd/tf/fit/fitting.py @@ -3,9 +3,6 @@ from abc import ( abstractmethod, ) -from typing import ( - Optional, -) from deepmd.common import ( j_get_type, @@ -134,8 +131,8 @@ def serialize_network( activation_function: str, resnet_dt: bool, variables: dict, - out_dim: Optional[int] = 1, - trainable: Optional[list[bool]] = None, + out_dim: int | None = 1, + trainable: list[bool] | None = None, suffix: str = "", ) -> dict: """Serialize network. diff --git a/deepmd/tf/fit/polar.py b/deepmd/tf/fit/polar.py index 31ccd18302..1e48a5fa59 100644 --- a/deepmd/tf/fit/polar.py +++ b/deepmd/tf/fit/polar.py @@ -1,8 +1,5 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import warnings -from typing import ( - Optional, -) import numpy as np @@ -112,19 +109,19 @@ def __init__( numb_fparam: int = 0, numb_aparam: int = 0, dim_case_embd: int = 0, - sel_type: Optional[list[int]] = None, + sel_type: list[int] | None = None, fit_diag: bool = True, - scale: Optional[list[float]] = None, + scale: list[float] | None = None, shift_diag: bool = True, # YWolfeee: will support the user to decide whether to use this function # diag_shift : list[float] = None, YWolfeee: will not support the user to assign a shift - seed: Optional[int] = None, + seed: int | None = None, activation_function: str = "tanh", precision: str = "default", uniform_seed: bool = False, mixed_types: bool = False, - type_map: Optional[list[str]] = None, # to be compat with input - default_fparam: Optional[list[float]] = None, # to be compat with input - trainable: Optional[list[bool]] = None, + type_map: list[str] | None = None, # to be compat with input + default_fparam: list[float] | None = None, # to be compat with input + trainable: list[bool] | None = None, **kwargs, ) -> None: """Constructor.""" @@ -425,8 +422,8 @@ def build( input_d: tf.Tensor, rot_mat: tf.Tensor, natoms: tf.Tensor, - input_dict: Optional[dict] = None, - reuse: Optional[bool] = None, + input_dict: dict | None = None, + reuse: bool | None = None, suffix: str = "", ): """Build the computational graph for fitting net. @@ -607,7 +604,7 @@ def init_variables( stacklevel=2, ) - def enable_mixed_precision(self, mixed_prec: Optional[dict] = None) -> None: + def enable_mixed_precision(self, mixed_prec: dict | None = None) -> None: """Receive the mixed precision setting. Parameters @@ -750,11 +747,11 @@ def __init__( descrpt: tf.Tensor, neuron: list[int] = [120, 120, 120], resnet_dt: bool = True, - sel_type: Optional[list[int]] = None, + sel_type: list[int] | None = None, fit_diag: bool = True, - scale: Optional[list[float]] = None, - diag_shift: Optional[list[float]] = None, - seed: Optional[int] = None, + scale: list[float] | None = None, + diag_shift: list[float] | None = None, + seed: int | None = None, activation_function: str = "tanh", precision: str = "default", ) -> None: @@ -789,7 +786,7 @@ def build( input_d, rot_mat, natoms, - input_dict: Optional[dict] = None, + input_dict: dict | None = None, reuse=None, suffix="", ) -> tf.Tensor: @@ -849,7 +846,7 @@ def init_variables( graph=graph, graph_def=graph_def, suffix=suffix ) - def enable_mixed_precision(self, mixed_prec: Optional[dict] = None) -> None: + def enable_mixed_precision(self, mixed_prec: dict | None = None) -> None: """Receive the mixed precision setting. Parameters diff --git a/deepmd/tf/infer/deep_dipole.py b/deepmd/tf/infer/deep_dipole.py index 09252d1d6c..834c0b7d8f 100644 --- a/deepmd/tf/infer/deep_dipole.py +++ b/deepmd/tf/infer/deep_dipole.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( TYPE_CHECKING, - Optional, ) from deepmd.infer.deep_dipole import ( @@ -50,7 +49,7 @@ def __init__( model_file: "Path", load_prefix: str = "load", default_tf_graph: bool = False, - input_map: Optional[dict] = None, + input_map: dict | None = None, neighbor_list=None, ) -> None: # use this in favor of dict update to move attribute from class to diff --git a/deepmd/tf/infer/deep_eval.py b/deepmd/tf/infer/deep_eval.py index 75440accb9..d1eca83055 100644 --- a/deepmd/tf/infer/deep_eval.py +++ b/deepmd/tf/infer/deep_eval.py @@ -1,14 +1,14 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import json +from collections.abc import ( + Callable, +) from functools import ( cached_property, ) from typing import ( TYPE_CHECKING, Any, - Callable, - Optional, - Union, ) import numpy as np @@ -93,8 +93,8 @@ def __init__( *args: list, load_prefix: str = "load", default_tf_graph: bool = False, - auto_batch_size: Union[bool, int, AutoBatchSize] = False, - input_map: Optional[dict] = None, + auto_batch_size: bool | int | AutoBatchSize = False, + input_map: dict | None = None, neighbor_list=None, **kwargs: dict, ) -> None: @@ -358,7 +358,7 @@ def _load_graph( frozen_graph_filename: "Path", prefix: str = "load", default_tf_graph: bool = False, - input_map: Optional[dict] = None, + input_map: dict | None = None, ): # We load the protobuf file from the disk and parse it to retrieve the # unserialized graph_def @@ -393,7 +393,7 @@ def _load_graph( def sort_input( coord: np.ndarray, atom_type: np.ndarray, - sel_atoms: Optional[list[int]] = None, + sel_atoms: list[int] | None = None, ): """Sort atoms in the system according their types. @@ -532,7 +532,7 @@ def eval_typeebd(self) -> np.ndarray: def build_neighbor_list( self, coords: np.ndarray, - cell: Optional[np.ndarray], + cell: np.ndarray | None, atype: np.ndarray, imap: np.ndarray, neighbor_list, @@ -634,7 +634,7 @@ def get_type_map(self) -> list[str]: """Get the type map (element name of the atom types) of this model.""" return self.tmap - def get_sel_type(self) -> Optional[np.ndarray]: + def get_sel_type(self) -> np.ndarray | None: """Get the selected atom types of this model. Only atoms with selected atom types have atomic contribution @@ -682,7 +682,7 @@ def eval_func(*args, **kwargs): def _get_natoms_and_nframes( self, coords: np.ndarray, - atom_types: Union[list[int], np.ndarray], + atom_types: list[int] | np.ndarray, ) -> tuple[int, int]: natoms = len(atom_types[0]) if natoms == 0: @@ -695,12 +695,12 @@ def _get_natoms_and_nframes( def eval( self, coords: np.ndarray, - cells: Optional[np.ndarray], + cells: np.ndarray | None, atom_types: np.ndarray, atomic: bool = False, - fparam: Optional[np.ndarray] = None, - aparam: Optional[np.ndarray] = None, - efield: Optional[np.ndarray] = None, + fparam: np.ndarray | None = None, + aparam: np.ndarray | None = None, + efield: np.ndarray | None = None, **kwargs: Any, ) -> dict[str, np.ndarray]: """Evaluate the energy, force and virial by using this DP. @@ -1025,11 +1025,11 @@ def _get_output_shape(self, odef, nframes, natoms): def eval_descriptor( self, coords: np.ndarray, - cells: Optional[np.ndarray], + cells: np.ndarray | None, atom_types: np.ndarray, - fparam: Optional[np.ndarray] = None, - aparam: Optional[np.ndarray] = None, - efield: Optional[np.ndarray] = None, + fparam: np.ndarray | None = None, + aparam: np.ndarray | None = None, + efield: np.ndarray | None = None, ) -> np.ndarray: """Evaluate descriptors by using this DP. @@ -1082,11 +1082,11 @@ def eval_descriptor( def _eval_descriptor_inner( self, coords: np.ndarray, - cells: Optional[np.ndarray], + cells: np.ndarray | None, atom_types: np.ndarray, - fparam: Optional[np.ndarray] = None, - aparam: Optional[np.ndarray] = None, - efield: Optional[np.ndarray] = None, + fparam: np.ndarray | None = None, + aparam: np.ndarray | None = None, + efield: np.ndarray | None = None, ) -> np.ndarray: natoms, nframes = self._get_natoms_and_nframes( coords, @@ -1166,8 +1166,8 @@ def __init__( model_file: "Path", load_prefix: str = "load", default_tf_graph: bool = False, - auto_batch_size: Union[bool, int, AutoBatchSize] = False, - input_map: Optional[dict] = None, + auto_batch_size: bool | int | AutoBatchSize = False, + input_map: dict | None = None, neighbor_list=None, ) -> None: self.graph = self._load_graph( @@ -1254,9 +1254,7 @@ def _graph_compatable(self) -> bool: else: return True - def _get_tensor( - self, tensor_name: str, attr_name: Optional[str] = None - ) -> tf.Tensor: + def _get_tensor(self, tensor_name: str, attr_name: str | None = None) -> tf.Tensor: """Get TF graph tensor and assign it to class namespace. Parameters @@ -1286,7 +1284,7 @@ def _load_graph( frozen_graph_filename: "Path", prefix: str = "load", default_tf_graph: bool = False, - input_map: Optional[dict] = None, + input_map: dict | None = None, ): # We load the protobuf file from the disk and parse it to retrieve the # unserialized graph_def @@ -1321,7 +1319,7 @@ def _load_graph( def sort_input( coord: np.ndarray, atom_type: np.ndarray, - sel_atoms: Optional[list[int]] = None, + sel_atoms: list[int] | None = None, mixed_type: bool = False, ): """Sort atoms in the system according their types. @@ -1478,7 +1476,7 @@ def eval_typeebd(self) -> np.ndarray: def build_neighbor_list( self, coords: np.ndarray, - cell: Optional[np.ndarray], + cell: np.ndarray | None, atype: np.ndarray, imap: np.ndarray, neighbor_list, diff --git a/deepmd/tf/infer/deep_tensor.py b/deepmd/tf/infer/deep_tensor.py index a4d30c4f2f..48fa811ff9 100644 --- a/deepmd/tf/infer/deep_tensor.py +++ b/deepmd/tf/infer/deep_tensor.py @@ -2,7 +2,6 @@ from typing import ( TYPE_CHECKING, ClassVar, - Optional, ) import numpy as np @@ -59,7 +58,7 @@ def __init__( model_file: "Path", load_prefix: str = "load", default_tf_graph: bool = False, - input_map: Optional[dict] = None, + input_map: dict | None = None, neighbor_list=None, ) -> None: """Constructor.""" @@ -143,12 +142,12 @@ def get_dim_aparam(self) -> int: def eval( self, coords: np.ndarray, - cells: Optional[np.ndarray], + cells: np.ndarray | None, atom_types: list[int], atomic: bool = True, - fparam: Optional[np.ndarray] = None, - aparam: Optional[np.ndarray] = None, - efield: Optional[np.ndarray] = None, + fparam: np.ndarray | None = None, + aparam: np.ndarray | None = None, + efield: np.ndarray | None = None, mixed_type: bool = False, ) -> np.ndarray: """Evaluate the model. @@ -273,12 +272,12 @@ def eval( def eval_full( self, coords: np.ndarray, - cells: Optional[np.ndarray], + cells: np.ndarray | None, atom_types: list[int], atomic: bool = False, - fparam: Optional[np.array] = None, - aparam: Optional[np.array] = None, - efield: Optional[np.array] = None, + fparam: np.ndarray | None = None, + aparam: np.ndarray | None = None, + efield: np.ndarray | None = None, mixed_type: bool = False, ) -> tuple[np.ndarray, ...]: """Evaluate the model with interface similar to the energy model. diff --git a/deepmd/tf/loss/ener.py b/deepmd/tf/loss/ener.py index ba65450613..cba60bb7f0 100644 --- a/deepmd/tf/loss/ener.py +++ b/deepmd/tf/loss/ener.py @@ -1,7 +1,4 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( - Optional, -) import numpy as np @@ -98,7 +95,7 @@ def __init__( limit_pref_ae: float = 0.0, start_pref_pf: float = 0.0, limit_pref_pf: float = 0.0, - relative_f: Optional[float] = None, + relative_f: float | None = None, enable_atom_ener_coeff: bool = False, start_pref_gf: float = 0.0, limit_pref_gf: float = 0.0, @@ -556,9 +553,9 @@ def __init__( limit_pref_ae: float = 0.0, start_pref_pf: float = 0.0, limit_pref_pf: float = 0.0, - relative_f: Optional[float] = None, + relative_f: float | None = None, enable_atom_ener_coeff: bool = False, - use_spin: Optional[list] = None, + use_spin: list | None = None, ) -> None: self.starter_learning_rate = starter_learning_rate self.start_pref_e = start_pref_e diff --git a/deepmd/tf/model/dos.py b/deepmd/tf/model/dos.py index 264c77d045..548756f3e9 100644 --- a/deepmd/tf/model/dos.py +++ b/deepmd/tf/model/dos.py @@ -1,8 +1,4 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( - Optional, - Union, -) from deepmd.tf.env import ( MODEL_VERSION, @@ -49,8 +45,8 @@ def __init__( self, descriptor: dict, fitting_net: dict, - type_embedding: Optional[Union[dict, TypeEmbedNet]] = None, - type_map: Optional[list[str]] = None, + type_embedding: dict | TypeEmbedNet | None = None, + type_map: list[str] | None = None, data_stat_nbatch: int = 10, data_stat_protect: float = 1e-2, **kwargs, @@ -137,7 +133,7 @@ def build( mesh, input_dict, frz_model=None, - ckpt_meta: Optional[str] = None, + ckpt_meta: str | None = None, suffix="", reuse=None, ): diff --git a/deepmd/tf/model/ener.py b/deepmd/tf/model/ener.py index 5b665511ef..7109ad9d73 100644 --- a/deepmd/tf/model/ener.py +++ b/deepmd/tf/model/ener.py @@ -1,8 +1,4 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( - Optional, - Union, -) import numpy as np @@ -75,16 +71,16 @@ def __init__( self, descriptor: dict, fitting_net: dict, - type_embedding: Optional[Union[dict, TypeEmbedNet]] = None, - type_map: Optional[list[str]] = None, + type_embedding: dict | TypeEmbedNet | None = None, + type_map: list[str] | None = None, data_stat_nbatch: int = 10, data_stat_protect: float = 1e-2, - use_srtab: Optional[str] = None, - smin_alpha: Optional[float] = None, - sw_rmin: Optional[float] = None, - sw_rmax: Optional[float] = None, + use_srtab: str | None = None, + smin_alpha: float | None = None, + sw_rmin: float | None = None, + sw_rmax: float | None = None, srtab_add_bias: bool = True, - spin: Optional[Spin] = None, + spin: Spin | None = None, data_bias_nsample: int = 10, **kwargs, ) -> None: @@ -182,7 +178,7 @@ def build( mesh, input_dict, frz_model=None, - ckpt_meta: Optional[str] = None, + ckpt_meta: str | None = None, suffix="", reuse=None, ): diff --git a/deepmd/tf/model/frozen.py b/deepmd/tf/model/frozen.py index 6ca18ed7bd..b63fe80a87 100644 --- a/deepmd/tf/model/frozen.py +++ b/deepmd/tf/model/frozen.py @@ -7,8 +7,6 @@ ) from typing import ( NoReturn, - Optional, - Union, ) from deepmd.entrypoints.convert_backend import ( @@ -83,10 +81,10 @@ def build( box: tf.Tensor, mesh: tf.Tensor, input_dict: dict, - frz_model: Optional[str] = None, - ckpt_meta: Optional[str] = None, + frz_model: str | None = None, + ckpt_meta: str | None = None, suffix: str = "", - reuse: Optional[Union[bool, Enum]] = None, + reuse: bool | Enum | None = None, ) -> dict: """Build the model. @@ -185,11 +183,11 @@ def build( "Contribution is welcome!" ) - def get_fitting(self) -> Union[Fitting, dict]: + def get_fitting(self) -> Fitting | dict: """Get the fitting(s).""" return {} - def get_loss(self, loss: dict, lr) -> Optional[Union[Loss, dict]]: + def get_loss(self, loss: dict, lr) -> Loss | dict | None: """Get the loss function(s).""" # loss should be never used for a frozen model return @@ -243,9 +241,9 @@ def get_type_map(self) -> list: def update_sel( cls, train_data: DeepmdDataSystem, - type_map: Optional[list[str]], + type_map: list[str] | None, local_jdata: dict, - ) -> tuple[dict, Optional[float]]: + ) -> tuple[dict, float | None]: """Update the selection and perform neighbor statistics. Parameters diff --git a/deepmd/tf/model/linear.py b/deepmd/tf/model/linear.py index 63f55eae9e..f036edbfc3 100644 --- a/deepmd/tf/model/linear.py +++ b/deepmd/tf/model/linear.py @@ -7,10 +7,6 @@ lru_cache, reduce, ) -from typing import ( - Optional, - Union, -) from deepmd.tf.env import ( GLOBAL_TF_FLOAT_PRECISION, @@ -64,13 +60,13 @@ def __init__(self, models: list[dict], weights: list[float], **kwargs) -> None: else: raise ValueError(f"Invalid weights {weights}") - def get_fitting(self) -> Union[Fitting, dict]: + def get_fitting(self) -> Fitting | dict: """Get the fitting(s).""" return { f"model{ii}": model.get_fitting() for ii, model in enumerate(self.models) } - def get_loss(self, loss: dict, lr) -> Optional[Union[Loss, dict]]: + def get_loss(self, loss: dict, lr) -> Loss | dict | None: """Get the loss function(s).""" # the first model that is not None, or None if all models are None for model in self.models: @@ -138,9 +134,9 @@ def get_type_map(self) -> list: def update_sel( cls, train_data: DeepmdDataSystem, - type_map: Optional[list[str]], + type_map: list[str] | None, local_jdata: dict, - ) -> tuple[dict, Optional[float]]: + ) -> tuple[dict, float | None]: """Update the selection and perform neighbor statistics. Parameters @@ -194,10 +190,10 @@ def build( box: tf.Tensor, mesh: tf.Tensor, input_dict: dict, - frz_model: Optional[str] = None, - ckpt_meta: Optional[str] = None, + frz_model: str | None = None, + ckpt_meta: str | None = None, suffix: str = "", - reuse: Optional[Union[bool, Enum]] = None, + reuse: bool | Enum | None = None, ) -> dict: """Build the model. diff --git a/deepmd/tf/model/model.py b/deepmd/tf/model/model.py index 95c9840cde..da99cf5e60 100644 --- a/deepmd/tf/model/model.py +++ b/deepmd/tf/model/model.py @@ -9,8 +9,6 @@ ) from typing import ( NoReturn, - Optional, - Union, ) import numpy as np @@ -115,13 +113,13 @@ def __new__(cls, *args, **kwargs): def __init__( self, - type_embedding: Optional[Union[dict, TypeEmbedNet]] = None, - type_map: Optional[list[str]] = None, + type_embedding: dict | TypeEmbedNet | None = None, + type_map: list[str] | None = None, data_stat_nbatch: int = 10, data_bias_nsample: int = 10, data_stat_protect: float = 1e-2, - spin: Optional[Spin] = None, - compress: Optional[dict] = None, + spin: Spin | None = None, + compress: dict | None = None, **kwargs, ) -> None: super().__init__() @@ -155,10 +153,10 @@ def build( box: tf.Tensor, mesh: tf.Tensor, input_dict: dict, - frz_model: Optional[str] = None, - ckpt_meta: Optional[str] = None, + frz_model: str | None = None, + ckpt_meta: str | None = None, suffix: str = "", - reuse: Optional[Union[bool, Enum]] = None, + reuse: bool | Enum | None = None, ): """Build the model. @@ -223,10 +221,10 @@ def build_descrpt( box: tf.Tensor, mesh: tf.Tensor, input_dict: dict, - frz_model: Optional[str] = None, - ckpt_meta: Optional[str] = None, + frz_model: str | None = None, + ckpt_meta: str | None = None, suffix: str = "", - reuse: Optional[Union[bool, Enum]] = None, + reuse: bool | Enum | None = None, ): """Build the descriptor part of the model. @@ -311,10 +309,10 @@ def build_descrpt( def build_type_embedding( self, ntypes: int, - frz_model: Optional[str] = None, - ckpt_meta: Optional[str] = None, + frz_model: str | None = None, + ckpt_meta: str | None = None, suffix: str = "", - reuse: Optional[Union[bool, Enum]] = None, + reuse: bool | Enum | None = None, ) -> tf.Tensor: """Build the type embedding part of the model. @@ -432,24 +430,24 @@ def enable_compression(self, suffix: str = "") -> NoReturn: """ raise RuntimeError("Not supported") - def get_numb_fparam(self) -> Union[int, dict]: + def get_numb_fparam(self) -> int | dict: """Get the number of frame parameters.""" return 0 - def get_numb_aparam(self) -> Union[int, dict]: + def get_numb_aparam(self) -> int | dict: """Get the number of atomic parameters.""" return 0 - def get_numb_dos(self) -> Union[int, dict]: + def get_numb_dos(self) -> int | dict: """Get the number of gridpoints in energy space.""" return 0 @abstractmethod - def get_fitting(self) -> Union[Fitting, dict]: + def get_fitting(self) -> Fitting | dict: """Get the fitting(s).""" @abstractmethod - def get_loss(self, loss: dict, lr) -> Optional[Union[Loss, dict]]: + def get_loss(self, loss: dict, lr) -> Loss | dict | None: """Get the loss function(s).""" @abstractmethod @@ -518,9 +516,9 @@ def get_feed_dict( def update_sel( cls, train_data: DeepmdDataSystem, - type_map: Optional[list[str]], + type_map: list[str] | None, local_jdata: dict, - ) -> tuple[dict, Optional[float]]: + ) -> tuple[dict, float | None]: """Update the selection and perform neighbor statistics. Notes @@ -647,10 +645,10 @@ def __new__(cls, *args, **kwargs): def __init__( self, - descriptor: Union[dict, Descriptor], - fitting_net: Union[dict, Fitting], - type_embedding: Optional[Union[dict, TypeEmbedNet]] = None, - type_map: Optional[list[str]] = None, + descriptor: dict | Descriptor, + fitting_net: dict | Fitting, + type_embedding: dict | TypeEmbedNet | None = None, + type_map: list[str] | None = None, **kwargs, ) -> None: super().__init__( @@ -806,11 +804,11 @@ def enable_compression(self, suffix: str = "") -> None: ): self.typeebd.init_variables(graph, graph_def, suffix=suffix) - def get_fitting(self) -> Union[Fitting, dict]: + def get_fitting(self) -> Fitting | dict: """Get the fitting(s).""" return self.fitting - def get_loss(self, loss: dict, lr) -> Union[Loss, dict]: + def get_loss(self, loss: dict, lr) -> Loss | dict: """Get the loss function(s).""" return self.fitting.get_loss(loss, lr) @@ -950,9 +948,9 @@ def _apply_out_bias_std(self, output, atype, natoms, coord, selected_atype=None) def update_sel( cls, train_data: DeepmdDataSystem, - type_map: Optional[list[str]], + type_map: list[str] | None, local_jdata: dict, - ) -> tuple[dict, Optional[float]]: + ) -> tuple[dict, float | None]: """Update the selection and perform neighbor statistics. Parameters diff --git a/deepmd/tf/model/pairtab.py b/deepmd/tf/model/pairtab.py index d91554b98d..8bdf48accb 100644 --- a/deepmd/tf/model/pairtab.py +++ b/deepmd/tf/model/pairtab.py @@ -2,10 +2,6 @@ from enum import ( Enum, ) -from typing import ( - Optional, - Union, -) import numpy as np @@ -67,7 +63,7 @@ class PairTabModel(Model): model_type = "ener" def __init__( - self, tab_file: str, rcut: float, sel: Union[int, list[int]], **kwargs + self, tab_file: str, rcut: float, sel: int | list[int], **kwargs ) -> None: super().__init__() self.tab_file = tab_file @@ -89,10 +85,10 @@ def build( box: tf.Tensor, mesh: tf.Tensor, input_dict: dict, - frz_model: Optional[str] = None, - ckpt_meta: Optional[str] = None, + frz_model: str | None = None, + ckpt_meta: str | None = None, suffix: str = "", - reuse: Optional[Union[bool, Enum]] = None, + reuse: bool | Enum | None = None, ): """Build the model. @@ -237,12 +233,12 @@ def init_variables( """ # skip. table can be initialized from the file - def get_fitting(self) -> Union[Fitting, dict]: + def get_fitting(self) -> Fitting | dict: """Get the fitting(s).""" # nothing needs to do return {} - def get_loss(self, loss: dict, lr) -> Optional[Union[Loss, dict]]: + def get_loss(self, loss: dict, lr) -> Loss | dict | None: """Get the loss function(s).""" # nothing needs to do return @@ -273,9 +269,9 @@ def enable_compression(self, suffix: str = "") -> None: def update_sel( cls, train_data: DeepmdDataSystem, - type_map: Optional[list[str]], + type_map: list[str] | None, local_jdata: dict, - ) -> tuple[dict, Optional[float]]: + ) -> tuple[dict, float | None]: """Update the selection and perform neighbor statistics. Notes diff --git a/deepmd/tf/model/pairwise_dprc.py b/deepmd/tf/model/pairwise_dprc.py index 5ed98f0c49..45d0c754b9 100644 --- a/deepmd/tf/model/pairwise_dprc.py +++ b/deepmd/tf/model/pairwise_dprc.py @@ -1,8 +1,4 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( - Optional, - Union, -) from deepmd.tf.common import ( make_default_mesh, @@ -49,17 +45,17 @@ def __init__( self, qm_model: dict, qmmm_model: dict, - type_embedding: Union[dict, TypeEmbedNet], + type_embedding: dict | TypeEmbedNet, type_map: list[str], data_stat_nbatch: int = 10, data_stat_nsample: int = 10, data_stat_protect: float = 1e-2, - use_srtab: Optional[str] = None, - smin_alpha: Optional[float] = None, - sw_rmin: Optional[float] = None, - sw_rmax: Optional[float] = None, - spin: Optional[Spin] = None, - compress: Optional[dict] = None, + use_srtab: str | None = None, + smin_alpha: float | None = None, + sw_rmin: float | None = None, + sw_rmax: float | None = None, + spin: Spin | None = None, + compress: dict | None = None, **kwargs, ) -> None: # internal variable to compare old and new behavior @@ -116,9 +112,9 @@ def build( mesh: tf.Tensor, input_dict: dict, frz_model=None, - ckpt_meta: Optional[str] = None, + ckpt_meta: str | None = None, suffix: str = "", - reuse: Optional[bool] = None, + reuse: bool | None = None, ): feed_dict = self.get_feed_dict( coord_, atype_, natoms, box_, mesh, aparam=input_dict["aparam"] @@ -300,14 +296,14 @@ def build( model_dict["atype"] = atype_ return model_dict - def get_fitting(self) -> Union[str, dict]: + def get_fitting(self) -> str | dict: """Get the fitting(s).""" return { "qm": self.qm_model.get_fitting(), "qmmm": self.qmmm_model.get_fitting(), } - def get_loss(self, loss: dict, lr) -> Union[Loss, dict]: + def get_loss(self, loss: dict, lr) -> Loss | dict: """Get the loss function(s).""" return self.qm_model.get_loss(loss, lr) @@ -413,9 +409,9 @@ def get_feed_dict( def update_sel( cls, train_data: DeepmdDataSystem, - type_map: Optional[list[str]], + type_map: list[str] | None, local_jdata: dict, - ) -> tuple[dict, Optional[float]]: + ) -> tuple[dict, float | None]: """Update the selection and perform neighbor statistics. Parameters diff --git a/deepmd/tf/model/tensor.py b/deepmd/tf/model/tensor.py index 58af997464..70d9a3d26c 100644 --- a/deepmd/tf/model/tensor.py +++ b/deepmd/tf/model/tensor.py @@ -1,8 +1,4 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( - Optional, - Union, -) import numpy as np @@ -55,8 +51,8 @@ def __init__( tensor_name: str, descriptor: dict, fitting_net: dict, - type_embedding: Optional[Union[dict, TypeEmbedNet]] = None, - type_map: Optional[list[str]] = None, + type_embedding: dict | TypeEmbedNet | None = None, + type_map: list[str] | None = None, data_stat_nbatch: int = 10, data_stat_protect: float = 1e-2, **kwargs, @@ -119,7 +115,7 @@ def build( mesh, input_dict, frz_model=None, - ckpt_meta: Optional[str] = None, + ckpt_meta: str | None = None, suffix="", reuse=None, ): diff --git a/deepmd/tf/nvnmd/entrypoints/mapt.py b/deepmd/tf/nvnmd/entrypoints/mapt.py index 2e6e56bf51..121263974a 100644 --- a/deepmd/tf/nvnmd/entrypoints/mapt.py +++ b/deepmd/tf/nvnmd/entrypoints/mapt.py @@ -1,8 +1,5 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import logging -from typing import ( - Optional, -) import numpy as np @@ -649,9 +646,9 @@ def build_davg_dstd(self): def mapt( *, - nvnmd_config: Optional[str] = "nvnmd/config.npy", - nvnmd_weight: Optional[str] = "nvnmd/weight.npy", - nvnmd_map: Optional[str] = "nvnmd/map.npy", + nvnmd_config: str | None = "nvnmd/config.npy", + nvnmd_weight: str | None = "nvnmd/weight.npy", + nvnmd_map: str | None = "nvnmd/map.npy", **kwargs, ) -> None: # build mapping table diff --git a/deepmd/tf/nvnmd/entrypoints/train.py b/deepmd/tf/nvnmd/entrypoints/train.py index c690190c0d..1650061c94 100644 --- a/deepmd/tf/nvnmd/entrypoints/train.py +++ b/deepmd/tf/nvnmd/entrypoints/train.py @@ -1,9 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import logging import os -from typing import ( - Optional, -) from deepmd.tf.entrypoints.freeze import ( freeze, @@ -118,8 +115,8 @@ def normalized_input_qnn(jdata, PATH_QNN, CONFIG_CNN, WEIGHT_CNN, MAP_CNN): def train_nvnmd( *, INPUT: str, - init_model: Optional[str], - restart: Optional[str], + init_model: str | None, + restart: str | None, step: str, skip_neighbor_stat: bool = False, **kwargs, diff --git a/deepmd/tf/nvnmd/entrypoints/wrap.py b/deepmd/tf/nvnmd/entrypoints/wrap.py index ced97bdbf1..e1878aa254 100755 --- a/deepmd/tf/nvnmd/entrypoints/wrap.py +++ b/deepmd/tf/nvnmd/entrypoints/wrap.py @@ -1,8 +1,5 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import logging -from typing import ( - Optional, -) import numpy as np @@ -608,10 +605,10 @@ def wrap_lut(self): def wrap( *, - nvnmd_config: Optional[str] = "nvnmd/config.npy", - nvnmd_weight: Optional[str] = "nvnmd/weight.npy", - nvnmd_map: Optional[str] = "nvnmd/map.npy", - nvnmd_model: Optional[str] = "nvnmd/model.pb", + nvnmd_config: str | None = "nvnmd/config.npy", + nvnmd_weight: str | None = "nvnmd/weight.npy", + nvnmd_map: str | None = "nvnmd/map.npy", + nvnmd_model: str | None = "nvnmd/model.pb", **kwargs, ) -> None: wrapObj = Wrap(nvnmd_config, nvnmd_weight, nvnmd_map, nvnmd_model) diff --git a/deepmd/tf/train/run_options.py b/deepmd/tf/train/run_options.py index fb6e09e8a9..0b5c3b1b43 100644 --- a/deepmd/tf/train/run_options.py +++ b/deepmd/tf/train/run_options.py @@ -96,7 +96,7 @@ class RunOptions: device type - gpu or cpu """ - gpus: Optional[list[int]] + gpus: list[int] | None world_size: int my_rank: int nodename: str @@ -108,11 +108,11 @@ class RunOptions: def __init__( self, - init_model: Optional[str] = None, - init_frz_model: Optional[str] = None, - finetune: Optional[str] = None, - restart: Optional[str] = None, - log_path: Optional[str] = None, + init_model: str | None = None, + init_frz_model: str | None = None, + finetune: str | None = None, + restart: str | None = None, + log_path: str | None = None, log_level: int = 0, mpi_log: str = "master", ) -> None: @@ -151,9 +151,9 @@ def print_resource_summary(self) -> None: def _setup_logger( self, - log_path: Optional[Path], + log_path: Path | None, log_level: int, - mpi_log: Optional[str], + mpi_log: str | None, ) -> None: """Set up package loggers. diff --git a/deepmd/tf/utils/convert.py b/deepmd/tf/utils/convert.py index 461d870f80..5575770257 100644 --- a/deepmd/tf/utils/convert.py +++ b/deepmd/tf/utils/convert.py @@ -2,9 +2,6 @@ import logging import os import textwrap -from typing import ( - Optional, -) from google.protobuf import ( text_format, @@ -60,7 +57,7 @@ def detect_model_version(input_model: str): def convert_to_21( - input_model: str, output_model: str, version: Optional[str] = None + input_model: str, output_model: str, version: str | None = None ) -> None: """Convert DP graph to 2.1 graph. diff --git a/deepmd/tf/utils/learning_rate.py b/deepmd/tf/utils/learning_rate.py index fee73ca9a3..64427e185d 100644 --- a/deepmd/tf/utils/learning_rate.py +++ b/deepmd/tf/utils/learning_rate.py @@ -1,7 +1,4 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( - Optional, -) import numpy as np @@ -50,9 +47,7 @@ def __init__( self.cd["decay_rate"] = decay_rate self.start_lr_ = self.cd["start_lr"] - def build( - self, global_step: tf.Tensor, stop_step: Optional[int] = None - ) -> tf.Tensor: + def build(self, global_step: tf.Tensor, stop_step: int | None = None) -> tf.Tensor: """Build the learning rate. Parameters diff --git a/deepmd/tf/utils/neighbor_stat.py b/deepmd/tf/utils/neighbor_stat.py index 37028b23bc..798c40621a 100644 --- a/deepmd/tf/utils/neighbor_stat.py +++ b/deepmd/tf/utils/neighbor_stat.py @@ -3,9 +3,6 @@ from collections.abc import ( Iterator, ) -from typing import ( - Optional, -) import numpy as np @@ -252,7 +249,7 @@ def _execute( self, coord: np.ndarray, atype: np.ndarray, - box: Optional[np.ndarray], + box: np.ndarray | None, pbc: bool, ): """Execute the operation. diff --git a/deepmd/tf/utils/parallel_op.py b/deepmd/tf/utils/parallel_op.py index c23f347dc1..46f9fb87b7 100644 --- a/deepmd/tf/utils/parallel_op.py +++ b/deepmd/tf/utils/parallel_op.py @@ -1,11 +1,10 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from collections.abc import ( + Callable, Generator, ) from typing import ( Any, - Callable, - Optional, ) from deepmd.tf.env import ( @@ -46,8 +45,8 @@ class ParallelOp: def __init__( self, builder: Callable[..., tuple[dict[str, tf.Tensor], tuple[tf.Tensor]]], - nthreads: Optional[int] = None, - config: Optional[tf.ConfigProto] = None, + nthreads: int | None = None, + config: tf.ConfigProto | None = None, ) -> None: if nthreads is not None: self.nthreads = nthreads diff --git a/deepmd/tf/utils/spin.py b/deepmd/tf/utils/spin.py index 8919bbd16a..c36a1ee21b 100644 --- a/deepmd/tf/utils/spin.py +++ b/deepmd/tf/utils/spin.py @@ -1,7 +1,4 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( - Optional, -) from deepmd.tf.env import ( GLOBAL_TF_FLOAT_PRECISION, @@ -24,9 +21,9 @@ class Spin: def __init__( self, - use_spin: Optional[list[bool]] = None, - spin_norm: Optional[list[float]] = None, - virtual_len: Optional[list[float]] = None, + use_spin: list[bool] | None = None, + spin_norm: list[float] | None = None, + virtual_len: list[float] | None = None, ) -> None: """Constructor.""" self.use_spin = use_spin diff --git a/deepmd/tf/utils/tabulate.py b/deepmd/tf/utils/tabulate.py index 67ae2f96bd..4e6c027f1b 100644 --- a/deepmd/tf/utils/tabulate.py +++ b/deepmd/tf/utils/tabulate.py @@ -1,11 +1,11 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import logging +from collections.abc import ( + Callable, +) from functools import ( cached_property, ) -from typing import ( - Callable, -) import numpy as np from scipy.special import ( diff --git a/deepmd/tf/utils/type_embed.py b/deepmd/tf/utils/type_embed.py index 9b7b17528d..34f3ae225b 100644 --- a/deepmd/tf/utils/type_embed.py +++ b/deepmd/tf/utils/type_embed.py @@ -1,10 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import logging import re -from typing import ( - Optional, - Union, -) import numpy as np @@ -116,15 +112,15 @@ def __init__( ntypes: int, neuron: list[int], resnet_dt: bool = False, - activation_function: Union[str, None] = "tanh", + activation_function: str | None = "tanh", precision: str = "default", trainable: bool = True, - seed: Optional[int] = None, + seed: int | None = None, uniform_seed: bool = False, padding: bool = False, use_econf_tebd: bool = False, use_tebd_bias: bool = False, - type_map: Optional[list[str]] = None, + type_map: list[str] | None = None, **kwargs, ) -> None: """Constructor.""" diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index 5878ea473d..7fcc117ab5 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -2,11 +2,11 @@ import json import logging import warnings +from collections.abc import ( + Callable, +) from typing import ( Any, - Callable, - Optional, - Union, ) from dargs import ( @@ -89,7 +89,7 @@ def deprecate_argument_extra_check(key: str) -> Callable[[dict], bool]: The name of the deprecated argument. """ - def deprecate_something(data: Optional[dict]) -> bool: + def deprecate_something(data: dict | None) -> bool: if data is not None and key in data: warnings.warn(f"{key} has been removed and takes no effect.", FutureWarning) data.pop(key) @@ -183,10 +183,10 @@ def __init__(self) -> None: self.__plugin = Plugin() def register( - self, name: str, alias: Optional[list[str]] = None, doc: str = "" + self, name: str, alias: list[str] | None = None, doc: str = "" ) -> Callable[ - [Union[Callable[[], Argument], Callable[[], list[Argument]]]], - Union[Callable[[], Argument], Callable[[], list[Argument]]], + [Callable[[], Argument] | Callable[[], list[Argument]]], + Callable[[], Argument] | Callable[[], list[Argument]], ]: """Register a descriptor argument plugin. @@ -1820,7 +1820,7 @@ def fitting_ener() -> list[Argument]: Argument("seed", [int, None], optional=True, doc=doc_seed), Argument( "atom_ener", - list[Optional[float]], + list[float | None], optional=True, default=[], doc=doc_atom_ener, @@ -2284,7 +2284,7 @@ def model_args(exclude_hybrid: bool = False) -> list[Argument]: ), Argument( "preset_out_bias", - dict[str, list[Optional[Union[float, list[float]]]]], + dict[str, list[float | list[float] | None]], optional=True, default=None, doc=doc_only_pt_supported + doc_preset_out_bias, @@ -2544,9 +2544,7 @@ def learning_rate_args(fold_subdoc: bool = False) -> Argument: # --- Loss configurations: --- # -def start_pref( - item: str, label: Optional[str] = None, abbr: Optional[str] = None -) -> str: +def start_pref(item: str, label: str | None = None, abbr: str | None = None) -> str: if label is None: label = item if abbr is None: diff --git a/deepmd/utils/batch_size.py b/deepmd/utils/batch_size.py index 860da030ba..012ffb4260 100644 --- a/deepmd/utils/batch_size.py +++ b/deepmd/utils/batch_size.py @@ -5,9 +5,11 @@ ABC, abstractmethod, ) +from collections.abc import ( + Callable, +) from typing import ( Any, - Callable, ) import array_api_compat diff --git a/deepmd/utils/compat.py b/deepmd/utils/compat.py index f93c30e047..f15c2b4909 100644 --- a/deepmd/utils/compat.py +++ b/deepmd/utils/compat.py @@ -11,8 +11,6 @@ ) from typing import ( Any, - Optional, - Union, ) import numpy as np @@ -23,7 +21,7 @@ def convert_input_v0_v1( - jdata: dict[str, Any], warning: bool = True, dump: Optional[Union[str, Path]] = None + jdata: dict[str, Any], warning: bool = True, dump: str | Path | None = None ) -> dict[str, Any]: """Convert input from v0 format to v1. @@ -54,7 +52,7 @@ def convert_input_v0_v1( return output -def _warning_input_v0_v1(fname: Optional[Union[str, Path]]) -> None: +def _warning_input_v0_v1(fname: str | Path | None) -> None: msg = ( "It seems that you are using a deepmd-kit input of version 0.x.x, " "which is deprecated. we have converted the input to >2.0.0 compatible" @@ -288,7 +286,7 @@ def remove_decay_rate(jdata: dict[str, Any]) -> None: def convert_input_v1_v2( - jdata: dict[str, Any], warning: bool = True, dump: Optional[Union[str, Path]] = None + jdata: dict[str, Any], warning: bool = True, dump: str | Path | None = None ) -> dict[str, Any]: tr_cfg = jdata["training"] tr_data_keys = { @@ -324,7 +322,7 @@ def convert_input_v1_v2( return jdata -def _warning_input_v1_v2(fname: Optional[Union[str, Path]]) -> None: +def _warning_input_v1_v2(fname: str | Path | None) -> None: msg = ( "It seems that you are using a deepmd-kit input of version 1.x.x, " "which is deprecated. we have converted the input to >2.0.0 compatible" @@ -335,7 +333,7 @@ def _warning_input_v1_v2(fname: Optional[Union[str, Path]]) -> None: def deprecate_numb_test( - jdata: dict[str, Any], warning: bool = True, dump: Optional[Union[str, Path]] = None + jdata: dict[str, Any], warning: bool = True, dump: str | Path | None = None ) -> dict[str, Any]: """Deprecate `numb_test` since v2.1. It has taken no effect since v2.0. @@ -373,7 +371,7 @@ def deprecate_numb_test( def update_deepmd_input( - jdata: dict[str, Any], warning: bool = True, dump: Optional[Union[str, Path]] = None + jdata: dict[str, Any], warning: bool = True, dump: str | Path | None = None ) -> dict[str, Any]: def is_deepmd_v0_input(jdata: dict[str, Any]) -> bool: return "model" not in jdata.keys() diff --git a/deepmd/utils/data.py b/deepmd/utils/data.py index 26a27c82d7..287107a7ff 100644 --- a/deepmd/utils/data.py +++ b/deepmd/utils/data.py @@ -13,8 +13,6 @@ ) from typing import ( Any, - Optional, - Union, ) import numpy as np @@ -64,9 +62,9 @@ def __init__( sys_path: str, set_prefix: str = "set", shuffle_test: bool = True, - type_map: Optional[list[str]] = None, + type_map: list[str] | None = None, optional_type_map: bool = True, - modifier: Optional[Any] = None, + modifier: Any | None = None, trn_all_set: bool = False, sort_atoms: bool = True, ) -> None: @@ -148,10 +146,10 @@ def add( atomic: bool = False, must: bool = False, high_prec: bool = False, - type_sel: Optional[list[int]] = None, + type_sel: list[int] | None = None, repeat: int = 1, default: float = 0.0, - dtype: Optional[np.dtype] = None, + dtype: np.dtype | None = None, output_natoms_for_type_sel: bool = False, ) -> "DeepmdData": """Add a data item that to be loaded. @@ -518,7 +516,7 @@ def _get_memmap(self, path: DPPath) -> np.memmap: return self._create_memmap(str(abs_path), str(file_mtime)) def _get_subdata( - self, data: dict[str, Any], idx: Optional[np.ndarray] = None + self, data: dict[str, Any], idx: np.ndarray | None = None ) -> dict[str, Any]: new_data = {} for ii in data: @@ -580,7 +578,7 @@ def _shuffle_data(self, data: dict[str, Any]) -> dict[str, Any]: ret[kk] = data[kk] return ret, idx - def _get_nframes(self, set_name: Union[DPPath, str]) -> int: + def _get_nframes(self, set_name: DPPath | str) -> int: if not isinstance(set_name, DPPath): set_name = DPPath(set_name) path = set_name / "coord.npy" @@ -717,9 +715,9 @@ def _load_data( must: bool = True, repeat: int = 1, high_prec: bool = False, - type_sel: Optional[list[int]] = None, + type_sel: list[int] | None = None, default: float = 0.0, - dtype: Optional[np.dtype] = None, + dtype: np.dtype | None = None, output_natoms_for_type_sel: bool = False, ) -> np.ndarray: if atomic: @@ -977,7 +975,7 @@ def _make_idx_map(self, atom_type: np.ndarray) -> np.ndarray: idx_map = idx return idx_map - def _load_type_map(self, sys_path: DPPath) -> Optional[list[str]]: + def _load_type_map(self, sys_path: DPPath) -> list[str] | None: fname = sys_path / "type_map.raw" if fname.is_file(): return fname.load_txt(dtype=str, ndmin=1).tolist() @@ -1059,10 +1057,10 @@ def __init__( atomic: bool = False, must: bool = False, high_prec: bool = False, - type_sel: Optional[list[int]] = None, + type_sel: list[int] | None = None, repeat: int = 1, default: float = 0.0, - dtype: Optional[np.dtype] = None, + dtype: np.dtype | None = None, output_natoms_for_type_sel: bool = False, ) -> None: self.key = key diff --git a/deepmd/utils/data_system.py b/deepmd/utils/data_system.py index 4f22b3c380..82ef5ec25d 100644 --- a/deepmd/utils/data_system.py +++ b/deepmd/utils/data_system.py @@ -7,8 +7,6 @@ ) from typing import ( Any, - Optional, - Union, ) import numpy as np @@ -44,14 +42,14 @@ def __init__( systems: list[str], batch_size: int, test_size: int, - rcut: Optional[float] = None, + rcut: float | None = None, set_prefix: str = "set", shuffle_test: bool = True, - type_map: Optional[list[str]] = None, + type_map: list[str] | None = None, optional_type_map: bool = True, - modifier: Optional[Any] = None, + modifier: Any | None = None, trn_all_set: bool = False, - sys_probs: Optional[list[float]] = None, + sys_probs: list[float] | None = None, auto_prob_style: str = "prob_sys_size", sort_atoms: bool = True, ) -> None: @@ -287,7 +285,7 @@ def default_mesh(self) -> list[np.ndarray]: ] def compute_energy_shift( - self, rcond: Optional[float] = None, key: str = "energy" + self, rcond: float | None = None, key: str = "energy" ) -> tuple[np.ndarray, np.ndarray]: sys_ener = [] for ss in self.data_systems: @@ -348,10 +346,10 @@ def add( atomic: bool = False, must: bool = False, high_prec: bool = False, - type_sel: Optional[list[int]] = None, + type_sel: list[int] | None = None, repeat: int = 1, default: float = 0.0, - dtype: Optional[np.dtype] = None, + dtype: np.dtype | None = None, output_natoms_for_type_sel: bool = False, ) -> None: """Add a data item that to be loaded. @@ -414,7 +412,7 @@ def get_data_dict(self, ii: int = 0) -> dict: def set_sys_probs( self, - sys_probs: Optional[list[float]] = None, + sys_probs: list[float] | None = None, auto_prob_style: str = "prob_sys_size", ) -> None: if sys_probs is None: @@ -435,7 +433,7 @@ def set_sys_probs( probs = process_sys_probs(sys_probs, self.nbatches) self.sys_probs = probs - def get_batch(self, sys_idx: Optional[int] = None) -> dict: + def get_batch(self, sys_idx: int | None = None) -> dict: # batch generation style altered by Ziyao Li: # one should specify the "sys_prob" and "auto_prob_style" params # via set_sys_prob() function. The sys_probs this function uses is @@ -462,7 +460,7 @@ def get_batch(self, sys_idx: Optional[int] = None) -> dict: b_data = self.get_batch_mixed() return b_data - def get_batch_standard(self, sys_idx: Optional[int] = None) -> dict: + def get_batch_standard(self, sys_idx: int | None = None) -> dict: """Get a batch of data from the data systems in the standard way. Parameters @@ -562,7 +560,7 @@ def _merge_batch_data(self, batch_data: list[dict]) -> dict: # ! altered by Marián Rynik def get_test( - self, sys_idx: Optional[int] = None, n_test: int = -1 + self, sys_idx: int | None = None, n_test: int = -1 ) -> dict[str, np.ndarray]: # depreciated """Get test data from the the data systems. @@ -588,7 +586,7 @@ def get_test( test_system_data["default_mesh"] = self.default_mesh[idx] return test_system_data - def get_sys_ntest(self, sys_idx: Optional[int] = None) -> int: + def get_sys_ntest(self, sys_idx: int | None = None) -> int: """Get number of tests for the currently selected system, or one defined by sys_idx. """ @@ -654,7 +652,7 @@ def _make_auto_ts(self, percent: float) -> list[int]: return ts def _check_type_map_consistency( - self, type_map_list: list[Optional[list[str]]] + self, type_map_list: list[list[str] | None] ) -> list[str]: ret = [] for ii in type_map_list: @@ -785,7 +783,7 @@ def prob_sys_size_ext(keywords: str, nsystems: int, nbatch: int) -> list[float]: def process_systems( - systems: Union[str, list[str]], patterns: Optional[list[str]] = None + systems: str | list[str], patterns: list[str] | None = None ) -> list[str]: """Process the user-input systems. @@ -832,8 +830,8 @@ def process_systems( def get_data( jdata: dict[str, Any], rcut: float, - type_map: Optional[list[str]], - modifier: Optional[Any], + type_map: list[str] | None, + modifier: Any | None, multi_task_mode: bool = False, ) -> DeepmdDataSystem: """Get the data system. diff --git a/deepmd/utils/econf_embd.py b/deepmd/utils/econf_embd.py index ce7e6a1aaf..191e804506 100644 --- a/deepmd/utils/econf_embd.py +++ b/deepmd/utils/econf_embd.py @@ -1,8 +1,5 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( - Union, -) import numpy as np from mendeleev import ( @@ -273,7 +270,7 @@ def print_econf_embedding(res: dict[str, np.ndarray]) -> None: def sort_element_type(elements: list[str]) -> list[str]: """Sort element types based on their atomic number.""" - def get_atomic_number(symbol: str) -> Union[int, float]: + def get_atomic_number(symbol: str) -> int | float: try: return element(symbol).atomic_number except ValueError: diff --git a/deepmd/utils/env_mat_stat.py b/deepmd/utils/env_mat_stat.py index a1f2ba6966..b10842fb41 100644 --- a/deepmd/utils/env_mat_stat.py +++ b/deepmd/utils/env_mat_stat.py @@ -10,9 +10,6 @@ from collections.abc import ( Iterator, ) -from typing import ( - Optional, -) import numpy as np @@ -169,7 +166,7 @@ def load_stats(self, path: DPPath) -> None: ) def load_or_compute_stats( - self, data: list[dict[str, np.ndarray]], path: Optional[DPPath] = None + self, data: list[dict[str, np.ndarray]], path: DPPath | None = None ) -> None: """Load the statistics of the environment matrix if it exists, otherwise compute and save it. diff --git a/deepmd/utils/model_branch_dict.py b/deepmd/utils/model_branch_dict.py index 2b3390a85b..501f5287b6 100644 --- a/deepmd/utils/model_branch_dict.py +++ b/deepmd/utils/model_branch_dict.py @@ -5,7 +5,6 @@ ) from typing import ( Any, - Optional, ) @@ -100,7 +99,7 @@ def __init__( # Construct table header: fixed columns + dynamic info keys self.headers: list[str] = ["Model Branch", "Alias", *self.info_keys] - def _wrap_cell(self, text: Any, width: Optional[int] = None) -> list[str]: + def _wrap_cell(self, text: Any, width: int | None = None) -> list[str]: """ Convert a cell value into a list of wrapped text lines. diff --git a/deepmd/utils/out_stat.py b/deepmd/utils/out_stat.py index ecbd379e2d..6aabf4a4f4 100644 --- a/deepmd/utils/out_stat.py +++ b/deepmd/utils/out_stat.py @@ -1,10 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later """Output statistics.""" -from typing import ( - Optional, -) - import numpy as np from deepmd.env import ( @@ -15,8 +11,8 @@ def compute_stats_from_redu( output_redu: np.ndarray, natoms: np.ndarray, - assigned_bias: Optional[np.ndarray] = None, - rcond: Optional[float] = None, + assigned_bias: np.ndarray | None = None, + rcond: float | None = None, ) -> tuple[np.ndarray, np.ndarray]: """Compute the output statistics. @@ -135,7 +131,7 @@ def compute_stats_from_atomic( def compute_stats_do_not_distinguish_types( output_redu: np.ndarray, natoms: np.ndarray, - assigned_bias: Optional[np.ndarray] = None, + assigned_bias: np.ndarray | None = None, intensive: bool = False, ) -> tuple[np.ndarray, np.ndarray]: """Compute element-independent statistics for property fitting. diff --git a/deepmd/utils/pair_tab.py b/deepmd/utils/pair_tab.py index b261f62a40..503f721c98 100644 --- a/deepmd/utils/pair_tab.py +++ b/deepmd/utils/pair_tab.py @@ -4,7 +4,6 @@ import logging from typing import ( Any, - Optional, ) import numpy as np @@ -35,12 +34,12 @@ class PairTab: cutoff raduis for the tabulated potential """ - def __init__(self, filename: str, rcut: Optional[float] = None) -> None: + def __init__(self, filename: str, rcut: float | None = None) -> None: """Constructor.""" self.data_type = np.float64 self.reinit(filename, rcut) - def reinit(self, filename: str, rcut: Optional[float] = None) -> None: + def reinit(self, filename: str, rcut: float | None = None) -> None: """Initialize the tabulated interaction. Parameters diff --git a/deepmd/utils/path.py b/deepmd/utils/path.py index e6b00cdf80..5f156298e3 100644 --- a/deepmd/utils/path.py +++ b/deepmd/utils/path.py @@ -14,8 +14,6 @@ from typing import ( Any, ClassVar, - Optional, - Union, ) import h5py @@ -169,7 +167,7 @@ class DPOSPath(DPPath): mode, by default "r" """ - def __init__(self, path: Union[str, Path], mode: str = "r") -> None: + def __init__(self, path: str | Path, mode: str = "r") -> None: super().__init__() self.mode = mode self.path = Path(path) @@ -342,7 +340,7 @@ def load_numpy(self) -> np.ndarray: """ return self.root[self._name][:] - def load_txt(self, dtype: Optional[np.dtype] = None, **kwargs: Any) -> np.ndarray: + def load_txt(self, dtype: np.dtype | None = None, **kwargs: Any) -> np.ndarray: """Load NumPy array from text. Returns diff --git a/deepmd/utils/plugin.py b/deepmd/utils/plugin.py index 3c27750bbb..e1817327a3 100644 --- a/deepmd/utils/plugin.py +++ b/deepmd/utils/plugin.py @@ -6,10 +6,11 @@ from abc import ( ABCMeta, ) +from collections.abc import ( + Callable, +) from typing import ( Any, - Callable, - Optional, ) @@ -98,7 +99,7 @@ class PluginVariant(metaclass=VariantABCMeta): pass -def make_plugin_registry(name: Optional[str] = None) -> type[object]: +def make_plugin_registry(name: str | None = None) -> type[object]: """Make a plugin registry. Parameters diff --git a/deepmd/utils/random.py b/deepmd/utils/random.py index d6cc327034..8710693607 100644 --- a/deepmd/utils/random.py +++ b/deepmd/utils/random.py @@ -1,8 +1,4 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( - Optional, - Union, -) import numpy as np @@ -10,11 +6,11 @@ def choice( - a: Union[np.ndarray, int], - size: Optional[Union[int, tuple[int, ...]]] = None, + a: np.ndarray | int, + size: int | tuple[int, ...] | None = None, replace: bool = True, - p: Optional[np.ndarray] = None, -) -> Union[np.ndarray, int]: + p: np.ndarray | None = None, +) -> np.ndarray | int: """Generates a random sample from a given 1-D array. Parameters @@ -41,8 +37,8 @@ def choice( def random( - size: Optional[Union[int, tuple[int, ...]]] = None, -) -> Union[float, np.ndarray]: + size: int | tuple[int, ...] | None = None, +) -> float | np.ndarray: """Return random floats in the half-open interval [0.0, 1.0). Parameters @@ -58,7 +54,7 @@ def random( return _RANDOM_GENERATOR.random_sample(size) -def seed(val: Optional[Union[int, list[int]]] = None) -> None: +def seed(val: int | list[int] | None = None) -> None: """Seed the generator. Parameters diff --git a/deepmd/utils/spin.py b/deepmd/utils/spin.py index 040de6b95b..aed82cae8b 100644 --- a/deepmd/utils/spin.py +++ b/deepmd/utils/spin.py @@ -1,9 +1,5 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import copy -from typing import ( - Optional, - Union, -) import numpy as np @@ -39,7 +35,7 @@ class Spin: def __init__( self, use_spin: list[bool], - virtual_scale: Union[list[float], float], + virtual_scale: list[float] | float, ) -> None: type_dtype = np.int32 self.ntypes_real = len(use_spin) @@ -137,7 +133,7 @@ def init_atom_exclude_types_placeholder(self) -> None: self.atom_exclude_types_p = self.placeholder_type.tolist() def get_pair_exclude_types( - self, exclude_types: Optional[list[tuple[int, int]]] = None + self, exclude_types: list[tuple[int, int]] | None = None ) -> list[tuple[int, int]]: """ Return the pair-wise exclusion types for descriptor. @@ -155,7 +151,7 @@ def get_pair_exclude_types( return _exclude_types def get_atom_exclude_types( - self, exclude_types: Optional[list[int]] = None + self, exclude_types: list[int] | None = None ) -> list[int]: """ Return the atom-wise exclusion types for fitting before out_def. @@ -170,7 +166,7 @@ def get_atom_exclude_types( return _exclude_types def get_atom_exclude_types_placeholder( - self, exclude_types: Optional[list[int]] = None + self, exclude_types: list[int] | None = None ) -> list[int]: """ Return the atom-wise exclusion types for fitting after out_def. diff --git a/deepmd/utils/tabulate.py b/deepmd/utils/tabulate.py index 4daae5d471..fb40f798e2 100644 --- a/deepmd/utils/tabulate.py +++ b/deepmd/utils/tabulate.py @@ -9,7 +9,6 @@ ) from typing import ( Any, - Optional, ) import numpy as np @@ -446,7 +445,7 @@ def _get_table_size(self) -> int: raise RuntimeError("Unsupported descriptor") return table_size - def _get_data_type(self) -> Optional[type]: + def _get_data_type(self) -> type | None: for item in self.matrix["layer_" + str(self.layer_size)]: if len(item) != 0: return type(item[0][0]) diff --git a/deepmd/utils/update_sel.py b/deepmd/utils/update_sel.py index 1d5d9bef01..616b69c2ca 100644 --- a/deepmd/utils/update_sel.py +++ b/deepmd/utils/update_sel.py @@ -6,8 +6,6 @@ ) from typing import ( Any, - Optional, - Union, ) from deepmd.utils.data_system import ( @@ -26,9 +24,9 @@ class BaseUpdateSel(ABC): def update_one_sel( self, train_data: DeepmdDataSystem, - type_map: Optional[list[str]], + type_map: list[str] | None, rcut: float, - sel: Union[int, list[int], str], + sel: int | list[int] | str, mixed_type: bool = False, ) -> tuple[float, list[int]]: min_nbor_dist, tmp_sel = self.get_nbor_stat( @@ -84,10 +82,10 @@ def wrap_up_4(self, xx: int) -> int: def get_nbor_stat( self, train_data: DeepmdDataSystem, - type_map: Optional[list[str]], + type_map: list[str] | None, rcut: float, mixed_type: bool = False, - ) -> tuple[float, Union[int, list[int]]]: + ) -> tuple[float, int | list[int]]: """Get the neighbor statistics of the data. Parameters diff --git a/doc/backend.md b/doc/backend.md index cab8dc2128..b2f7dc4826 100644 --- a/doc/backend.md +++ b/doc/backend.md @@ -12,7 +12,7 @@ In the documentation, TensorFlow {{ tensorflow_icon }}, PyTorch {{ pytorch_icon - Model filename extension: `.pb` - Checkpoint filename extension: `.meta`, `.index`, `.data-00000-of-00001` -[TensorFlow](https://tensorflow.org) 2.7 or above is required, since NumPy 1.21 or above is required. +[TensorFlow](https://tensorflow.org) 2.8 is the first version to support Python 3.10. DeePMD-kit does not use the TensorFlow v2 API but uses the TensorFlow v1 API (`tf.compat.v1`) in the graph mode. ### PyTorch {{ pytorch_icon }} @@ -28,7 +28,7 @@ While `.pth` and `.pt` are the same in the PyTorch package, they have different - Model filename extension: `.xlo`, `.savedmodel` - Checkpoint filename extension: `.jax` -[JAX](https://jax.readthedocs.io/) 0.4.33 (which requires Python 3.10 or above) or above is required. +[JAX](https://jax.readthedocs.io/) 0.4.33 or above is required. Both `.xlo` and `.jax` are customized format extensions defined in DeePMD-kit, since JAX has no convention for file extensions. `.savedmodel` is the TensorFlow [SavedModel format](https://www.tensorflow.org/guide/saved_model) generated by [JAX2TF](https://www.tensorflow.org/guide/jax2tf), which needs the installation of TensorFlow. Only the `.savedmodel` format supports C++ inference, which needs the TensorFlow C++ interface. diff --git a/doc/development/coding-conventions.rst b/doc/development/coding-conventions.rst index 4f82b34a60..80d8f915d1 100644 --- a/doc/development/coding-conventions.rst +++ b/doc/development/coding-conventions.rst @@ -30,7 +30,7 @@ Rules ----- The code must be compatible with the oldest supported version of python -which is 3.9. +which is 3.10. The project follows the generic coding conventions as specified in the `Style Guide for Python Code`_, `Docstring diff --git a/doc/environment.yml b/doc/environment.yml index 85d5a97c5b..fc96b56485 100644 --- a/doc/environment.yml +++ b/doc/environment.yml @@ -4,7 +4,7 @@ channels: - defaults dependencies: - doxygen>=1.9.1 - - python=3.9 + - python=3.10 - pip>=20.1 - pip: - ..[docs,cpu,torch] diff --git a/doc/install/easy-install.md b/doc/install/easy-install.md index 9f4d6d8861..0a28f5786c 100644 --- a/doc/install/easy-install.md +++ b/doc/install/easy-install.md @@ -10,7 +10,7 @@ You can refer to [DeepModeling conda FAQ](https://docs.deepmodeling.com/faq/cond ::: :::{note} -Python 3.9 or above is required for Python interface. +Python 3.10 or above is required for Python interface. ::: - [Install off-line packages](#install-off-line-packages) diff --git a/doc/install/install-from-source.md b/doc/install/install-from-source.md index 2acad95b69..3ed72ba033 100644 --- a/doc/install/install-from-source.md +++ b/doc/install/install-from-source.md @@ -21,7 +21,7 @@ deepmd_source_dir=`pwd` ### Install Backend's Python interface First, check the Python version on your machine. -Python 3.9 or above is required. +Python 3.10 or above is required. ```bash python --version @@ -136,7 +136,7 @@ deactivate If one has multiple python interpreters named something like python3.x, it can be specified by, for example ```bash -virtualenv -p python3.9 $deepmd_venv +virtualenv -p python3.12 $deepmd_venv ``` One should remember to activate the virtual environment every time he/she uses DeePMD-kit. diff --git a/doc/install/install-tf.1.12.md b/doc/install/install-tf.1.12.md deleted file mode 100644 index 13abd8f7a7..0000000000 --- a/doc/install/install-tf.1.12.md +++ /dev/null @@ -1,131 +0,0 @@ -# Install TensorFlow's C++ interface - -The TensorFlow's C++ interface will be compiled from the source code. Firstly one installs bazel. It is highly recommended that the bazel version 0.15.0 is used. A full instruction of bazel installation can be found [here](https://docs.bazel.build/versions/master/install.html). - -```bash -cd /some/workspace -wget https://github.com/bazelbuild/bazel/releases/download/0.15.0/bazel-0.15.0-dist.zip -mkdir bazel-0.15.0 -cd bazel-0.15.0 -unzip ../bazel-0.15.0-dist.zip -./compile.sh -export PATH=`pwd`/output:$PATH -``` - -Firstly get the source code of the TensorFlow - -```bash -cd /some/workspace -git clone https://github.com/tensorflow/tensorflow tensorflow -b v1.12.0 --depth=1 -cd tensorflow -``` - -DeePMD-kit is compiled by CMake, so we need to compile and integrate TensorFlow with CMake projects. The rest of this section follows [the instruction provided by Tuatini](http://tuatini.me/building-tensorflow-as-a-standalone-project/). Now execute - -```bash -./configure -``` - -You will answer a list of questions that help configure the building of TensorFlow. It is recommended to build for Python3. You may want to answer the question like this (please replace `$tensorflow_venv` with the virtual environment directory): - -```bash -Please specify the location of python. [Default is $tensorflow_venv/bin/python]: -``` - -The library path for Python should be set accordingly. - -Now build the shared library of TensorFlow: - -```bash -bazel build -c opt --verbose_failures //tensorflow:libtensorflow_cc.so -``` - -You may want to add options `--copt=-msse4.2`, `--copt=-mavx`, `--copt=-mavx2` and `--copt=-mfma` to enable SSE4.2, AVX, AVX2 and FMA SIMD accelerations, respectively. It is noted that these options should be chosen according to the CPU architecture. If the RAM becomes an issue for your machine, you may limit the RAM usage by using `--local_resources 2048,.5,1.0`. - -Now I assume you want to install TensorFlow in directory `$tensorflow_root`. Create the directory if it does not exist - -```bash -mkdir -p $tensorflow_root -``` - -Before moving on, we need to compile the dependencies of TensorFlow, including Protobuf, Eigen, nsync and absl. Firstly, protobuf - -```bash -mkdir /tmp/proto -sed -i 's;PROTOBUF_URL=.*;PROTOBUF_URL=\"https://mirror.bazel.build/github.com/google/protobuf/archive/v3.6.0.tar.gz\";g' tensorflow/contrib/makefile/download_dependencies.sh -tensorflow/contrib/makefile/download_dependencies.sh -cd tensorflow/contrib/makefile/downloads/protobuf/ -./autogen.sh -./configure --prefix=/tmp/proto/ -make -make install -``` - -Then Eigen - -```bash -mkdir /tmp/eigen -cd ../eigen -mkdir build_dir -cd build_dir -cmake -DCMAKE_INSTALL_PREFIX=/tmp/eigen/ ../ -make install -``` - -nsync - -```bash -mkdir /tmp/nsync -cd ../../nsync -mkdir build_dir -cd build_dir -cmake -DCMAKE_INSTALL_PREFIX=/tmp/nsync/ ../ -make -make install -``` - -And absl - -```bash -cd ../../absl -bazel build -mkdir -p $tensorflow_root/include/ -rsync -avzh --include '*/' --include '*.h' --exclude '*' absl $tensorflow_root/include/ -cd ../../../../.. -``` - -Now, copy the libraries to the tensorflow's installation directory: - -```bash -mkdir $tensorflow_root/lib -cp bazel-bin/tensorflow/libtensorflow_cc.so $tensorflow_root/lib/ -cp bazel-bin/tensorflow/libtensorflow_framework.so $tensorflow_root/lib/ -cp /tmp/proto/lib/libprotobuf.a $tensorflow_root/lib/ -cp /tmp/nsync/lib64/libnsync.a $tensorflow_root/lib/ -``` - -Then copy the headers - -```bash -mkdir -p $tensorflow_root/include/tensorflow -cp -r bazel-genfiles/* $tensorflow_root/include/ -cp -r tensorflow/cc $tensorflow_root/include/tensorflow -cp -r tensorflow/core $tensorflow_root/include/tensorflow -cp -r third_party $tensorflow_root/include -cp -r /tmp/proto/include/* $tensorflow_root/include -cp -r /tmp/eigen/include/eigen3/* $tensorflow_root/include -cp -r /tmp/nsync/include/*h $tensorflow_root/include -``` - -Now clean up the source files in the header directories: - -```bash -cd $tensorflow_root/include -find . -name "*.cc" -type f -delete -``` - -The temporary installation directories for the dependencies can be removed: - -```bash -rm -fr /tmp/proto /tmp/eigen /tmp/nsync -``` diff --git a/doc/install/install-tf.1.14-gpu.md b/doc/install/install-tf.1.14-gpu.md deleted file mode 100644 index 5850af24ba..0000000000 --- a/doc/install/install-tf.1.14-gpu.md +++ /dev/null @@ -1,164 +0,0 @@ -# Install TensorFlow-GPU's C++ interface - -TensorFlow's C++ interface will be compiled from the source code. Firstly one installs Bazel. It is highly recommended that the Bazel version 0.24.1 is used. Full instructions on Bazel installation can be found [here](https://docs.bazel.build/versions/master/install.html). - -```bash -cd /some/workspace -wget https://github.com/bazelbuild/bazel/releases/download/0.24.1/bazel-0.24.1-dist.zip -mkdir bazel-0.24.1 -cd bazel-0.24.1 -unzip ../bazel-0.24.1-dist.zip -./compile.sh -export PATH=`pwd`/output:$PATH -``` - -Firstly get the source code of the TensorFlow - -```bash -cd /some/workspace -git clone https://github.com/tensorflow/tensorflow tensorflow -b v1.14.0 --depth=1 -cd tensorflow -``` - -DeePMD-kit is compiled by CMake, so we need to compile and integrate TensorFlow with CMake projects. The rest of this section follows [the instruction provided by Tuatini](http://tuatini.me/building-tensorflow-as-a-standalone-project/). Now execute - -You will answer a list of questions that help configure the building of TensorFlow. It is recommended to build for Python3. You may want to answer the question like this (please replace `$tensorflow_venv` with the virtual environment directory): - -```bash -./configure -Please specify the location of python. [Default is xxx]: - -Traceback (most recent call last): - File "", line 1, in -AttributeError: module 'site' has no attribute 'getsitepackages' -Found possible Python library paths: - /xxx/deepmd_gpu/tensorflow_venv/lib/python3.7/site-packages -Please input the desired Python library path to use. Default is [xxx] - -Do you wish to build TensorFlow with XLA JIT support? [Y/n]: -XLA JIT support will be enabled for TensorFlow. - -Do you wish to build TensorFlow with OpenCL SYCL support? [y/N]: -No OpenCL SYCL support will be enabled for TensorFlow. - -Do you wish to build TensorFlow with ROCm support? [y/N]: -No ROCm support will be enabled for TensorFlow. - -Do you wish to build TensorFlow with CUDA support? [y/N]: y -CUDA support will be enabled for TensorFlow. - -Do you wish to build TensorFlow with TensorRT support? [y/N]: -No TensorRT support will be enabled for TensorFlow. - -Found CUDA 10.1 in: - /usr/local/cuda/lib64 - /usr/local/cuda/include -Found cuDNN 7 in: - /usr/local/cuda/lib64 - /usr/local/cuda/include - -Please specify a list of comma-separated CUDA compute capabilities you want to build with. -You can find the compute capability of your device at: https://developer.nvidia.com/cuda-gpus. -Please note that each additional compute capability significantly increases your build time and binary size, and that TensorFlow only supports compute capabilities >= 3.5 [Default is: 6.1,6.1]: - -Do you want to use clang as CUDA compiler? [y/N]: -nvcc will be used as CUDA compiler. - -Please specify which gcc should be used by nvcc as the host compiler. [Default is /usr/bin/gcc]: - - -Do you wish to build TensorFlow with MPI support? [y/N]: -No MPI support will be enabled for TensorFlow. - -Please specify optimization flags to use during compilation when bazel option "--config=opt" is specified [Default is -march=native -Wno-sign-compare]: - -Would you like to interactively configure ./WORKSPACE for Android builds? [y/N]: -Not configuring the WORKSPACE for Android builds. - -Preconfigured Bazel build configs. You can use any of the below by adding "--config=<>" to your build command. See .bazelrc for more details. - --config=mkl # Build with MKL support. - --config=monolithic # Config for mostly static monolithic build. - --config=gdr # Build with GDR support. - --config=verbs # Build with libverbs support. - --config=ngraph # Build with Intel nGraph support. - --config=numa # Build with NUMA support. - --config=dynamic_kernels # (Experimental) Build kernels into separate shared objects. - --config=v2 # Build TensorFlow 2.x instead of 1.x. -Preconfigured Bazel build configs to DISABLE default on features: - --config=noaws # Disable AWS S3 filesystem support. - --config=nogcp # Disable GCP support. - --config=nohdfs # Disable HDFS support. - --config=noignite # Disable Apache Ignite support. - --config=nokafka # Disable Apache Kafka support. - --config=nonccl # Disable NVIDIA NCCL support. -Configuration finished -``` - -The library path for Python should be set accordingly. - -Now build the shared library of TensorFlow: - -```bash -bazel build -c opt --verbose_failures //tensorflow:libtensorflow_cc.so -``` - -You may want to add options `--copt=-msse4.2`, `--copt=-mavx`, `--copt=-mavx2` and `--copt=-mfma` to enable SSE4.2, AVX, AVX2 and FMA SIMD accelerations, respectively. It is noted that these options should be chosen according to the CPU architecture. If the RAM becomes an issue for your machine, you may limit the RAM usage by using `--local_resources 2048,.5,1.0`. - -Now I assume you want to install TensorFlow in directory `$tensorflow_root`. Create the directory if it does not exist - -```bash -mkdir -p $tensorflow_root -``` - -Now, copy the libraries to the TensorFlow's installation directory: - -```bash -mkdir $tensorflow_root/lib -cp -d bazel-bin/tensorflow/libtensorflow_cc.so* $tensorflow_root/lib/ -cp -d bazel-bin/tensorflow/libtensorflow_framework.so* $tensorflow_root/lib/ -cp -d $tensorflow_root/lib/libtensorflow_framework.so.1 $tensorflow_root/lib/libtensorflow_framework.so -``` - -Then copy the headers - -```bash -mkdir -p $tensorflow_root/include/tensorflow -cp -r bazel-genfiles/* $tensorflow_root/include/ -cp -r tensorflow/cc $tensorflow_root/include/tensorflow -cp -r tensorflow/core $tensorflow_root/include/tensorflow -cp -r third_party $tensorflow_root/include -cp -r bazel-tensorflow/external/eigen_archive/Eigen/ $tensorflow_root/include -cp -r bazel-tensorflow/external/eigen_archive/unsupported/ $tensorflow_root/include -rsync -avzh --include '*/' --include '*.h' --include '*.inc' --exclude '*' bazel-tensorflow/external/protobuf_archive/src/ $tensorflow_root/include/ -rsync -avzh --include '*/' --include '*.h' --include '*.inc' --exclude '*' bazel-tensorflow/external/com_google_absl/absl/ $tensorflow_root/include/absl -``` - -Now clean up the source files in the header directories: - -```bash -cd $tensorflow_root/include -find . -name "*.cc" -type f -delete -``` - -# Troubleshooting - -```bash -git: unknown command -C ... -``` - -This may be your git version issue because the low version of Git does not support this command. Upgrading your Git may be helpful. - -```bash -CMake Error: The following variables are used in this project, but they are set to NOTFOUND. -Please set them or make sure they are set and tested correctly in the CMake files: -FFTW_LIB (ADVANCED) - linked by target "FFTW" in directory xxx -``` - -Currently, when building the Eigen package, you can delete the FFTW in the CMake file. - -```bash -fatal error: absl/numeric/int128_have_intrinsic.inc: No such file or directory -``` - -Basically, you could build an empty file named "int128_have_intrinsic.inc" in the same directory of "int128.h". diff --git a/doc/install/install-tf.1.14.md b/doc/install/install-tf.1.14.md deleted file mode 100644 index 6457d484ad..0000000000 --- a/doc/install/install-tf.1.14.md +++ /dev/null @@ -1,79 +0,0 @@ -# Install tensorflow's C++ interface - -The tensorflow's C++ interface will be compiled from the source code. Firstly one installs bazel. It is highly recommended that the bazel version 0.24.1 is used. A full instruction of bazel installation can be found [here](https://docs.bazel.build/versions/master/install.html). - -```bash -cd /some/workspace -wget https://github.com/bazelbuild/bazel/releases/download/0.24.1/bazel-0.24.1-dist.zip -mkdir bazel-0.24.1 -cd bazel-0.24.1 -unzip ../bazel-0.24.1-dist.zip -./compile.sh -export PATH=`pwd`/output:$PATH -``` - -Firstly get the source code of the tensorflow - -```bash -cd /some/workspace -git clone https://github.com/tensorflow/tensorflow tensorflow -b v1.14.0 --depth=1 -cd tensorflow -``` - -DeePMD-kit is compiled by cmake, so we need to compile and integrate tensorflow with cmake projects. The rest of this section basically follows [the instruction provided by Tuatini](http://tuatini.me/building-tensorflow-as-a-standalone-project/). Now execute - -```bash -./configure -``` - -You will answer a list of questions that help configure the building of tensorflow. It is recommended to build for Python3. You may want to answer the question like this (please replace `$tensorflow_venv` by the virtual environment directory): - -```bash -Please specify the location of python. [Default is $tensorflow_venv/bin/python]: -``` - -The library path for Python should be set accordingly. - -Now build the shared library of tensorflow: - -```bash -bazel build -c opt --verbose_failures //tensorflow:libtensorflow_cc.so -``` - -You may want to add options `--copt=-msse4.2`, `--copt=-mavx`, `--copt=-mavx2` and `--copt=-mfma` to enable SSE4.2, AVX, AVX2 and FMA SIMD accelerations, respectively. It is noted that these options should be chosen according to the CPU architecture. If the RAM becomes an issue of your machine, you may limit the RAM usage by using `--local_resources 2048,.5,1.0`. - -Now I assume you want to install tensorflow in directory `$tensorflow_root`. Create the directory if it does not exists - -```bash -mkdir -p $tensorflow_root -``` - -Now, copy the libraries to the tensorflow's installation directory: - -```bash -mkdir $tensorflow_root/lib -cp -d bazel-bin/tensorflow/libtensorflow_cc.so* $tensorflow_root/lib/ -cp -d bazel-bin/tensorflow/libtensorflow_framework.so* $tensorflow_root/lib/ -cp -d $tensorflow_root/lib/libtensorflow_framework.so.1 $tensorflow_root/lib/libtensorflow_framework.so -``` - -Then copy the headers - -```bash -mkdir -p $tensorflow_root/include/tensorflow -cp -r bazel-genfiles/* $tensorflow_root/include/ -cp -r tensorflow/cc $tensorflow_root/include/tensorflow -cp -r tensorflow/core $tensorflow_root/include/tensorflow -cp -r third_party $tensorflow_root/include -cp -r bazel-tensorflow/external/eigen_archive/Eigen/ $tensorflow_root/include -cp -r bazel-tensorflow/external/eigen_archive/unsupported/ $tensorflow_root/include -rsync -avzh --include '*/' --include '*.h' --include '*.inc' --exclude '*' bazel-tensorflow/external/protobuf_archive/src/ $tensorflow_root/include/ -rsync -avzh --include '*/' --include '*.h' --include '*.inc' --exclude '*' bazel-tensorflow/external/com_google_absl/absl/ $tensorflow_root/include/absl -``` - -Now clean up the source files in the header directories: - -```bash -cd $tensorflow_root/include -find . -name "*.cc" -type f -delete -``` diff --git a/doc/install/install-tf.1.8.md b/doc/install/install-tf.1.8.md deleted file mode 100644 index f9554f9348..0000000000 --- a/doc/install/install-tf.1.8.md +++ /dev/null @@ -1,121 +0,0 @@ -# Install tensorflow's C++ interface - -The tensorflow's C++ interface will be compiled from the source code. Firstly one installs bazel. It is highly recommended that the bazel version 0.10.0 is used. A full instruction of bazel installation can be found [here](https://docs.bazel.build/versions/master/install.html). - -```bash -cd /some/workspace -wget https://github.com/bazelbuild/bazel/releases/download/0.10.0/bazel-0.10.0-dist.zip -mkdir bazel-0.10.0 -cd bazel-0.10.0 -unzip ../bazel-0.10.0-dist.zip -./compile.sh -export PATH=`pwd`/output:$PATH -``` - -Firstly get the source code of the TensorFlow - -```bash -cd /some/workspace -git clone https://github.com/tensorflow/tensorflow tensorflow -b v1.8.0 --depth=1 -cd tensorflow -``` - -DeePMD-kit is compiled by CMake, so we need to compile and integrate TensorFlow with CMake projects. The rest of this section basically follows [the instruction provided by Tuatini](http://tuatini.me/building-tensorflow-as-a-standalone-project/). Now execute - -```bash -./configure -``` - -You will answer a list of questions that help configure the building of TensorFlow. It is recommended to build for Python3. You may want to answer the question like this (please replace `$tensorflow_venv` with the virtual environment directory): - -```bash -Please specify the location of python. [Default is $tensorflow_venv/bin/python]: -``` - -The library path for Python should be set accordingly. - -Now build the shared library of TensorFlow: - -```bash -bazel build -c opt --verbose_failures //tensorflow:libtensorflow_cc.so -``` - -You may want to add options `--copt=-msse4.2`, `--copt=-mavx`, `--copt=-mavx2` and `--copt=-mfma` to enable SSE4.2, AVX, AVX2 and FMA SIMD accelerations, respectively. It is noted that these options should be chosen according to the CPU architecture. If the RAM becomes an issue of your machine, you may limit the RAM usage by using `--local_resources 2048,.5,1.0`. - -Now I assume you want to install TensorFlow in directory `$tensorflow_root`. Create the directory if it does not exist - -```bash -mkdir -p $tensorflow_root -``` - -Before moving on, we need to compile the dependencies of TensorFlow, including Protobuf, Eigen and nsync. Firstly, protobuf - -```bash -mkdir /tmp/proto -tensorflow/contrib/makefile/download_dependencies.sh -cd tensorflow/contrib/makefile/downloads/protobuf/ -./autogen.sh -./configure --prefix=/tmp/proto/ -make -make install -``` - -Then Eigen - -```bash -mkdir /tmp/eigen -cd ../eigen -mkdir build_dir -cd build_dir -cmake -DCMAKE_INSTALL_PREFIX=/tmp/eigen/ ../ -make install -``` - -And nsync - -```bash -mkdir /tmp/nsync -cd ../../nsync -mkdir build_dir -cd build_dir -cmake -DCMAKE_INSTALL_PREFIX=/tmp/nsync/ ../ -make -make install -cd ../../../../../.. -``` - -Now, copy the libraries to the TensorFlow's installation directory: - -```bash -mkdir $tensorflow_root/lib -cp bazel-bin/tensorflow/libtensorflow_cc.so $tensorflow_root/lib/ -cp bazel-bin/tensorflow/libtensorflow_framework.so $tensorflow_root/lib/ -cp /tmp/proto/lib/libprotobuf.a $tensorflow_root/lib/ -cp /tmp/nsync/lib/libnsync.a $tensorflow_root/lib/ -``` - -Then copy the headers - -```bash -mkdir -p $tensorflow_root/include/tensorflow -cp -r bazel-genfiles/* $tensorflow_root/include/ -cp -r tensorflow/cc $tensorflow_root/include/tensorflow -cp -r tensorflow/core $tensorflow_root/include/tensorflow -cp -r third_party $tensorflow_root/include -cp -r /tmp/proto/include/* $tensorflow_root/include -cp -r /tmp/eigen/include/eigen3/* $tensorflow_root/include -cp -r /tmp/nsync/include/*h $tensorflow_root/include -``` - -Now clean up the source files in the header directories: - -```bash -cd $tensorflow_root/include -find . -name "*.cc" -type f -delete -``` - -The temporary installation directories for the dependencies can be removed: - -```bash -rm -fr /tmp/proto /tmp/eigen /tmp/nsync -``` diff --git a/doc/install/install-tf.2.3.md b/doc/install/install-tf.2.3.md deleted file mode 100644 index 2fc7b35f2c..0000000000 --- a/doc/install/install-tf.2.3.md +++ /dev/null @@ -1,123 +0,0 @@ -# Install TensorFlow's C++ interface - -The tensorflow's C++ interface will be compiled from the source code. Firstly one installs bazel. The bazel version 3.1.0 should be used. A full instruction of bazel installation can be found [here](https://docs.bazel.build/versions/master/install.html). - -```bash -cd /some/workspace -wget https://github.com/bazelbuild/bazel/releases/download/3.1.0/bazel-3.1.0-installer-linux-x86_64.sh -chmod +x bazel-3.1.0-installer-linux-x86_64.sh -./bazel-3.1.0-installer-linux-x86_64.sh --prefix /some/workspace/bazel -export PATH=/some/workspace/bazel/bin:$PATH -``` - -Firstly get the source code of the TensorFlow - -```bash -git clone https://github.com/tensorflow/tensorflow tensorflow -b v2.3.0 --depth=1 -cd tensorflow -./configure -``` - -You will answer a list of questions that help configure the building of TensorFlow. You may want to answer the question like the following. If you do not want to add CUDA support, please answer no. - -``` -Please specify the location of python. [Default is xxx]: - -Found possible Python library paths: - xxx -Please input the desired Python library path to use. Default is [xxx] - -Do you wish to build TensorFlow with OpenCL SYCL support? [y/N]: -No OpenCL SYCL support will be enabled for TensorFlow. - -Do you wish to build TensorFlow with ROCm support? [y/N]: -No ROCm support will be enabled for TensorFlow. - -Do you wish to build TensorFlow with CUDA support? [y/N]: y -CUDA support will be enabled for TensorFlow. - -Do you wish to build TensorFlow with TensorRT support? [y/N]: -No TensorRT support will be enabled for TensorFlow. - -Found CUDA 10.2 in: - /usr/local/cuda/lib64 - /usr/local/cuda/include -Found cuDNN 7 in: - /usr/local/cuda/lib64 - /usr/local/cuda/include - -Please specify a list of comma-separated CUDA compute capabilities you want to build with. -You can find the compute capability of your device at: https://developer.nvidia.com/cuda-gpus. -Please note that each additional compute capability significantly increases your build time and binary size, and that TensorFlow only supports compute capabilities >= 3.5 [Default is: 7.5,7.5]: - -Do you want to use clang as CUDA compiler? [y/N]: -nvcc will be used as CUDA compiler. - -Please specify which gcc should be used by nvcc as the host compiler. [Default is /usr/bin/gcc]: - -Please specify optimization flags to use during compilation when bazel option "--config=opt" is specified [Default is -march=native -Wno-sign-compare]: - -Would you like to interactively configure ./WORKSPACE for Android builds? [y/N]: -Not configuring the WORKSPACE for Android builds. - -Preconfigured Bazel build configs. You can use any of the below by adding "--config=<>" to your build command. See .bazelrc for more details. - --config=mkl # Build with MKL support. - --config=monolithic # Config for mostly static monolithic build. - --config=ngraph # Build with Intel nGraph support. - --config=numa # Build with NUMA support. - --config=dynamic_kernels # (Experimental) Build kernels into separate shared objects. - --config=v2 # Build TensorFlow 2.x instead of 1.x. -Preconfigured Bazel build configs to DISABLE default on features: - --config=noaws # Disable AWS S3 filesystem support. - --config=nogcp # Disable GCP support. - --config=nohdfs # Disable HDFS support. - --config=nonccl # Disable NVIDIA NCCL support. -Configuration finished -``` - -The library path for Python should be set accordingly. - -Now build the shared library of tensorflow: - -```bash -bazel build -c opt --verbose_failures //tensorflow:libtensorflow_cc.so -``` - -You may want to add options `--copt=-msse4.2`, `--copt=-mavx`, `--copt=-mavx2` and `--copt=-mfma` to enable SSE4.2, AVX, AVX2 and FMA SIMD accelerations, respectively. It is noted that these options should be chosen according to the CPU architecture. If the RAM becomes an issue of your machine, you may limit the RAM usage by using `--local_resources 2048,.5,1.0`. - -Now I assume you want to install TensorFlow in directory `$tensorflow_root`. Create the directory if it does not exist - -```bash -mkdir -p $tensorflow_root -``` - -Now, copy the libraries to the tensorflow's installation directory: - -```bash -mkdir -p $tensorflow_root/lib -cp -d bazel-bin/tensorflow/libtensorflow_cc.so* $tensorflow_root/lib/ -cp -d bazel-bin/tensorflow/libtensorflow_framework.so* $tensorflow_root/lib/ -cp -d $tensorflow_root/lib/libtensorflow_framework.so.2 $tensorflow_root/lib/libtensorflow_framework.so -``` - -Then copy the headers - -```bash -mkdir -p $tensorflow_root/include/tensorflow -rsync -avzh --exclude '_virtual_includes/' --include '*/' --include '*.h' --include '*.inc' --exclude '*' bazel-bin/ $tensorflow_root/include/ -rsync -avzh --include '*/' --include '*.h' --include '*.inc' --exclude '*' tensorflow/cc $tensorflow_root/include/tensorflow/ -rsync -avzh --include '*/' --include '*.h' --include '*.inc' --exclude '*' tensorflow/core $tensorflow_root/include/tensorflow/ -rsync -avzh --include '*/' --include '*' --exclude '*.cc' third_party/ $tensorflow_root/include/third_party/ -rsync -avzh --include '*/' --include '*' --exclude '*.txt' bazel-tensorflow/external/eigen_archive/Eigen/ $tensorflow_root/include/Eigen/ -rsync -avzh --include '*/' --include '*' --exclude '*.txt' bazel-tensorflow/external/eigen_archive/unsupported/ $tensorflow_root/include/unsupported/ -rsync -avzh --include '*/' --include '*.h' --include '*.inc' --exclude '*' bazel-tensorflow/external/com_google_protobuf/src/google/ $tensorflow_root/include/google/ -rsync -avzh --include '*/' --include '*.h' --include '*.inc' --exclude '*' bazel-tensorflow/external/com_google_absl/absl/ $tensorflow_root/include/absl/ -``` - -# Troubleshooting - -```bash -git: unknown command -C ... -``` - -This may be an issue with your git version issue. Early versions of git do not support this command, in this case upgrading your git to a newer version may resolve any issues. diff --git a/pyproject.toml b/pyproject.toml index ec9c5e5903..e61b63eb80 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -29,7 +29,6 @@ classifiers = [ "Programming Language :: Python :: 3 :: Only", "Environment :: GPU :: NVIDIA CUDA :: 12 :: 12.2", "Intended Audience :: Science/Research", - "Programming Language :: Python :: 3.9", "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", "Programming Language :: Python :: 3.12", @@ -55,7 +54,7 @@ dependencies = [ 'mendeleev', 'array-api-compat', ] -requires-python = ">=3.9" +requires-python = ">=3.10" keywords = ["deepmd"] [project.entry-points."lammps.plugins"] diff --git a/source/install/build_tf.py b/source/install/build_tf.py index d4fcc357e7..a5ca2d7411 100755 --- a/source/install/build_tf.py +++ b/source/install/build_tf.py @@ -57,7 +57,6 @@ ) from typing import ( NoReturn, - Optional, ) # default config @@ -133,9 +132,9 @@ def __init__( self, filename: str, url: str, - sha256: Optional[str] = None, + sha256: str | None = None, executable: bool = False, - gzip: Optional[str] = None, + gzip: str | None = None, ) -> None: self.filename = filename self.url = url @@ -856,7 +855,7 @@ class RawTextArgumentDefaultsHelpFormatter( pass -def parse_args(args: Optional[list[str]] = None): +def parse_args(args: list[str] | None = None): """TensorFlow C++ Library Installer commandline options argument parser. Parameters diff --git a/source/tests/array_api_strict/common.py b/source/tests/array_api_strict/common.py index 28f67a97f6..50109ded86 100644 --- a/source/tests/array_api_strict/common.py +++ b/source/tests/array_api_strict/common.py @@ -1,13 +1,10 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( - Optional, -) import array_api_strict import numpy as np -def to_array_api_strict_array(array: Optional[np.ndarray]): +def to_array_api_strict_array(array: np.ndarray | None): """Convert a numpy array to a JAX array. Parameters diff --git a/source/tests/common/test_argument_parser.py b/source/tests/common/test_argument_parser.py index 4aebb7dafc..df0d06e4b5 100644 --- a/source/tests/common/test_argument_parser.py +++ b/source/tests/common/test_argument_parser.py @@ -15,7 +15,6 @@ from typing import ( TYPE_CHECKING, Any, - Union, ) from deepmd.main import ( @@ -30,7 +29,7 @@ from typing_extensions import TypedDict # python<=3.7 class DATA(TypedDict): - type: Union[type, tuple[type]] + type: type | tuple[type] value: Any TEST_DICT = dict[str, DATA] diff --git a/source/tests/consistent/common.py b/source/tests/consistent/common.py index 9a9c585c44..c1e6644be6 100644 --- a/source/tests/consistent/common.py +++ b/source/tests/consistent/common.py @@ -8,6 +8,9 @@ ABC, abstractmethod, ) +from collections.abc import ( + Callable, +) from enum import ( Enum, ) @@ -16,10 +19,7 @@ ) from typing import ( Any, - Callable, ClassVar, - Optional, - Union, ) from uuid import ( uuid4, @@ -80,18 +80,18 @@ class CommonTest(ABC): """Arguments data.""" additional_data: ClassVar[dict] = {} """Additional data that will not be checked.""" - tf_class: ClassVar[Optional[type]] + tf_class: ClassVar[type | None] """TensorFlow model class.""" - dp_class: ClassVar[Optional[type]] + dp_class: ClassVar[type | None] """Native DP model class.""" - pt_class: ClassVar[Optional[type]] + pt_class: ClassVar[type | None] """PyTorch model class.""" - jax_class: ClassVar[Optional[type]] + jax_class: ClassVar[type | None] """JAX model class.""" - pd_class: ClassVar[Optional[type]] + pd_class: ClassVar[type | None] """Paddle model class.""" - array_api_strict_class: ClassVar[Optional[type]] - args: ClassVar[Optional[Union[Argument, list[Argument]]]] + array_api_strict_class: ClassVar[type | None] + args: ClassVar[Argument | list[Argument] | None] """Arguments that maps to the `data`.""" skip_dp: ClassVar[bool] = False """Whether to skip the native DP model.""" diff --git a/source/tests/consistent/descriptor/test_dpa1.py b/source/tests/consistent/descriptor/test_dpa1.py index d31cf289b9..a0a3d9c833 100644 --- a/source/tests/consistent/descriptor/test_dpa1.py +++ b/source/tests/consistent/descriptor/test_dpa1.py @@ -2,7 +2,6 @@ import unittest from typing import ( Any, - Optional, ) import numpy as np @@ -133,7 +132,7 @@ def data(self) -> dict: def is_meaningless_zero_attention_layer_tests( self, attn_layer: int, - temperature: Optional[float], + temperature: float | None, ) -> bool: return attn_layer == 0 and (temperature is not None) diff --git a/source/tests/consistent/descriptor/test_se_atten_v2.py b/source/tests/consistent/descriptor/test_se_atten_v2.py index 459fac037b..129343c177 100644 --- a/source/tests/consistent/descriptor/test_se_atten_v2.py +++ b/source/tests/consistent/descriptor/test_se_atten_v2.py @@ -2,7 +2,6 @@ import unittest from typing import ( Any, - Optional, ) import numpy as np @@ -132,7 +131,7 @@ def is_meaningless_zero_attention_layer_tests( attn_layer: int, attn_dotr: bool, normalize: bool, - temperature: Optional[float], + temperature: float | None, ) -> bool: return attn_layer == 0 and (attn_dotr or normalize or temperature is not None) diff --git a/source/tests/infer/case.py b/source/tests/infer/case.py index fc7dd30e9c..828974c6e6 100644 --- a/source/tests/infer/case.py +++ b/source/tests/infer/case.py @@ -26,9 +26,6 @@ from pathlib import ( Path, ) -from typing import ( - Optional, -) import numpy as np import yaml @@ -156,7 +153,7 @@ def __init__(self, filename: str) -> None: self.model_def_script = config.get("model_def_script") @lru_cache - def get_model(self, suffix: str, out_file: Optional[str] = None) -> str: + def get_model(self, suffix: str, out_file: str | None = None) -> str: """Get the model file with the specified suffix. Parameters diff --git a/source/tests/jax/test_dp_hessian_model.py b/source/tests/jax/test_dp_hessian_model.py index 798b893651..89c066e980 100644 --- a/source/tests/jax/test_dp_hessian_model.py +++ b/source/tests/jax/test_dp_hessian_model.py @@ -7,25 +7,23 @@ from deepmd.dpmodel.common import ( to_numpy_array, ) +from deepmd.jax.common import ( + to_jax_array, +) +from deepmd.jax.descriptor.se_e2_a import ( + DescrptSeA, +) +from deepmd.jax.env import ( + jnp, +) +from deepmd.jax.fitting.fitting import ( + EnergyFittingNet, +) +from deepmd.jax.model.ener_model import ( + EnergyModel, +) -if sys.version_info >= (3, 10): - from deepmd.jax.common import ( - to_jax_array, - ) - from deepmd.jax.descriptor.se_e2_a import ( - DescrptSeA, - ) - from deepmd.jax.env import ( - jnp, - ) - from deepmd.jax.fitting.fitting import ( - EnergyFittingNet, - ) - from deepmd.jax.model.ener_model import ( - EnergyModel, - ) - - dtype = jnp.float64 +dtype = jnp.float64 @unittest.skipIf( diff --git a/source/tests/jax/test_make_hessian_model.py b/source/tests/jax/test_make_hessian_model.py index 185660e2be..8666ff4ad4 100644 --- a/source/tests/jax/test_make_hessian_model.py +++ b/source/tests/jax/test_make_hessian_model.py @@ -10,30 +10,28 @@ from deepmd.dpmodel.output_def import ( OutputVariableCategory, ) +from deepmd.jax.common import ( + to_jax_array, +) +from deepmd.jax.descriptor.se_e2_a import ( + DescrptSeA, +) +from deepmd.jax.env import ( + jax, + jnp, +) +from deepmd.jax.fitting.fitting import ( + EnergyFittingNet, +) +from deepmd.jax.model import ( + EnergyModel, +) + +from ..seed import ( + GLOBAL_SEED, +) -if sys.version_info >= (3, 10): - from deepmd.jax.common import ( - to_jax_array, - ) - from deepmd.jax.descriptor.se_e2_a import ( - DescrptSeA, - ) - from deepmd.jax.env import ( - jax, - jnp, - ) - from deepmd.jax.fitting.fitting import ( - EnergyFittingNet, - ) - from deepmd.jax.model import ( - EnergyModel, - ) - - from ..seed import ( - GLOBAL_SEED, - ) - - dtype = jnp.float64 +dtype = jnp.float64 def finite_hessian(f, x, delta=1e-6): diff --git a/source/tests/jax/test_padding_atoms.py b/source/tests/jax/test_padding_atoms.py index 42e2ae527c..b63b464721 100644 --- a/source/tests/jax/test_padding_atoms.py +++ b/source/tests/jax/test_padding_atoms.py @@ -10,25 +10,23 @@ from deepmd.dpmodel.common import ( to_numpy_array, ) +from deepmd.jax.common import ( + to_jax_array, +) +from deepmd.jax.descriptor.se_e2_a import ( + DescrptSeA, +) +from deepmd.jax.env import ( + jnp, +) +from deepmd.jax.fitting.fitting import ( + PropertyFittingNet, +) +from deepmd.jax.model.property_model import ( + PropertyModel, +) -if sys.version_info >= (3, 10): - from deepmd.jax.common import ( - to_jax_array, - ) - from deepmd.jax.descriptor.se_e2_a import ( - DescrptSeA, - ) - from deepmd.jax.env import ( - jnp, - ) - from deepmd.jax.fitting.fitting import ( - PropertyFittingNet, - ) - from deepmd.jax.model.property_model import ( - PropertyModel, - ) - - dtype = jnp.float64 +dtype = jnp.float64 @unittest.skipIf( diff --git a/source/tests/pd/common.py b/source/tests/pd/common.py index ec36fd0eb9..f96777c9a3 100644 --- a/source/tests/pd/common.py +++ b/source/tests/pd/common.py @@ -1,9 +1,5 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import pathlib -from typing import ( - Optional, - Union, -) import numpy as np import paddle @@ -51,10 +47,10 @@ def run_dp(cmd: str) -> int: def eval_model( model, - coords: Union[np.ndarray, paddle.Tensor], - cells: Optional[Union[np.ndarray, paddle.Tensor]], - atom_types: Union[np.ndarray, paddle.Tensor, list[int]], - spins: Optional[Union[np.ndarray, paddle.Tensor]] = None, + coords: np.ndarray | paddle.Tensor, + cells: np.ndarray | paddle.Tensor | None, + atom_types: np.ndarray | paddle.Tensor | list[int], + spins: np.ndarray | paddle.Tensor | None = None, atomic: bool = False, infer_batch_size: int = 2, denoise: bool = False, diff --git a/source/tests/pd/model/test_atomic_model_atomic_stat.py b/source/tests/pd/model/test_atomic_model_atomic_stat.py index bfc86edc12..e80ef42c6e 100644 --- a/source/tests/pd/model/test_atomic_model_atomic_stat.py +++ b/source/tests/pd/model/test_atomic_model_atomic_stat.py @@ -6,7 +6,6 @@ ) from typing import ( NoReturn, - Optional, ) import h5py @@ -81,11 +80,11 @@ def forward( self, descriptor: paddle.Tensor, atype: paddle.Tensor, - gr: Optional[paddle.Tensor] = None, - g2: Optional[paddle.Tensor] = None, - h2: Optional[paddle.Tensor] = None, - fparam: Optional[paddle.Tensor] = None, - aparam: Optional[paddle.Tensor] = None, + gr: paddle.Tensor | None = None, + g2: paddle.Tensor | None = None, + h2: paddle.Tensor | None = None, + fparam: paddle.Tensor | None = None, + aparam: paddle.Tensor | None = None, ): nf, nloc, _ = descriptor.shape ret = {} diff --git a/source/tests/pd/model/test_atomic_model_global_stat.py b/source/tests/pd/model/test_atomic_model_global_stat.py index abd7928a0f..053596bce2 100644 --- a/source/tests/pd/model/test_atomic_model_global_stat.py +++ b/source/tests/pd/model/test_atomic_model_global_stat.py @@ -4,9 +4,6 @@ from pathlib import ( Path, ) -from typing import ( - Optional, -) import h5py import numpy as np @@ -92,11 +89,11 @@ def forward( self, descriptor: paddle.Tensor, atype: paddle.Tensor, - gr: Optional[paddle.Tensor] = None, - g2: Optional[paddle.Tensor] = None, - h2: Optional[paddle.Tensor] = None, - fparam: Optional[paddle.Tensor] = None, - aparam: Optional[paddle.Tensor] = None, + gr: paddle.Tensor | None = None, + g2: paddle.Tensor | None = None, + h2: paddle.Tensor | None = None, + fparam: paddle.Tensor | None = None, + aparam: paddle.Tensor | None = None, ): nf, nloc, _ = descriptor.shape ret = {} diff --git a/source/tests/pd/model/test_force_grad.py b/source/tests/pd/model/test_force_grad.py index d7b569ef38..eb6975afec 100644 --- a/source/tests/pd/model/test_force_grad.py +++ b/source/tests/pd/model/test_force_grad.py @@ -5,9 +5,6 @@ from pathlib import ( Path, ) -from typing import ( - Optional, -) import numpy as np import paddle @@ -31,7 +28,7 @@ class CheckSymmetry(DeepmdData): def __init__( self, sys_path: str, - type_map: Optional[list[str]] = None, + type_map: list[str] | None = None, ): super().__init__(sys_path=sys_path, type_map=type_map) self.add("energy", 1, atomic=False, must=False, high_prec=True) diff --git a/source/tests/pd/model/test_rotation.py b/source/tests/pd/model/test_rotation.py index 94e3442631..48f5ae8983 100644 --- a/source/tests/pd/model/test_rotation.py +++ b/source/tests/pd/model/test_rotation.py @@ -4,9 +4,6 @@ from pathlib import ( Path, ) -from typing import ( - Optional, -) import numpy as np import paddle @@ -29,7 +26,7 @@ class CheckSymmetry(DeepmdData): def __init__( self, sys_path: str, - type_map: Optional[list[str]] = None, + type_map: list[str] | None = None, ): super().__init__(sys_path=sys_path, type_map=type_map) self.add("energy", 1, atomic=False, must=False, high_prec=True) diff --git a/source/tests/pt/common.py b/source/tests/pt/common.py index 2dbfdb84ff..0592be1423 100644 --- a/source/tests/pt/common.py +++ b/source/tests/pt/common.py @@ -1,9 +1,5 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import pathlib -from typing import ( - Optional, - Union, -) import numpy as np import torch @@ -51,10 +47,10 @@ def run_dp(cmd: str) -> int: def eval_model( model, - coords: Union[np.ndarray, torch.Tensor], - cells: Optional[Union[np.ndarray, torch.Tensor]], - atom_types: Union[np.ndarray, torch.Tensor, list[int]], - spins: Optional[Union[np.ndarray, torch.Tensor]] = None, + coords: np.ndarray | torch.Tensor, + cells: np.ndarray | torch.Tensor | None, + atom_types: np.ndarray | torch.Tensor | list[int], + spins: np.ndarray | torch.Tensor | None = None, atomic: bool = False, infer_batch_size: int = 2, denoise: bool = False, diff --git a/source/tests/pt/model/test_atomic_model_atomic_stat.py b/source/tests/pt/model/test_atomic_model_atomic_stat.py index e6e873e6c1..876b363dc8 100644 --- a/source/tests/pt/model/test_atomic_model_atomic_stat.py +++ b/source/tests/pt/model/test_atomic_model_atomic_stat.py @@ -6,7 +6,6 @@ ) from typing import ( NoReturn, - Optional, ) import h5py @@ -81,11 +80,11 @@ def forward( self, descriptor: torch.Tensor, atype: torch.Tensor, - gr: Optional[torch.Tensor] = None, - g2: Optional[torch.Tensor] = None, - h2: Optional[torch.Tensor] = None, - fparam: Optional[torch.Tensor] = None, - aparam: Optional[torch.Tensor] = None, + gr: torch.Tensor | None = None, + g2: torch.Tensor | None = None, + h2: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, ): nf, nloc, _ = descriptor.shape ret = {} diff --git a/source/tests/pt/model/test_atomic_model_global_stat.py b/source/tests/pt/model/test_atomic_model_global_stat.py index fbb0279c09..acba3b56ad 100644 --- a/source/tests/pt/model/test_atomic_model_global_stat.py +++ b/source/tests/pt/model/test_atomic_model_global_stat.py @@ -6,7 +6,6 @@ ) from typing import ( NoReturn, - Optional, ) import h5py @@ -93,11 +92,11 @@ def forward( self, descriptor: torch.Tensor, atype: torch.Tensor, - gr: Optional[torch.Tensor] = None, - g2: Optional[torch.Tensor] = None, - h2: Optional[torch.Tensor] = None, - fparam: Optional[torch.Tensor] = None, - aparam: Optional[torch.Tensor] = None, + gr: torch.Tensor | None = None, + g2: torch.Tensor | None = None, + h2: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, ): nf, nloc, _ = descriptor.shape ret = {} diff --git a/source/tests/pt/model/test_force_grad.py b/source/tests/pt/model/test_force_grad.py index 44dbffa536..27bf660241 100644 --- a/source/tests/pt/model/test_force_grad.py +++ b/source/tests/pt/model/test_force_grad.py @@ -5,9 +5,6 @@ from pathlib import ( Path, ) -from typing import ( - Optional, -) import numpy as np import torch @@ -31,7 +28,7 @@ class CheckSymmetry(DeepmdData): def __init__( self, sys_path: str, - type_map: Optional[list[str]] = None, + type_map: list[str] | None = None, ) -> None: super().__init__(sys_path=sys_path, type_map=type_map) self.add("energy", 1, atomic=False, must=False, high_prec=True) diff --git a/source/tests/pt/model/test_linear_atomic_model_stat.py b/source/tests/pt/model/test_linear_atomic_model_stat.py index 0f92f1253d..fdc4e6dac6 100644 --- a/source/tests/pt/model/test_linear_atomic_model_stat.py +++ b/source/tests/pt/model/test_linear_atomic_model_stat.py @@ -4,9 +4,6 @@ from pathlib import ( Path, ) -from typing import ( - Optional, -) import h5py import numpy as np @@ -73,11 +70,11 @@ def forward( self, descriptor: torch.Tensor, atype: torch.Tensor, - gr: Optional[torch.Tensor] = None, - g2: Optional[torch.Tensor] = None, - h2: Optional[torch.Tensor] = None, - fparam: Optional[torch.Tensor] = None, - aparam: Optional[torch.Tensor] = None, + gr: torch.Tensor | None = None, + g2: torch.Tensor | None = None, + h2: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, ): nf, nloc, _ = descriptor.shape ret = {} @@ -125,11 +122,11 @@ def forward( self, descriptor: torch.Tensor, atype: torch.Tensor, - gr: Optional[torch.Tensor] = None, - g2: Optional[torch.Tensor] = None, - h2: Optional[torch.Tensor] = None, - fparam: Optional[torch.Tensor] = None, - aparam: Optional[torch.Tensor] = None, + gr: torch.Tensor | None = None, + g2: torch.Tensor | None = None, + h2: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, ): nf, nloc, _ = descriptor.shape ret = {} diff --git a/source/tests/pt/model/test_polar_atomic_model_stat.py b/source/tests/pt/model/test_polar_atomic_model_stat.py index 19fb2a4d3f..c8b60f1164 100644 --- a/source/tests/pt/model/test_polar_atomic_model_stat.py +++ b/source/tests/pt/model/test_polar_atomic_model_stat.py @@ -6,7 +6,6 @@ ) from typing import ( NoReturn, - Optional, ) import h5py @@ -46,11 +45,11 @@ def forward( self, descriptor: torch.Tensor, atype: torch.Tensor, - gr: Optional[torch.Tensor] = None, - g2: Optional[torch.Tensor] = None, - h2: Optional[torch.Tensor] = None, - fparam: Optional[torch.Tensor] = None, - aparam: Optional[torch.Tensor] = None, + gr: torch.Tensor | None = None, + g2: torch.Tensor | None = None, + h2: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, ): nf, nloc, _ = descriptor.shape ret = {} diff --git a/source/tests/pt/model/test_rotation.py b/source/tests/pt/model/test_rotation.py index 713a957e99..f2abc6f9fd 100644 --- a/source/tests/pt/model/test_rotation.py +++ b/source/tests/pt/model/test_rotation.py @@ -4,9 +4,6 @@ from pathlib import ( Path, ) -from typing import ( - Optional, -) import numpy as np import torch @@ -29,7 +26,7 @@ class CheckSymmetry(DeepmdData): def __init__( self, sys_path: str, - type_map: Optional[list[str]] = None, + type_map: list[str] | None = None, ) -> None: super().__init__(sys_path=sys_path, type_map=type_map) self.add("energy", 1, atomic=False, must=False, high_prec=True) diff --git a/source/tests/universal/common/cases/atomic_model/utils.py b/source/tests/universal/common/cases/atomic_model/utils.py index 054272b8c6..45b30c6454 100644 --- a/source/tests/universal/common/cases/atomic_model/utils.py +++ b/source/tests/universal/common/cases/atomic_model/utils.py @@ -1,8 +1,9 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +from collections.abc import ( + Callable, +) from typing import ( Any, - Callable, - Optional, ) import numpy as np @@ -41,11 +42,11 @@ class AtomicModelTestCase: """Expected whether having message passing.""" forward_wrapper: Callable[[Any], Any] """Class wrapper for forward method.""" - aprec_dict: dict[str, Optional[float]] + aprec_dict: dict[str, float | None] """Dictionary of absolute precision in each test.""" - rprec_dict: dict[str, Optional[float]] + rprec_dict: dict[str, float | None] """Dictionary of relative precision in each test.""" - epsilon_dict: dict[str, Optional[float]] + epsilon_dict: dict[str, float | None] """Dictionary of epsilons in each test.""" def test_get_type_map(self) -> None: diff --git a/source/tests/universal/common/cases/model/utils.py b/source/tests/universal/common/cases/model/utils.py index 08a369933d..8ec0edc180 100644 --- a/source/tests/universal/common/cases/model/utils.py +++ b/source/tests/universal/common/cases/model/utils.py @@ -1,12 +1,13 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import unittest +from collections.abc import ( + Callable, +) from copy import ( deepcopy, ) from typing import ( Any, - Callable, - Optional, ) import numpy as np @@ -54,11 +55,11 @@ class ModelTestCase: """Class wrapper for forward method.""" forward_wrapper_cpu_ref: Callable[[Any], Any] """Convert model to CPU method.""" - aprec_dict: dict[str, Optional[float]] + aprec_dict: dict[str, float | None] """Dictionary of absolute precision in each test.""" - rprec_dict: dict[str, Optional[float]] + rprec_dict: dict[str, float | None] """Dictionary of relative precision in each test.""" - epsilon_dict: dict[str, Optional[float]] + epsilon_dict: dict[str, float | None] """Dictionary of epsilons in each test.""" def test_get_type_map(self) -> None: