fix: cast input and output types at model's interface

deepmd/dpmodel/common.py

@@ -21,6 +21,13 @@
     "int64": np.int64,
     "default": GLOBAL_NP_FLOAT_PRECISION,
 }
+RESERVED_PRECISON_DICT = {
+    np.float16: "float16",
+    np.float32: "float32",
+    np.float64: "float64",
+    np.int32: "int32",
+    np.int64: "int64",
+}
 DEFAULT_PRECISION = "float64"
 
 

deepmd/dpmodel/model/make_model.py

@@ -1,12 +1,17 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 from typing import (
     Dict,
+    List,
     Optional,
+    Tuple,
 )
 
 import numpy as np
 
 from deepmd.dpmodel.common import (
+    GLOBAL_NP_FLOAT_PRECISION,
+    PRECISION_DICT,
+    RESERVED_PRECISON_DICT,
     NativeOP,
 )
 from deepmd.dpmodel.output_def import (
@@ -59,6 +64,9 @@
                 *args,
                 **kwargs,
             )
+            self.precision_dict = PRECISION_DICT
+            self.reverse_precision_dict = RESERVED_PRECISON_DICT
+            self.global_np_float_precision = GLOBAL_NP_FLOAT_PRECISION
 
         def model_output_def(self):
             """Get the output def for the model."""
