add category property to OutputVariableDef

deepmd/dpmodel/output_def.py

@@ -307,7 +307,23 @@
     -------
     int
         The new category of the variable definition.
+
+    Raises
+    ------
+    ValueError
+        If the operation has been applied to the variable definition,
+        and exceed the maximum limitation.
     """
+    if op == OutputVariableOperation.REDU or op == OutputVariableOperation.DERV_C:
+        if check_operation_applied(var_def, op):
+            raise ValueError(f"operation {op} has been applied")
+    elif op == OutputVariableOperation.DERV_R:
+        if check_operation_applied(var_def, OutputVariableOperation.DERV_R):
+            op = OutputVariableOperation.SEC_DERV_R
+            if check_operation_applied(var_def, OutputVariableOperation.SEC_DERV_R):
+                raise ValueError(f"operation {op} has been applied twice")
+    else:
+        raise ValueError(f"operation {op} not supported")
     return var_def.category | op.value
 
 
@@ -328,17 +344,6 @@
     bool
         True if the operation has been applied, False otherwise.
     """
-    if op in (OutputVariableOperation.DERV_REDU, OutputVariableOperation.DERV_C):
-        assert not check_operation_applied(var_def, op)
-    elif op == OutputVariableOperation.DERV_R:
-        if check_operation_applied(var_def, OutputVariableOperation.DERV_R):
-            op = OutputVariableOperation.SEC_DERV_R
-        else:
-            assert not check_operation_applied(
-                var_def, OutputVariableOperation.SEC_DERV_R
-            )
-    else:
-        raise ValueError(f"operation {op} not supported")
     return var_def.category & op.value == op.value
 
 

source/tests/common/dpmodel/test_output_def.py

@@ -17,6 +17,7 @@
 from deepmd.dpmodel.output_def import (
     OutputVariableCategory,
     OutputVariableOperation,
+    apply_operation,
     check_var,
 )
 
@@ -159,6 +160,48 @@ def test_model_output_def(self):
             OutputVariableOperation.DERV_C,
         )
 
+        # apply_operation
+        self.assertEqual(
+            apply_operation(md["energy"], OutputVariableOperation.REDU),
+            md["energy_redu"].category,
+        )
+        self.assertEqual(
+            apply_operation(md["energy"], OutputVariableOperation.DERV_R),
+            md["energy_derv_r"].category,
+        )
+        self.assertEqual(
+            apply_operation(md["energy"], OutputVariableOperation.DERV_C),
+            md["energy_derv_c"].category,
+        )
+        self.assertEqual(
+            apply_operation(md["energy_derv_c"], OutputVariableOperation.REDU),
+            md["energy_derv_c_redu"].category,
+        )
+        # raise ValueError
+        with self.assertRaises(ValueError):
+            apply_operation(md["energy_redu"], OutputVariableOperation.REDU)
+        with self.assertRaises(ValueError):
+            apply_operation(md["energy_derv_c"], OutputVariableOperation.DERV_C)
+        with self.assertRaises(ValueError):
+            apply_operation(md["energy_derv_c_redu"], OutputVariableOperation.REDU)
+        # hession
+        hession_cat = apply_operation(
+            md["energy_derv_r"], OutputVariableOperation.DERV_R
+        )
+        self.assertEqual(
+            hession_cat & OutputVariableOperation.DERV_R, OutputVariableOperation.DERV_R
+        )
+        self.assertEqual(
+            hession_cat & OutputVariableOperation.SEC_DERV_R,
+            OutputVariableOperation.SEC_DERV_R,
+        )
+        self.assertEqual(hession_cat, OutputVariableCategory.DERV_R_DERV_R)
+        hession_vardef = OutputVariableDef(
+            "energy_derv_r_derv_r", [1], False, False, category=hession_cat
+        )
+        with self.assertRaises(ValueError):
+            apply_operation(hession_vardef, OutputVariableOperation.DERV_R)
+
     def test_raise_no_redu_deriv(self):
         with self.assertRaises(ValueError) as context:
             (OutputVariableDef("energy", [1], False, True),)