Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions deepmd/dpmodel/atomic_model/base_atomic_model.py
Original file line number Diff line number Diff line change
Expand Up @@ -88,6 +88,10 @@ def get_type_map(self) -> list[str]:
"""Get the type map."""
return self.type_map

def has_default_fparam(self) -> bool:
"""Check if the model has default frame parameters."""
return False

def reinit_atom_exclude(
self,
exclude_types: list[int] = [],
Expand Down
4 changes: 4 additions & 0 deletions deepmd/dpmodel/atomic_model/dp_atomic_model.py
Original file line number Diff line number Diff line change
Expand Up @@ -233,6 +233,10 @@ def get_dim_aparam(self) -> int:
"""Get the number (dimension) of atomic parameters of this atomic model."""
return self.fitting.get_dim_aparam()

def has_default_fparam(self) -> bool:
"""Check if the model has default frame parameters."""
return self.fitting.has_default_fparam()

def get_sel_type(self) -> list[int]:
"""Get the selected atom types of this model.

Expand Down
7 changes: 6 additions & 1 deletion deepmd/dpmodel/fitting/dipole_fitting.py
Original file line number Diff line number Diff line change
Expand Up @@ -84,6 +84,9 @@ class DipoleFitting(GeneralFitting):
Only reducible variable are differentiable.
type_map: list[str], Optional
A list of strings. Give the name to each type of atoms.
default_fparam: list[float], optional
The default frame parameter. If set, when `fparam.npy` files are not included in the data system,
this value will be used as the default value for the frame parameter in the fitting net.
"""

def __init__(
Expand All @@ -110,6 +113,7 @@ def __init__(
c_differentiable: bool = True,
type_map: Optional[list[str]] = None,
seed: Optional[Union[int, list[int]]] = None,
default_fparam: Optional[list[float]] = None,
) -> None:
if tot_ener_zero:
raise NotImplementedError("tot_ener_zero is not implemented")
Expand Down Expand Up @@ -144,6 +148,7 @@ def __init__(
exclude_types=exclude_types,
type_map=type_map,
seed=seed,
default_fparam=default_fparam,
)

def _net_out_dim(self):
Expand All @@ -161,7 +166,7 @@ def serialize(self) -> dict:
@classmethod
def deserialize(cls, data: dict) -> "GeneralFitting":
data = data.copy()
check_version_compatibility(data.pop("@version", 1), 3, 1)
check_version_compatibility(data.pop("@version", 1), 4, 1)
var_name = data.pop("var_name", None)
assert var_name == "dipole"
return super().deserialize(data)
Expand Down
4 changes: 3 additions & 1 deletion deepmd/dpmodel/fitting/dos_fitting.py
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,7 @@ def __init__(
exclude_types: list[int] = [],
type_map: Optional[list[str]] = None,
seed: Optional[Union[int, list[int]]] = None,
default_fparam: Optional[list] = None,
) -> None:
if bias_dos is not None:
self.bias_dos = bias_dos
Expand All @@ -70,12 +71,13 @@ def __init__(
exclude_types=exclude_types,
type_map=type_map,
seed=seed,
default_fparam=default_fparam,
)

@classmethod
def deserialize(cls, data: dict) -> "GeneralFitting":
data = data.copy()
check_version_compatibility(data.pop("@version", 1), 3, 1)
check_version_compatibility(data.pop("@version", 1), 4, 1)
data["numb_dos"] = data.pop("dim_out")
data.pop("tot_ener_zero", None)
data.pop("var_name", None)
Expand Down
4 changes: 3 additions & 1 deletion deepmd/dpmodel/fitting/ener_fitting.py
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,7 @@ def __init__(
exclude_types: list[int] = [],
type_map: Optional[list[str]] = None,
seed: Optional[Union[int, list[int]]] = None,
default_fparam: Optional[list] = None,
) -> None:
super().__init__(
var_name="energy",
Expand All @@ -70,12 +71,13 @@ def __init__(
exclude_types=exclude_types,
type_map=type_map,
seed=seed,
default_fparam=default_fparam,
)

@classmethod
def deserialize(cls, data: dict) -> "GeneralFitting":
data = data.copy()
check_version_compatibility(data.pop("@version", 1), 3, 1)
check_version_compatibility(data.pop("@version", 1), 4, 1)
data.pop("var_name")
data.pop("dim_out")
return super().deserialize(data)
Expand Down
33 changes: 32 additions & 1 deletion deepmd/dpmodel/fitting/general_fitting.py
Original file line number Diff line number Diff line change
Expand Up @@ -94,6 +94,9 @@ class GeneralFitting(NativeOP, BaseFitting):
A list of strings. Give the name to each type of atoms.
seed: Optional[Union[int, list[int]]]
Random seed for initializing the network parameters.
default_fparam: list[float], optional
The default frame parameter. If set, when `fparam.npy` files are not included in the data system,
this value will be used as the default value for the frame parameter in the fitting net.
"""

def __init__(
Expand All @@ -120,6 +123,7 @@ def __init__(
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,
) -> None:
self.var_name = var_name
self.ntypes = ntypes
Expand All @@ -129,6 +133,7 @@ def __init__(
self.numb_fparam = numb_fparam
self.numb_aparam = numb_aparam
self.dim_case_embd = dim_case_embd
self.default_fparam = default_fparam
self.rcond = rcond
self.tot_ener_zero = tot_ener_zero
self.trainable = trainable
Expand Down Expand Up @@ -177,6 +182,15 @@ def __init__(
self.case_embd = np.zeros(self.dim_case_embd, dtype=self.prec)
else:
self.case_embd = None

if self.default_fparam is not None:
if self.numb_fparam > 0:
assert len(self.default_fparam) == self.numb_fparam, (
"default_fparam length mismatch!"
)
self.default_fparam_tensor = np.array(self.default_fparam, dtype=self.prec)
else:
self.default_fparam_tensor = None
# init networks
in_dim = (
self.dim_descrpt
Expand Down Expand Up @@ -217,6 +231,10 @@ def get_dim_aparam(self) -> int:
"""Get the number (dimension) of atomic parameters of this atomic model."""
return self.numb_aparam

def has_default_fparam(self) -> bool:
"""Check if the fitting has default frame parameters."""
return self.default_fparam is not None

def get_sel_type(self) -> list[int]:
"""Get the selected atom types of this model.

Expand Down Expand Up @@ -274,6 +292,8 @@ def __setitem__(self, key, value) -> None:
self.case_embd = value
elif key in ["scale"]:
self.scale = value
elif key in ["default_fparam_tensor"]:
self.default_fparam_tensor = value
else:
raise KeyError(key)

Expand All @@ -292,6 +312,8 @@ def __getitem__(self, key):
return self.case_embd
elif key in ["scale"]:
return self.scale
elif key in ["default_fparam_tensor"]:
return self.default_fparam_tensor
else:
raise KeyError(key)

Expand All @@ -306,7 +328,7 @@ def serialize(self) -> dict:
"""Serialize the fitting to dict."""
return {
"@class": "Fitting",
"@version": 3,
"@version": 4,
"var_name": self.var_name,
"ntypes": self.ntypes,
"dim_descrpt": self.dim_descrpt,
Expand All @@ -315,6 +337,7 @@ def serialize(self) -> dict:
"numb_fparam": self.numb_fparam,
"numb_aparam": self.numb_aparam,
"dim_case_embd": self.dim_case_embd,
"default_fparam": self.default_fparam,
"rcond": self.rcond,
"activation_function": self.activation_function,
"precision": self.precision,
Expand Down Expand Up @@ -403,6 +426,14 @@ def _call_common(
xx_zeros = xp.zeros_like(xx)
else:
xx_zeros = None

if self.numb_fparam > 0 and fparam is None:
# use default fparam
assert self.default_fparam_tensor is not None
fparam = xp.tile(
xp.reshape(self.default_fparam_tensor, (1, self.numb_fparam)), (nf, 1)
)

# check fparam dim, concate to input descriptor
if self.numb_fparam > 0:
assert fparam is not None, "fparam should not be None"
Expand Down
7 changes: 6 additions & 1 deletion deepmd/dpmodel/fitting/invar_fitting.py
Original file line number Diff line number Diff line change
Expand Up @@ -110,6 +110,9 @@ class InvarFitting(GeneralFitting):
Atomic contributions of the excluded atom types are set zero.
type_map: list[str], Optional
A list of strings. Give the name to each type of atoms.
default_fparam: list[float], optional
The default frame parameter. If set, when `fparam.npy` files are not included in the data system,
this value will be used as the default value for the frame parameter in the fitting net.

"""

Expand Down Expand Up @@ -138,6 +141,7 @@ def __init__(
exclude_types: list[int] = [],
type_map: Optional[list[str]] = None,
seed: Optional[Union[int, list[int]]] = None,
default_fparam: Optional[list[float]] = None,
) -> None:
if tot_ener_zero:
raise NotImplementedError("tot_ener_zero is not implemented")
Expand Down Expand Up @@ -173,6 +177,7 @@ def __init__(
else [x is not None for x in atom_ener],
type_map=type_map,
seed=seed,
default_fparam=default_fparam,
)

def serialize(self) -> dict:
Expand All @@ -185,7 +190,7 @@ def serialize(self) -> dict:
@classmethod
def deserialize(cls, data: dict) -> "GeneralFitting":
data = data.copy()
check_version_compatibility(data.pop("@version", 1), 3, 1)
check_version_compatibility(data.pop("@version", 1), 4, 1)
return super().deserialize(data)

def _net_out_dim(self):
Expand Down
9 changes: 7 additions & 2 deletions deepmd/dpmodel/fitting/polarizability_fitting.py
Original file line number Diff line number Diff line change
Expand Up @@ -90,6 +90,9 @@ class PolarFitting(GeneralFitting):
Whether to shift the diagonal part of the polarizability matrix. The shift operation is carried out after scale.
type_map: list[str], Optional
A list of strings. Give the name to each type of atoms.
default_fparam: list[float], optional
The default frame parameter. If set, when `fparam.npy` files are not included in the data system,
this value will be used as the default value for the frame parameter in the fitting net.
"""

def __init__(
Expand Down Expand Up @@ -117,6 +120,7 @@ def __init__(
shift_diag: bool = True,
type_map: Optional[list[str]] = None,
seed: Optional[Union[int, list[int]]] = None,
default_fparam: Optional[list[float]] = None,
) -> None:
if tot_ener_zero:
raise NotImplementedError("tot_ener_zero is not implemented")
Expand Down Expand Up @@ -164,6 +168,7 @@ def __init__(
exclude_types=exclude_types,
type_map=type_map,
seed=seed,
default_fparam=default_fparam,
)

def _net_out_dim(self):
Expand All @@ -189,7 +194,7 @@ def __getitem__(self, key):
def serialize(self) -> dict:
data = super().serialize()
data["type"] = "polar"
data["@version"] = 4
data["@version"] = 5
data["embedding_width"] = self.embedding_width
data["fit_diag"] = self.fit_diag
data["shift_diag"] = self.shift_diag
Expand All @@ -200,7 +205,7 @@ def serialize(self) -> dict:
@classmethod
def deserialize(cls, data: dict) -> "GeneralFitting":
data = data.copy()
check_version_compatibility(data.pop("@version", 1), 4, 1)
check_version_compatibility(data.pop("@version", 1), 5, 1)
var_name = data.pop("var_name", None)
assert var_name == "polar"
return super().deserialize(data)
Expand Down
9 changes: 7 additions & 2 deletions deepmd/dpmodel/fitting/property_fitting.py
Original file line number Diff line number Diff line change
Expand Up @@ -65,6 +65,9 @@ class PropertyFittingNet(InvarFitting):
Atomic contributions of the excluded atom types are set zero.
type_map: list[str], Optional
A list of strings. Give the name to each type of atoms.
default_fparam: list[float], optional
The default frame parameter. If set, when `fparam.npy` files are not included in the data system,
this value will be used as the default value for the frame parameter in the fitting net.
"""

def __init__(
Expand All @@ -87,6 +90,7 @@ def __init__(
mixed_types: bool = True,
exclude_types: list[int] = [],
type_map: Optional[list[str]] = None,
default_fparam: Optional[list] = None,
# not used
seed: Optional[int] = None,
) -> None:
Expand All @@ -110,6 +114,7 @@ def __init__(
mixed_types=mixed_types,
exclude_types=exclude_types,
type_map=type_map,
default_fparam=default_fparam,
)

def output_def(self) -> FittingOutputDef:
Expand All @@ -129,7 +134,7 @@ def output_def(self) -> FittingOutputDef:
@classmethod
def deserialize(cls, data: dict) -> "PropertyFittingNet":
data = data.copy()
check_version_compatibility(data.pop("@version"), 4, 1)
check_version_compatibility(data.pop("@version"), 5, 1)
data.pop("dim_out")
data["property_name"] = data.pop("var_name")
data.pop("tot_ener_zero")
Expand All @@ -149,6 +154,6 @@ def serialize(self) -> dict:
"task_dim": self.task_dim,
"intensive": self.intensive,
}
dd["@version"] = 4
dd["@version"] = 5

return dd
4 changes: 4 additions & 0 deletions deepmd/dpmodel/infer/deep_eval.py
Original file line number Diff line number Diff line change
Expand Up @@ -120,6 +120,10 @@ def get_dim_aparam(self) -> int:
"""Get the number (dimension) of atomic parameters of this DP."""
return self.dp.get_dim_aparam()

def has_default_fparam(self) -> bool:
"""Check if the model has default frame parameters."""
return self.dp.has_default_fparam()

@property
def model_type(self) -> type["DeepEvalWrapper"]:
"""The the evaluator of the model type."""
Expand Down
4 changes: 4 additions & 0 deletions deepmd/dpmodel/model/make_model.py
Original file line number Diff line number Diff line change
Expand Up @@ -564,6 +564,10 @@ def get_dim_aparam(self) -> int:
"""Get the number (dimension) of atomic parameters of this atomic model."""
return self.atomic_model.get_dim_aparam()

def has_default_fparam(self) -> bool:
"""Check if the model has default frame parameters."""
return self.atomic_model.has_default_fparam()

def get_sel_type(self) -> list[int]:
"""Get the selected atom types of this model.

Expand Down
8 changes: 6 additions & 2 deletions deepmd/entrypoints/test.py
Original file line number Diff line number Diff line change
Expand Up @@ -343,7 +343,11 @@ def test_ener(
data.add("atom_ener", 1, atomic=True, must=True, high_prec=False)
if dp.get_dim_fparam() > 0:
data.add(
"fparam", dp.get_dim_fparam(), atomic=False, must=True, high_prec=False
"fparam",
dp.get_dim_fparam(),
atomic=False,
must=not dp.has_default_fparam(),
high_prec=False,
)
if dp.get_dim_aparam() > 0:
data.add("aparam", dp.get_dim_aparam(), atomic=True, must=True, high_prec=False)
Expand Down Expand Up @@ -380,7 +384,7 @@ def test_ener(
atype = test_data["type"][:numb_test].reshape([numb_test, -1])
else:
atype = test_data["type"][0]
if dp.get_dim_fparam() > 0:
if dp.get_dim_fparam() > 0 and test_data["find_fparam"] != 0.0:
fparam = test_data["fparam"][:numb_test]
else:
fparam = None
Expand Down
8 changes: 8 additions & 0 deletions deepmd/infer/deep_eval.py
Original file line number Diff line number Diff line change
Expand Up @@ -162,6 +162,10 @@ def get_type_map(self) -> list[str]:
def get_dim_fparam(self) -> int:
"""Get the number (dimension) of frame parameters of this DP."""

def has_default_fparam(self) -> bool:
"""Check if the model has default frame parameters."""
return False

@abstractmethod
def get_dim_aparam(self) -> int:
"""Get the number (dimension) of atomic parameters of this DP."""
Expand Down Expand Up @@ -432,6 +436,10 @@ def get_dim_fparam(self) -> int:
"""Get the number (dimension) of frame parameters of this DP."""
return self.deep_eval.get_dim_fparam()

def has_default_fparam(self) -> bool:
"""Check if the model has default frame parameters."""
return self.deep_eval.has_default_fparam()

def get_dim_aparam(self) -> int:
"""Get the number (dimension) of atomic parameters of this DP."""
return self.deep_eval.get_dim_aparam()
Expand Down
Loading