@@ -115,15 +123,19 @@
 
             """
             nframes, nloc = atype.shape[:2]
-            if box is not None:
+            cc, bb, fp, ap, input_prec = self.input_type_cast(
+                coord, box=box, fparam=fparam, aparam=aparam
+            )
+            del coord, box, fparam, aparam
+            if bb is not None:
                 coord_normalized = normalize_coord(
-                    coord.reshape(nframes, nloc, 3),
-                    box.reshape(nframes, 3, 3),
+                    cc.reshape(nframes, nloc, 3),
+                    bb.reshape(nframes, 3, 3),
                 )
             else:
-                coord_normalized = coord.copy()
+                coord_normalized = cc.copy()
             extended_coord, extended_atype, mapping = extend_coord_with_ghosts(
-                coord_normalized, atype, box, self.get_rcut()
+                coord_normalized, atype, bb, self.get_rcut()
             )
             nlist = build_neighbor_list(
                 extended_coord,
@@ -139,8 +151,8 @@
                 extended_atype,
                 nlist,
                 mapping,
-                fparam=fparam,
-                aparam=aparam,
+                fparam=fp,
+                aparam=ap,
                 do_atomic_virial=do_atomic_virial,
             )
             model_predict = communicate_extended_output(
@@ -149,6 +161,7 @@
                 mapping,
                 do_atomic_virial=do_atomic_virial,
             )
+            model_predict = self.output_type_cast(model_predict, input_prec)
             return model_predict
 
         def call_lower(
@@ -192,22 +205,82 @@
             nframes, nall = extended_atype.shape[:2]
             extended_coord = extended_coord.reshape(nframes, -1, 3)
             nlist = self.format_nlist(extended_coord, extended_atype, nlist)
+            cc_ext, _, fp, ap, input_prec = self.input_type_cast(
+                extended_coord, fparam=fparam, aparam=aparam
+            )
+            del extended_coord, fparam, aparam
             atomic_ret = self.forward_atomic(
-                extended_coord,
+                cc_ext,
                 extended_atype,
                 nlist,
                 mapping=mapping,
-                fparam=fparam,
-                aparam=aparam,
+                fparam=fp,
+                aparam=ap,
             )
             model_predict = fit_output_to_model_output(
                 atomic_ret,
                 self.fitting_output_def(),
-                extended_coord,
+                cc_ext,
                 do_atomic_virial=do_atomic_virial,
             )
+            model_predict = self.output_type_cast(model_predict, input_prec)
             return model_predict
 
+        def input_type_cast(
+            self,
+            coord: np.ndarray,
+            box: Optional[np.ndarray] = None,
+            fparam: Optional[np.ndarray] = None,
+            aparam: Optional[np.ndarray] = None,
+        ) -> Tuple[
+            np.ndarray,
+            Optional[np.ndarray],
+            Optional[np.ndarray],
+            Optional[np.ndarray],
+            str,
+        ]:
+            """Cast the input data to global float type."""
+            input_prec = self.reverse_precision_dict[
+                self.precision_dict[coord.dtype.name]
+            ]
+            ###
+            ### type checking would not pass jit, convert to coord prec anyway
+            ###
+            _lst: List[Optional[np.ndarray]] = [
+                vv.astype(coord.dtype) if vv is not None else None
+                for vv in [box, fparam, aparam]
+            ]
+            box, fparam, aparam = _lst
+            if (
+                input_prec
+                == self.reverse_precision_dict[self.global_np_float_precision]
+            ):
+                return coord, box, fparam, aparam, input_prec
+            else:
+                pp = self.global_np_float_precision
+                return (
+                    coord.astype(pp),
+                    box.astype(pp) if box is not None else None,
+                    fparam.astype(pp) if fparam is not None else None,
+                    aparam.astype(pp) if aparam is not None else None,
+                    input_prec,
+                )
+
+        def output_type_cast(
+            self,
+            model_ret: Dict[str, np.ndarray],
+            input_prec: str,
+        ) -> Dict[str, np.ndarray]:
+            """Convert the model output to the input prec."""
+            if (
+                input_prec
+                != self.reverse_precision_dict[self.global_np_float_precision]
+            ):
+                pp = self.precision_dict[input_prec]
+                for kk, vv in model_ret.items():
+                    model_ret[kk] = vv.astype(pp) if vv is not None else None
+            return model_ret
+
         def format_nlist(
             self,
             extended_coord: np.ndarray,

deepmd/pt/model/model/make_model.py

@@ -3,6 +3,7 @@
     Dict,
     List,
     Optional,
+    Tuple,
 )
 
 import torch
@@ -17,6 +18,11 @@
     communicate_extended_output,
     fit_output_to_model_output,
 )
+from deepmd.pt.utils.env import (
+    GLOBAL_PT_FLOAT_PRECISION,
+    PRECISION_DICT,
+    RESERVED_PRECISON_DICT,
+)
 from deepmd.pt.utils.nlist import (
     extend_input_and_build_neighbor_list,
     nlist_distinguish_types,
@@ -56,6 +62,9 @@
                 *args,
                 **kwargs,
             )
+            self.precision_dict = PRECISION_DICT
+            self.reverse_precision_dict = RESERVED_PRECISON_DICT
+            self.global_pt_float_precision = GLOBAL_PT_FLOAT_PRECISION
 
         def model_output_def(self):
             """Get the output def for the model."""
@@ -115,34 +124,39 @@
                 The keys are defined by the `ModelOutputDef`.
 
             """
+            cc, bb, fp, ap, input_prec = self.input_type_cast(
+                coord, box=box, fparam=fparam, aparam=aparam
+            )
+            del coord, box, fparam, aparam
             (
                 extended_coord,
                 extended_atype,
                 mapping,
                 nlist,
             ) = extend_input_and_build_neighbor_list(
-                coord,
+                cc,
                 atype,
                 self.get_rcut(),
                 self.get_sel(),
                 mixed_types=self.mixed_types(),
-                box=box,
+                box=bb,
             )
             model_predict_lower = self.forward_common_lower(
                 extended_coord,
                 extended_atype,
                 nlist,
                 mapping,
                 do_atomic_virial=do_atomic_virial,
-                fparam=fparam,
-                aparam=aparam,
+                fparam=fp,
+                aparam=ap,
             )
             model_predict = communicate_extended_output(
                 model_predict_lower,
                 self.model_output_def(),
                 mapping,
                 do_atomic_virial=do_atomic_virial,
             )
