diff --git a/deepmd/dpmodel/infer/deep_eval.py b/deepmd/dpmodel/infer/deep_eval.py index aa4c9fe79b..e86322866b 100644 --- a/deepmd/dpmodel/infer/deep_eval.py +++ b/deepmd/dpmodel/infer/deep_eval.py @@ -47,6 +47,9 @@ from deepmd.infer.deep_pot import ( DeepPot, ) +from deepmd.infer.deep_property import ( + DeepProperty, +) from deepmd.infer.deep_wfc import ( DeepWFC, ) @@ -134,8 +137,9 @@ def has_default_fparam(self) -> bool: @property def model_type(self) -> type["DeepEvalWrapper"]: - """The the evaluator of the model type.""" - model_output_type = self.dp.model_output_type() + """The evaluator of the model type.""" + model = self.get_model() + model_output_type = model.model_output_type() if "energy" in model_output_type: return DeepPot elif "dos" in model_output_type: @@ -146,6 +150,8 @@ def model_type(self) -> type["DeepEvalWrapper"]: return DeepPolar elif "wfc" in model_output_type: return DeepWFC + elif self._get_property_var_name(model) in model_output_type: + return DeepProperty else: raise RuntimeError("Unknown model type") @@ -238,6 +244,12 @@ def eval( out = self._eval_func(self._eval_model, numb_test, natoms)( coords, cells, atom_types, fparam, aparam, request_defs ) + # ``AutoBatchSize.execute_all`` unwraps a single-output result out of + # its tuple, which would make ``zip`` iterate over the array's frame + # axis. Re-wrap so the request-def names line up (a single request def + # arises for global-only DOS/property inference at atomic=False). + if not isinstance(out, tuple): + out = (out,) return dict( zip( [x.name for x in request_defs], diff --git a/deepmd/dpmodel/model/property_model.py b/deepmd/dpmodel/model/property_model.py index d3153b92f9..6d85e3b80a 100644 --- a/deepmd/dpmodel/model/property_model.py +++ b/deepmd/dpmodel/model/property_model.py @@ -43,6 +43,14 @@ def get_var_name(self) -> str: """Get the name of the property.""" return self.get_fitting_net().var_name + def get_task_dim(self) -> int: + """Get the output dimension of the property.""" + return self.get_fitting_net().dim_out + + def get_intensive(self) -> bool: + """Whether the property is intensive.""" + return self.model_output_def()[self.get_var_name()].intensive + def call( self, coord: Array, diff --git a/deepmd/infer/deep_eval.py b/deepmd/infer/deep_eval.py index ae76f45212..05e40854d1 100644 --- a/deepmd/infer/deep_eval.py +++ b/deepmd/infer/deep_eval.py @@ -375,7 +375,31 @@ def _check_mixed_types(self, atom_types: np.ndarray) -> bool: @property @abstractmethod def model_type(self) -> type["DeepEval"]: - """The the evaluator of the model type.""" + """The evaluator of the model type. + + Each backend implements the dispatch on its own module so it can import + the concrete ``Deep*`` wrapper classes at the top level. Those wrappers + import ``DeepEval`` from this module, so a dispatch here would form an + import cycle (flagged by CodeQL). :meth:`_get_property_var_name` is + provided for the shared property branch. + """ + + @staticmethod + def _get_property_var_name(model: Any) -> str | None: + """Return the property variable name of ``model``, or ``None``. + + Used by every backend's ``model_type`` to detect a property model. + ``get_var_name`` may be absent (dpmodel/pt live models expose it only on + property models) or present-but-unimplemented (jax/tf2 artifacts always + define it and raise ``NotImplementedError`` otherwise), so probe + defensively. + """ + if not hasattr(model, "get_var_name"): + return None + try: + return model.get_var_name() + except NotImplementedError: + return None @abstractmethod def get_sel_type(self) -> list[int]: @@ -414,7 +438,24 @@ def get_has_hessian(self) -> bool: return False def get_var_name(self) -> str: - """Get the name of the fitting property.""" + """Get the name of the fitting property (property models only).""" + model = self.get_model() + if hasattr(model, "get_var_name"): + return model.get_var_name() + raise NotImplementedError + + def get_task_dim(self) -> int: + """Get the output dimension of the property (property models only).""" + model = self.get_model() + if hasattr(model, "get_task_dim"): + return model.get_task_dim() + raise NotImplementedError + + def get_intensive(self) -> bool: + """Whether the property is intensive (property models only).""" + model = self.get_model() + if hasattr(model, "get_intensive"): + return model.get_intensive() raise NotImplementedError @abstractmethod diff --git a/deepmd/infer/deep_property.py b/deepmd/infer/deep_property.py index 5e35dcd781..aac65ce2e7 100644 --- a/deepmd/infer/deep_property.py +++ b/deepmd/infer/deep_property.py @@ -134,14 +134,13 @@ def eval( aparam=aparam, **kwargs, ) - atomic_property = results[self.get_var_name()].reshape( - nframes, natoms, self.get_task_dim() - ) property = results[f"{self.get_var_name()}_redu"].reshape( nframes, self.get_task_dim() ) - if atomic: + atomic_property = results[self.get_var_name()].reshape( + nframes, natoms, self.get_task_dim() + ) return ( property, atomic_property, diff --git a/deepmd/jax/infer/deep_eval.py b/deepmd/jax/infer/deep_eval.py index d5d769784a..ef65f93a38 100644 --- a/deepmd/jax/infer/deep_eval.py +++ b/deepmd/jax/infer/deep_eval.py @@ -41,6 +41,9 @@ from deepmd.infer.deep_pot import ( DeepPot, ) +from deepmd.infer.deep_property import ( + DeepProperty, +) from deepmd.infer.deep_wfc import ( DeepWFC, ) @@ -152,7 +155,8 @@ def get_dim_aparam(self) -> int: @property def model_type(self) -> type["DeepEvalWrapper"]: """The evaluator of the model type.""" - model_output_type = self.dp.model_output_type() + model = self.get_model() + model_output_type = model.model_output_type() if "energy" in model_output_type: return DeepPot elif "dos" in model_output_type: @@ -163,6 +167,8 @@ def model_type(self) -> type["DeepEvalWrapper"]: return DeepPolar elif "wfc" in model_output_type: return DeepWFC + elif self._get_property_var_name(model) in model_output_type: + return DeepProperty else: raise RuntimeError("Unknown model type") @@ -270,6 +276,12 @@ def eval( out = self._eval_func(self._eval_model, numb_test, natoms)( coords, cells, atom_types, fparam, aparam, request_defs ) + # ``AutoBatchSize.execute_all`` unwraps a single-output result out of + # its tuple, which would make ``zip`` iterate over the array's frame + # axis. Re-wrap so the request-def names line up (a single request def + # arises for global-only DOS/property inference at atomic=False). + if not isinstance(out, tuple): + out = (out,) return dict( zip( [x.name for x in request_defs], diff --git a/deepmd/jax/jax2tf/serialization.py b/deepmd/jax/jax2tf/serialization.py index dd496dedf0..baf3c65b84 100644 --- a/deepmd/jax/jax2tf/serialization.py +++ b/deepmd/jax/jax2tf/serialization.py @@ -318,6 +318,28 @@ def get_default_fparam() -> tf.Tensor: tf_model.get_default_fparam = get_default_fparam + # property models: persist the output name/dimension/intensiveness so + # the evaluator can dispatch to DeepProperty and reshape the output. + if hasattr(model, "get_var_name"): + + @tf.function + def get_var_name() -> tf.Tensor: + return tf.constant(model.get_var_name(), dtype=tf.string) + + tf_model.get_var_name = get_var_name + + @tf.function + def get_task_dim() -> tf.Tensor: + return tf.constant(model.get_task_dim(), dtype=tf.int64) + + tf_model.get_task_dim = get_task_dim + + @tf.function + def get_intensive() -> tf.Tensor: + return tf.constant(model.get_intensive(), dtype=tf.bool) + + tf_model.get_intensive = get_intensive + tf.saved_model.save( tf_model, model_file, diff --git a/deepmd/jax/jax2tf/tfmodel.py b/deepmd/jax/jax2tf/tfmodel.py index 2d820810cc..6b9a1fcab9 100644 --- a/deepmd/jax/jax2tf/tfmodel.py +++ b/deepmd/jax/jax2tf/tfmodel.py @@ -78,6 +78,15 @@ def __init__( self.default_fparam = self.model.get_default_fparam().numpy().tolist() else: self.default_fparam = None + # property models only (absent for other model types). + if hasattr(self.model, "get_var_name"): + self._var_name = self.model.get_var_name().numpy().decode() + self._task_dim = self.model.get_task_dim().numpy().item() + self._intensive = self.model.get_intensive().numpy().item() + else: + self._var_name = None + self._task_dim = None + self._intensive = False def __call__( self, @@ -175,9 +184,27 @@ def call( def model_output_def(self) -> ModelOutputDef: return ModelOutputDef( - FittingOutputDef([OUTPUT_DEFS[tt] for tt in self.model_output_type()]) + FittingOutputDef( + [self._output_var_def(tt) for tt in self.model_output_type()] + ) ) + def _output_var_def(self, name: str) -> OutputVariableDef: + if name in OUTPUT_DEFS: + return OUTPUT_DEFS[name] + # property models carry a user-defined output name (``var_name``) that + # is not in the fixed table; rebuild its def from the persisted metadata. + if self._var_name is not None and name == self._var_name: + return OutputVariableDef( + self._var_name, + shape=[self._task_dim], + reducible=True, + r_differentiable=False, + c_differentiable=False, + intensive=self._intensive, + ) + raise KeyError(f"Unknown model output variable {name!r}") + def call_lower( self, extended_coord: jnp.ndarray, @@ -349,3 +376,19 @@ def has_default_fparam(self) -> bool: def get_default_fparam(self) -> list[float] | None: """Get the default frame parameters.""" return self.default_fparam + + def get_var_name(self) -> str: + """Get the name of the property (property models only).""" + if self._var_name is None: + raise NotImplementedError + return self._var_name + + def get_task_dim(self) -> int: + """Get the output dimension of the property (property models only).""" + if self._task_dim is None: + raise NotImplementedError + return self._task_dim + + def get_intensive(self) -> bool: + """Whether the property is intensive (property models only).""" + return self._intensive diff --git a/deepmd/jax/model/hlo.py b/deepmd/jax/model/hlo.py index 0162850eaa..72bc943c56 100644 --- a/deepmd/jax/model/hlo.py +++ b/deepmd/jax/model/hlo.py @@ -69,6 +69,10 @@ def __init__( has_default_fparam: bool = False, default_fparam: list[float] | None = None, numb_dos: int = 0, + # property models only + var_name: str | None = None, + task_dim: int | None = None, + intensive: bool = False, ) -> None: self._call_lower = jax_export.deserialize(stablehlo).call self._call_lower_atomic_virial = jax_export.deserialize( @@ -93,6 +97,9 @@ def __init__( self._has_default_fparam = has_default_fparam self.default_fparam = default_fparam self.numb_dos = numb_dos + self._var_name = var_name + self._task_dim = task_dim + self._intensive = intensive def __call__( self, @@ -180,9 +187,27 @@ def call( def model_output_def(self) -> ModelOutputDef: return ModelOutputDef( - FittingOutputDef([OUTPUT_DEFS[tt] for tt in self.model_output_type()]) + FittingOutputDef( + [self._output_var_def(tt) for tt in self.model_output_type()] + ) ) + def _output_var_def(self, name: str) -> OutputVariableDef: + if name in OUTPUT_DEFS: + return OUTPUT_DEFS[name] + # property models carry a user-defined output name (``var_name``) that + # is not in the fixed table; rebuild its def from the persisted metadata. + if self._var_name is not None and name == self._var_name: + return OutputVariableDef( + self._var_name, + shape=[self._task_dim], + reducible=True, + r_differentiable=False, + c_differentiable=False, + intensive=self._intensive, + ) + raise KeyError(f"Unknown model output variable {name!r}") + def call_lower( self, extended_coord: jnp.ndarray, @@ -233,6 +258,22 @@ def get_dim_aparam(self) -> int: """Get the number (dimension) of atomic parameters of this atomic model.""" return self.dim_aparam + def get_var_name(self) -> str: + """Get the name of the property (property models only).""" + if self._var_name is None: + raise NotImplementedError + return self._var_name + + def get_task_dim(self) -> int: + """Get the output dimension of the property (property models only).""" + if self._task_dim is None: + raise NotImplementedError + return self._task_dim + + def get_intensive(self) -> bool: + """Whether the property is intensive (property models only).""" + return self._intensive + def get_sel_type(self) -> list[int]: """Get the selected atom types of this model. diff --git a/deepmd/jax/utils/serialization.py b/deepmd/jax/utils/serialization.py index d77fabe6a4..39354cc1fe 100644 --- a/deepmd/jax/utils/serialization.py +++ b/deepmd/jax/utils/serialization.py @@ -321,6 +321,18 @@ def call_lower_with_fixed_do_atomic_virial( "sel": model.get_sel(), "has_default_fparam": model.has_default_fparam(), "default_fparam": model.get_default_fparam(), + # property models: the output name/dimension/intensiveness cannot be + # recovered from the StableHLO alone, so persist them for the + # evaluator (None for non-property models). + "var_name": model.get_var_name() + if hasattr(model, "get_var_name") + else None, + "task_dim": model.get_task_dim() + if hasattr(model, "get_task_dim") + else None, + "intensive": model.get_intensive() + if hasattr(model, "get_intensive") + else False, } save_dp_model(filename=model_file, model_dict=data) elif model_file.endswith(".savedmodel"): diff --git a/deepmd/pt_expt/infer/deep_eval.py b/deepmd/pt_expt/infer/deep_eval.py index 58b62aaf56..f03b6facfb 100644 --- a/deepmd/pt_expt/infer/deep_eval.py +++ b/deepmd/pt_expt/infer/deep_eval.py @@ -48,6 +48,9 @@ from deepmd.infer.deep_pot import ( DeepPot, ) +from deepmd.infer.deep_property import ( + DeepProperty, +) from deepmd.infer.deep_wfc import ( DeepWFC, ) @@ -740,6 +743,12 @@ def model_type(self) -> type["DeepEvalWrapper"]: return DeepPolar elif "wfc" in model_output_type: return DeepWFC + elif ( + self._dpmodel is not None + and hasattr(self._dpmodel, "get_var_name") + and self._dpmodel.get_var_name() in model_output_type + ): + return DeepProperty else: raise RuntimeError("Unknown model type") @@ -761,6 +770,33 @@ def get_numb_dos(self) -> int: """Get the number of DOS.""" return 0 + def get_var_name(self) -> str: + """Get the name of the property (property models only).""" + if self._dpmodel is not None and hasattr(self._dpmodel, "get_var_name"): + return self._dpmodel.get_var_name() + raise NotImplementedError( + "get_var_name is only available for property models with the " + "reconstructed dpmodel (not in metadata-only mode)." + ) + + def get_task_dim(self) -> int: + """Get the output dimension of the property (property models only).""" + if self._dpmodel is not None and hasattr(self._dpmodel, "get_task_dim"): + return self._dpmodel.get_task_dim() + raise NotImplementedError( + "get_task_dim is only available for property models with the " + "reconstructed dpmodel (not in metadata-only mode)." + ) + + def get_intensive(self) -> bool: + """Whether the property is intensive (property models only).""" + if self._dpmodel is not None and hasattr(self._dpmodel, "get_intensive"): + return self._dpmodel.get_intensive() + raise NotImplementedError( + "get_intensive is only available for property models with the " + "reconstructed dpmodel (not in metadata-only mode)." + ) + def get_has_efield(self) -> bool: """Check if the model has efield.""" return False diff --git a/deepmd/pt_expt/model/property_model.py b/deepmd/pt_expt/model/property_model.py index 6d8470f142..868eb81c04 100644 --- a/deepmd/pt_expt/model/property_model.py +++ b/deepmd/pt_expt/model/property_model.py @@ -41,6 +41,14 @@ def get_var_name(self) -> str: """Get the name of the property.""" return self.get_fitting_net().var_name + def get_task_dim(self) -> int: + """Get the output dimension of the property.""" + return self.get_fitting_net().dim_out + + def get_intensive(self) -> bool: + """Whether the property is intensive.""" + return self.model_output_def()[self.get_var_name()].intensive + def forward( self, coord: torch.Tensor, diff --git a/deepmd/tf2/infer/deep_eval.py b/deepmd/tf2/infer/deep_eval.py index 4d6a8c5e2f..af5037e2e6 100644 --- a/deepmd/tf2/infer/deep_eval.py +++ b/deepmd/tf2/infer/deep_eval.py @@ -39,6 +39,9 @@ from deepmd.infer.deep_pot import ( DeepPot, ) +from deepmd.infer.deep_property import ( + DeepProperty, +) from deepmd.infer.deep_wfc import ( DeepWFC, ) @@ -93,6 +96,22 @@ def __init__(self, model: str) -> None: if hasattr(self.model, "get_default_fparam") else None ) + # property models only (absent for other model types). + self._var_name = ( + self.model.get_var_name().numpy().decode() + if hasattr(self.model, "get_var_name") + else None + ) + self._task_dim = ( + self.model.get_task_dim().numpy().item() + if hasattr(self.model, "get_task_dim") + else None + ) + self._intensive = ( + self.model.get_intensive().numpy().item() + if hasattr(self.model, "get_intensive") + else False + ) def __call__( self, @@ -163,6 +182,22 @@ def has_default_fparam(self) -> bool: def get_default_fparam(self) -> list[float] | None: return self.default_fparam + def get_var_name(self) -> str: + """Get the name of the property (property models only).""" + if self._var_name is None: + raise NotImplementedError + return self._var_name + + def get_task_dim(self) -> int: + """Get the output dimension of the property (property models only).""" + if self._task_dim is None: + raise NotImplementedError + return self._task_dim + + def get_intensive(self) -> bool: + """Whether the property is intensive (property models only).""" + return self._intensive + class DeepEval(DeepEvalBackend): """TensorFlow 2 SavedModel backend implementation of DeepEval.""" @@ -212,7 +247,8 @@ def has_default_fparam(self) -> bool: @property def model_type(self) -> type["DeepEvalWrapper"]: - model_output_type = self.dp.model_output_type() + model = self.get_model() + model_output_type = model.model_output_type() if "energy" in model_output_type: return DeepPot if "dos" in model_output_type: @@ -223,6 +259,8 @@ def model_type(self) -> type["DeepEvalWrapper"]: return DeepPolar if "wfc" in model_output_type: return DeepWFC + if self._get_property_var_name(model) in model_output_type: + return DeepProperty raise RuntimeError("Unknown model type") def get_sel_type(self) -> list[int]: @@ -258,6 +296,12 @@ def eval( out = self._eval_func(self._eval_model, numb_test, natoms)( coords, cells, atom_types, fparam, aparam, request_defs ) + # ``AutoBatchSize.execute_all`` unwraps a single-output result out of + # its tuple, which would make ``zip`` iterate over the array's frame + # axis. Re-wrap so the request-def names line up (a single request def + # arises for global-only DOS/property inference at atomic=False). + if not isinstance(out, tuple): + out = (out,) return dict(zip([x.name for x in request_defs], out, strict=True)) def _get_request_defs(self, atomic: bool) -> list[OutputVariableDef]: diff --git a/deepmd/tf2/utils/serialization.py b/deepmd/tf2/utils/serialization.py index 5de81a2bdf..0fcd6e7ec8 100644 --- a/deepmd/tf2/utils/serialization.py +++ b/deepmd/tf2/utils/serialization.py @@ -533,6 +533,28 @@ def get_default_fparam() -> tf.Tensor: tf_model.get_default_fparam = get_default_fparam + # property models: persist the output name/dimension/intensiveness so the + # evaluator can dispatch to DeepProperty and reshape the output. + if hasattr(model, "get_var_name"): + + @tf.function + def get_var_name() -> tf.Tensor: + return tf.constant(model.get_var_name(), dtype=tf.string) + + tf_model.get_var_name = get_var_name + + @tf.function + def get_task_dim() -> tf.Tensor: + return tf.constant(model.get_task_dim(), dtype=tf.int64) + + tf_model.get_task_dim = get_task_dim + + @tf.function + def get_intensive() -> tf.Tensor: + return tf.constant(model.get_intensive(), dtype=tf.bool) + + tf_model.get_intensive = get_intensive + tf.saved_model.save( tf_model, model_file, diff --git a/source/tests/consistent/io/test_io.py b/source/tests/consistent/io/test_io.py index 9601ac7607..6ccb906bae 100644 --- a/source/tests/consistent/io/test_io.py +++ b/source/tests/consistent/io/test_io.py @@ -5,6 +5,9 @@ from pathlib import ( Path, ) +from typing import ( + ClassVar, +) import numpy as np @@ -32,6 +35,9 @@ class IOTest: data: dict + # backends that cannot represent this model type (e.g. tf v1 has no + # property model), skipped by the cross-backend round trips below. + skip_backends: ClassVar[set[str]] = set() def get_data_from_model(self, model_file: str) -> dict: """Get data from a model file. @@ -81,6 +87,8 @@ def test_data_equal(self) -> None: ("dpmodel", 0), ): with self.subTest(backend_name=backend_name): + if backend_name in self.skip_backends: + continue backend = Backend.get_backend(backend_name)() if not backend.is_available(): continue @@ -151,7 +159,7 @@ def test_deep_eval(self) -> None: ("dpmodel", 0), ("jax", 0) if DP_TEST_TF2_ONLY else (None, None), ): - if backend_name is None: + if backend_name is None or backend_name in self.skip_backends: continue backend = Backend.get_backend(backend_name)() if not backend.is_available(): @@ -180,6 +188,10 @@ def test_deep_eval(self) -> None: fparam=fparam, aparam=aparam, ) + # the non-atomic eval returns exactly the global outputs; the atomic + # eval appends the per-atom outputs. Split by that count so this + # generalizes across model types (energy: 3 global, property: 1). + n_global = len(ret) rets.append(ret) ret = deep_eval.eval( self.coords, @@ -189,8 +201,8 @@ def test_deep_eval(self) -> None: aparam=aparam, atomic=True, ) - rets.append(ret[:3]) - rets_atomic.append(ret[3:]) + rets.append(ret[:n_global]) + rets_atomic.append(ret[n_global:]) ret = deep_eval.eval( self.coords, None, @@ -207,8 +219,8 @@ def test_deep_eval(self) -> None: aparam=aparam, atomic=True, ) - rets_nopbc.append(ret[:3]) - rets_nopbc_atomic.append(ret[3:]) + rets_nopbc.append(ret[:n_global]) + rets_nopbc_atomic.append(ret[n_global:]) for rets_idx, rets_x in enumerate( (rets, rets_atomic, rets_nopbc, rets_nopbc_atomic) @@ -311,3 +323,50 @@ def setUp(self) -> None: def tearDown(self) -> None: IOTest.tearDown(self) + + +class TestDeepProperty(unittest.TestCase, IOTest): + # tf v1 has no property model, so the property dict cannot round trip + # through the tensorflow backend. + skip_backends: ClassVar[set[str]] = {"tensorflow"} + + def setUp(self) -> None: + model_def_script = { + "type_map": ["O", "H"], + "descriptor": { + "type": "se_e2_a", + "sel": [20, 20], + "rcut_smth": 0.50, + "rcut": 6.00, + "neuron": [ + 3, + 6, + ], + "resnet_dt": False, + "axis_neuron": 2, + "precision": "float64", + "type_one_side": True, + "seed": 1, + }, + "fitting_net": { + "type": "property", + "neuron": [ + 5, + 5, + ], + "property_name": "foo", + "task_dim": 3, + "resnet_dt": True, + "precision": "float64", + "seed": 1, + }, + } + model = get_model(copy.deepcopy(model_def_script)) + self.data = { + "model": model.serialize(), + "backend": "test", + "model_def_script": model_def_script, + } + + def tearDown(self) -> None: + IOTest.tearDown(self) diff --git a/source/tests/pt_expt/infer/test_deep_eval_property.py b/source/tests/pt_expt/infer/test_deep_eval_property.py new file mode 100644 index 0000000000..4c85440d12 --- /dev/null +++ b/source/tests/pt_expt/infer/test_deep_eval_property.py @@ -0,0 +1,98 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Tests for pt_expt property-model inference via the DeepProperty interface. + +Verifies the full pipeline: + PropertyModel.serialize() -> deserialize_to_file(.pte) -> DeepEval(.pte) + +A pt_expt property checkpoint could be trained but not evaluated: the pt_expt +DeepEval.model_type did not dispatch property models (raising "Unknown model +type") and did not expose the property metadata getters. This checks the +dispatch and the get_var_name / get_task_dim / get_intensive getters. +""" + +import os +import tempfile +import unittest + +import numpy as np +import torch + +from deepmd.infer import ( + DeepEval, +) +from deepmd.infer.deep_property import ( + DeepProperty, +) +from deepmd.pt_expt.descriptor.se_e2_a import ( + DescrptSeA, +) +from deepmd.pt_expt.fitting import ( + PropertyFittingNet, +) +from deepmd.pt_expt.model import ( + PropertyModel, +) +from deepmd.pt_expt.utils.serialization import ( + deserialize_to_file, +) + +from ...seed import ( + GLOBAL_SEED, +) + + +class TestDeepEvalProperty(unittest.TestCase): + @classmethod + def setUpClass(cls) -> None: + cls.rcut = 4.0 + cls.rcut_smth = 0.5 + cls.sel = [8, 6] + cls.nt = 2 + cls.task_dim = 3 + cls.type_map = ["foo", "bar"] + + ds = DescrptSeA(cls.rcut, cls.rcut_smth, cls.sel) + ft = PropertyFittingNet( + cls.nt, + ds.get_dim_out(), + task_dim=cls.task_dim, + mixed_types=ds.mixed_types(), + seed=GLOBAL_SEED, + ) + cls.model = PropertyModel(ds, ft, type_map=cls.type_map).to(torch.float64) + cls.model.eval() + + cls.tmpfile = tempfile.NamedTemporaryFile(suffix=".pte", delete=False) + cls.tmpfile.close() + deserialize_to_file(cls.tmpfile.name, {"model": cls.model.serialize()}) + cls.dp = DeepEval(cls.tmpfile.name) + + @classmethod + def tearDownClass(cls) -> None: + os.unlink(cls.tmpfile.name) + + def test_model_type_is_property(self) -> None: + self.assertIs(self.dp.deep_eval.model_type, DeepProperty) + self.assertIsInstance(self.dp, DeepProperty) + + def test_get_var_name(self) -> None: + self.assertEqual(self.dp.deep_eval.get_var_name(), self.model.get_var_name()) + + def test_get_task_dim(self) -> None: + self.assertEqual(self.dp.deep_eval.get_task_dim(), self.task_dim) + + def test_get_intensive(self) -> None: + self.assertEqual(self.dp.deep_eval.get_intensive(), self.model.get_intensive()) + + def test_eval_shape(self) -> None: + rng = np.random.default_rng(0) + coords = rng.random([1, 4, 3]) * 3.0 + cells = (np.eye(3) * 10.0).reshape(1, 9) + atypes = np.array([[0, 1, 0, 1]], dtype=np.int32) + prop = self.dp.eval(coords, cells, atypes)[0] + self.assertEqual(prop.shape[0], 1) + self.assertEqual(prop.reshape(1, -1).shape[1], self.task_dim) + + +if __name__ == "__main__": + unittest.main()