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

deepmd/dpmodel/common.py

@@ -6,7 +6,8 @@
 
 import numpy as np
 
-from deepmd.common import (
+from deepmd.env import (
+    GLOBAL_ENER_FLOAT_PRECISION,
     GLOBAL_NP_FLOAT_PRECISION,
 )
 
@@ -21,6 +22,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"
 
 
@@ -35,3 +43,13 @@ def call(self, *args, **kwargs):
     def __call__(self, *args, **kwargs):
         """Forward pass in NumPy implementation."""
         return self.call(*args, **kwargs)
+
+
+__all__ = [
+    "GLOBAL_NP_FLOAT_PRECISION",
+    "GLOBAL_ENER_FLOAT_PRECISION",
+    "PRECISION_DICT",
+    "RESERVED_PRECISON_DICT",
+    "DEFAULT_PRECISION",
+    "NativeOP",
+]

deepmd/dpmodel/model/make_model.py

@@ -1,17 +1,25 @@
 # 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_ENER_FLOAT_PRECISION,
+    GLOBAL_NP_FLOAT_PRECISION,
+    PRECISION_DICT,
+    RESERVED_PRECISON_DICT,
     NativeOP,
 )
 from deepmd.dpmodel.output_def import (
     ModelOutputDef,
     OutputVariableCategory,
+    OutputVariableOperation,
+    check_operation_applied,
 )
 from deepmd.dpmodel.utils import (
     build_neighbor_list,
@@ -59,6 +67,10 @@
                 *args,
                 **kwargs,
             )
+            self.precision_dict = PRECISION_DICT
+            self.reverse_precision_dict = RESERVED_PRECISON_DICT
+            self.global_np_float_precision = GLOBAL_NP_FLOAT_PRECISION
+            self.global_ener_float_precision = GLOBAL_ENER_FLOAT_PRECISION
 
         def model_output_def(self):
             """Get the output def for the model."""
@@ -115,15 +127,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 +155,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 +165,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 +209,95 @@
             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."""
+            do_cast = (
+                input_prec
+                != self.reverse_precision_dict[self.global_np_float_precision]
+            )
+            pp = self.precision_dict[input_prec]
+            odef = self.model_output_def()
+            for kk in odef.keys():
+                if kk not in model_ret.keys():
+                    # do not return energy_derv_c if not do_atomic_virial
+                    continue
+                if check_operation_applied(odef[kk], OutputVariableOperation.REDU):
+                    model_ret[kk] = (
+                        model_ret[kk].astype(self.global_ener_float_precision)
+                        if model_ret[kk] is not None
+                        else None
+                    )
+                elif do_cast:
+                    model_ret[kk] = (
+                        model_ret[kk].astype(pp) if model_ret[kk] is not None else None
+                    )
+            return model_ret
+
         def format_nlist(
             self,
             extended_coord: np.ndarray,

deepmd/dpmodel/model/transform_output.py

@@ -5,6 +5,9 @@
 
 import numpy as np
 
+from deepmd.dpmodel.common import (
+    GLOBAL_ENER_FLOAT_PRECISION,
+)
 from deepmd.dpmodel.output_def import (
     FittingOutputDef,
     ModelOutputDef,
@@ -30,7 +33,10 @@
         atom_axis = -(len(shap) + 1)
         if vdef.reduciable:
             kk_redu = get_reduce_name(kk)
-            model_ret[kk_redu] = np.sum(vv, axis=atom_axis)
+            # cast to energy prec brefore reduction
+            model_ret[kk_redu] = np.sum(
+                vv.astype(GLOBAL_ENER_FLOAT_PRECISION), axis=atom_axis
+            )
             if vdef.r_differentiable:
                 kk_derv_r, kk_derv_c = get_deriv_name(kk)
                 # name-holders