+            model_predict = self.output_type_cast(model_predict, input_prec)
             return model_predict
 
         def forward_common_lower(
@@ -186,22 +200,87 @@
             nframes, nall = extended_atype.shape[:2]
             extended_coord = extended_coord.view(nframes, -1, 3)
             nlist = self.format_nlist(extended_coord, extended_atype, nlist)
+            cc_ext, _, fp, ap, input_prec = self.input_type_cast(
+                extended_coord, fparam=fparam, aparam=aparam
+            )
+            del extended_coord, fparam, aparam
             atomic_ret = self.forward_atomic(
-                extended_coord,
+                cc_ext,
                 extended_atype,
                 nlist,
                 mapping=mapping,
-                fparam=fparam,
-                aparam=aparam,
+                fparam=fp,
+                aparam=ap,
             )
             model_predict = fit_output_to_model_output(
                 atomic_ret,
                 self.fitting_output_def(),
-                extended_coord,
+                cc_ext,
                 do_atomic_virial=do_atomic_virial,
             )
+            model_predict = self.output_type_cast(model_predict, input_prec)
             return model_predict
 
+        def input_type_cast(
+            self,
+            coord: torch.Tensor,
+            box: Optional[torch.Tensor] = None,
+            fparam: Optional[torch.Tensor] = None,
+            aparam: Optional[torch.Tensor] = None,
+        ) -> Tuple[
+            torch.Tensor,
+            Optional[torch.Tensor],
+            Optional[torch.Tensor],
+            Optional[torch.Tensor],
+            str,
+        ]:
+            """Cast the input data to global float type."""
+            input_prec = self.reverse_precision_dict[coord.dtype]
+            ###
+            ### type checking would not pass jit, convert to coord prec anyway
+            ###
+            # for vv, kk in zip([fparam, aparam], ["frame", "atomic"]):
+            #     if vv is not None and self.reverse_precision_dict[vv.dtype] != input_prec:
+            #         log.warning(
+            #           f"type of {kk} parameter {self.reverse_precision_dict[vv.dtype]}"
+            #           " does not match"
+            #           f" that of the coordinate {input_prec}"
+            #         )
+            _lst: List[Optional[torch.Tensor]] = [
+                vv.to(coord.dtype) if vv is not None else None
+                for vv in [box, fparam, aparam]
+            ]
+            box, fparam, aparam = _lst
+            if (
+                input_prec
+                == self.reverse_precision_dict[self.global_pt_float_precision]
+            ):
+                return coord, box, fparam, aparam, input_prec
+            else:
+                pp = self.global_pt_float_precision
+                return (
+                    coord.to(pp),
+                    box.to(pp) if box is not None else None,
+                    fparam.to(pp) if fparam is not None else None,
+                    aparam.to(pp) if aparam is not None else None,
+                    input_prec,
+                )
+
+        def output_type_cast(
+            self,
+            model_ret: Dict[str, torch.Tensor],
+            input_prec: str,
+        ) -> Dict[str, torch.Tensor]:
+            """Convert the model output to the input prec."""
+            if (
+                input_prec
+                != self.reverse_precision_dict[self.global_pt_float_precision]
+            ):
+                pp = self.precision_dict[input_prec]
+                for kk, vv in model_ret.items():
+                    model_ret[kk] = vv.to(pp) if vv is not None else None
+            return model_ret
+
         @torch.jit.export
         def format_nlist(
             self,

deepmd/pt/utils/exclude_mask.py

@@ -22,6 +22,13 @@
         exclude_types: List[int] = [],
     ):
         super().__init__()
+        self.reinit(ntypes, exclude_types)
+
+    def reinit(
+        self,
+        ntypes: int,
+        exclude_types: List[int] = [],
+    ):
         self.ntypes = ntypes
         self.exclude_types = exclude_types
         self.type_mask = np.array(
@@ -62,6 +69,13 @@
         exclude_types: List[Tuple[int, int]] = [],
     ):
         super().__init__()
+        self.reinit(ntypes, exclude_types)
+
+    def reinit(
+        self,
+        ntypes: int,
+        exclude_types: List[Tuple[int, int]] = [],
+    ):
         self.ntypes = ntypes
         self._exclude_types: Set[Tuple[int, int]] = set()
         for tt in exclude_types: