diff --git a/deepmd/dpmodel/output_def.py b/deepmd/dpmodel/output_def.py index 8b190ed5de..fac24534eb 100644 --- a/deepmd/dpmodel/output_def.py +++ b/deepmd/dpmodel/output_def.py @@ -172,8 +172,12 @@ class OutputVariableDef: are differentiable. Virial, the transposed negative gradient with cell tensor times cell tensor, will be calculated, see eq 40 JCP 159, 054801 (2023). + atomic : bool + If the variable is defined for each atom. category : int The category of the output variable. + hessian : bool + If hessian is requred """ def __init__( @@ -185,6 +189,7 @@ def __init__( c_differentiable: bool = False, atomic: bool = True, category: int = OutputVariableCategory.OUT.value, + r_hessian: bool = False, ): self.name = name self.shape = list(shape) @@ -194,13 +199,15 @@ def __init__( self.c_differentiable = c_differentiable if self.c_differentiable and not self.r_differentiable: raise ValueError("c differentiable requires r_differentiable") - if not self.reduciable and self.r_differentiable: - raise ValueError("only reduciable variable are r differentiable") - if not self.reduciable and self.c_differentiable: - raise ValueError("only reduciable variable are c differentiable") if self.reduciable and not self.atomic: raise ValueError("a reduciable variable should be atomic") self.category = category + self.r_hessian = r_hessian + if self.r_hessian: + if not self.reduciable: + raise ValueError("only reduciable variable can calculate hessian") + if not self.r_differentiable: + raise ValueError("only r_differentiable variable can calculate hessian") class FittingOutputDef: @@ -257,6 +264,7 @@ def __init__( self.def_outp = fit_defs self.def_redu = do_reduce(self.def_outp.get_data()) self.def_derv_r, self.def_derv_c = do_derivative(self.def_outp.get_data()) + self.def_hess_r, _ = do_derivative(self.def_derv_r) self.def_derv_c_redu = do_reduce(self.def_derv_c) self.var_defs: Dict[str, OutputVariableDef] = {} for ii in [ @@ -265,6 +273,7 @@ def __init__( self.def_derv_c, self.def_derv_r, self.def_derv_c_redu, + self.def_hess_r, ]: self.var_defs.update(ii) @@ -292,6 +301,9 @@ def keys_redu(self): def keys_derv_r(self): return self.def_derv_r.keys() + def keys_hess_r(self): + return self.def_hess_r.keys() + def keys_derv_c(self): return self.def_derv_c.keys() @@ -392,7 +404,9 @@ def do_derivative( rkr, vv.shape + [3], # noqa: RUF005 reduciable=False, - r_differentiable=False, + r_differentiable=( + vv.r_hessian and vv.category == OutputVariableCategory.OUT.value + ), c_differentiable=False, atomic=True, category=apply_operation(vv, OutputVariableOperation.DERV_R), diff --git a/source/tests/common/dpmodel/test_output_def.py b/source/tests/common/dpmodel/test_output_def.py index 9e794e0f32..272082513c 100644 --- a/source/tests/common/dpmodel/test_output_def.py +++ b/source/tests/common/dpmodel/test_output_def.py @@ -44,6 +44,16 @@ def test_model_output_def(self): r_differentiable=True, c_differentiable=True, atomic=True, + r_hessian=False, + ), + OutputVariableDef( + "energy2", + [1], + reduciable=True, + r_differentiable=True, + c_differentiable=True, + atomic=True, + r_hessian=True, ), OutputVariableDef( "dos", @@ -64,26 +74,33 @@ def test_model_output_def(self): ] # fitting definition fd = FittingOutputDef(defs) - expected_keys = ["energy", "dos", "foo"] + expected_keys = ["energy", "energy2", "dos", "foo"] self.assertEqual( set(expected_keys), set(fd.keys()), ) # shape self.assertEqual(fd["energy"].shape, [1]) + self.assertEqual(fd["energy2"].shape, [1]) self.assertEqual(fd["dos"].shape, [10]) self.assertEqual(fd["foo"].shape, [3]) # atomic self.assertEqual(fd["energy"].atomic, True) + self.assertEqual(fd["energy2"].atomic, True) self.assertEqual(fd["dos"].atomic, True) self.assertEqual(fd["foo"].atomic, True) # reduce self.assertEqual(fd["energy"].reduciable, True) + self.assertEqual(fd["energy2"].reduciable, True) self.assertEqual(fd["dos"].reduciable, True) self.assertEqual(fd["foo"].reduciable, False) # derivative self.assertEqual(fd["energy"].r_differentiable, True) self.assertEqual(fd["energy"].c_differentiable, True) + self.assertEqual(fd["energy"].r_hessian, False) + self.assertEqual(fd["energy2"].r_differentiable, True) + self.assertEqual(fd["energy2"].c_differentiable, True) + self.assertEqual(fd["energy2"].r_hessian, True) self.assertEqual(fd["dos"].r_differentiable, False) self.assertEqual(fd["foo"].r_differentiable, False) self.assertEqual(fd["dos"].c_differentiable, False) @@ -92,12 +109,18 @@ def test_model_output_def(self): md = ModelOutputDef(fd) expected_keys = [ "energy", + "energy2", "dos", "foo", "energy_redu", "energy_derv_r", "energy_derv_c", "energy_derv_c_redu", + "energy2_redu", + "energy2_derv_r", + "energy2_derv_r_derv_r", + "energy2_derv_c", + "energy2_derv_c_redu", "dos_redu", ] self.assertEqual( @@ -108,33 +131,51 @@ def test_model_output_def(self): self.assertEqual(md[kk].name, kk) # reduce self.assertEqual(md["energy"].reduciable, True) + self.assertEqual(md["energy2"].reduciable, True) self.assertEqual(md["dos"].reduciable, True) self.assertEqual(md["foo"].reduciable, False) # derivative self.assertEqual(md["energy"].r_differentiable, True) self.assertEqual(md["energy"].c_differentiable, True) + self.assertEqual(md["energy"].r_hessian, False) + self.assertEqual(md["energy2"].r_differentiable, True) + self.assertEqual(md["energy2"].c_differentiable, True) + self.assertEqual(md["energy2"].r_hessian, True) self.assertEqual(md["dos"].r_differentiable, False) self.assertEqual(md["foo"].r_differentiable, False) self.assertEqual(md["dos"].c_differentiable, False) self.assertEqual(md["foo"].c_differentiable, False) # shape self.assertEqual(md["energy"].shape, [1]) + self.assertEqual(md["energy2"].shape, [1]) self.assertEqual(md["dos"].shape, [10]) self.assertEqual(md["foo"].shape, [3]) self.assertEqual(md["energy_redu"].shape, [1]) self.assertEqual(md["energy_derv_r"].shape, [1, 3]) self.assertEqual(md["energy_derv_c"].shape, [1, 9]) self.assertEqual(md["energy_derv_c_redu"].shape, [1, 9]) + self.assertEqual(md["energy2_redu"].shape, [1]) + self.assertEqual(md["energy2_derv_r"].shape, [1, 3]) + self.assertEqual(md["energy2_derv_c"].shape, [1, 9]) + self.assertEqual(md["energy2_derv_c_redu"].shape, [1, 9]) + self.assertEqual(md["energy2_derv_r_derv_r"].shape, [1, 3, 3]) # atomic self.assertEqual(md["energy"].atomic, True) + self.assertEqual(md["energy2"].atomic, True) self.assertEqual(md["dos"].atomic, True) self.assertEqual(md["foo"].atomic, True) self.assertEqual(md["energy_redu"].atomic, False) self.assertEqual(md["energy_derv_r"].atomic, True) self.assertEqual(md["energy_derv_c"].atomic, True) self.assertEqual(md["energy_derv_c_redu"].atomic, False) + self.assertEqual(md["energy2_redu"].atomic, False) + self.assertEqual(md["energy2_derv_r"].atomic, True) + self.assertEqual(md["energy2_derv_c"].atomic, True) + self.assertEqual(md["energy2_derv_c_redu"].atomic, False) + self.assertEqual(md["energy2_derv_r_derv_r"].atomic, True) # category self.assertEqual(md["energy"].category, OutputVariableCategory.OUT) + self.assertEqual(md["energy2"].category, OutputVariableCategory.OUT) self.assertEqual(md["dos"].category, OutputVariableCategory.OUT) self.assertEqual(md["foo"].category, OutputVariableCategory.OUT) self.assertEqual(md["energy_redu"].category, OutputVariableCategory.REDU) @@ -143,101 +184,159 @@ def test_model_output_def(self): self.assertEqual( md["energy_derv_c_redu"].category, OutputVariableCategory.DERV_C_REDU ) - # flag - self.assertEqual(md["energy"].category & OutputVariableOperation.REDU, 0) - self.assertEqual(md["energy"].category & OutputVariableOperation.DERV_R, 0) - self.assertEqual(md["energy"].category & OutputVariableOperation.DERV_C, 0) - self.assertEqual(md["dos"].category & OutputVariableOperation.REDU, 0) - self.assertEqual(md["dos"].category & OutputVariableOperation.DERV_R, 0) - self.assertEqual(md["dos"].category & OutputVariableOperation.DERV_C, 0) - self.assertEqual(md["foo"].category & OutputVariableOperation.REDU, 0) - self.assertEqual(md["foo"].category & OutputVariableOperation.DERV_R, 0) - self.assertEqual(md["foo"].category & OutputVariableOperation.DERV_C, 0) - self.assertEqual( - md["energy_redu"].category & OutputVariableOperation.REDU, - OutputVariableOperation.REDU, - ) - self.assertEqual(md["energy_redu"].category & OutputVariableOperation.DERV_R, 0) - self.assertEqual(md["energy_redu"].category & OutputVariableOperation.DERV_C, 0) - self.assertEqual(md["energy_derv_r"].category & OutputVariableOperation.REDU, 0) + self.assertEqual(md["energy2_redu"].category, OutputVariableCategory.REDU) + self.assertEqual(md["energy2_derv_r"].category, OutputVariableCategory.DERV_R) + self.assertEqual(md["energy2_derv_c"].category, OutputVariableCategory.DERV_C) self.assertEqual( - md["energy_derv_r"].category & OutputVariableOperation.DERV_R, - OutputVariableOperation.DERV_R, + md["energy2_derv_c_redu"].category, OutputVariableCategory.DERV_C_REDU ) self.assertEqual( - md["energy_derv_r"].category & OutputVariableOperation.DERV_C, 0 + md["energy2_derv_r_derv_r"].category, OutputVariableCategory.DERV_R_DERV_R ) - self.assertEqual(md["energy_derv_c"].category & OutputVariableOperation.REDU, 0) + # flag + OVO = OutputVariableOperation + self.assertEqual(md["energy"].category & OVO.REDU, 0) + self.assertEqual(md["energy"].category & OVO.DERV_R, 0) + self.assertEqual(md["energy"].category & OVO.DERV_C, 0) + self.assertEqual(md["energy2"].category & OVO.REDU, 0) + self.assertEqual(md["energy2"].category & OVO.DERV_R, 0) + self.assertEqual(md["energy2"].category & OVO.DERV_C, 0) + self.assertEqual(md["dos"].category & OVO.REDU, 0) + self.assertEqual(md["dos"].category & OVO.DERV_R, 0) + self.assertEqual(md["dos"].category & OVO.DERV_C, 0) + self.assertEqual(md["foo"].category & OVO.REDU, 0) + self.assertEqual(md["foo"].category & OVO.DERV_R, 0) + self.assertEqual(md["foo"].category & OVO.DERV_C, 0) + # flag: energy self.assertEqual( - md["energy_derv_c"].category & OutputVariableOperation.DERV_R, 0 + md["energy_redu"].category & OVO.REDU, + OVO.REDU, ) + self.assertEqual(md["energy_redu"].category & OVO.DERV_R, 0) + self.assertEqual(md["energy_redu"].category & OVO.DERV_C, 0) + self.assertEqual(md["energy_derv_r"].category & OVO.REDU, 0) self.assertEqual( - md["energy_derv_c"].category & OutputVariableOperation.DERV_C, - OutputVariableOperation.DERV_C, + md["energy_derv_r"].category & OVO.DERV_R, + OVO.DERV_R, ) + self.assertEqual(md["energy_derv_r"].category & OVO.DERV_C, 0) + self.assertEqual(md["energy_derv_c"].category & OVO.REDU, 0) + self.assertEqual(md["energy_derv_c"].category & OVO.DERV_R, 0) self.assertEqual( - md["energy_derv_c_redu"].category & OutputVariableOperation.REDU, - OutputVariableOperation.REDU, + md["energy_derv_c"].category & OVO.DERV_C, + OVO.DERV_C, ) self.assertEqual( - md["energy_derv_c_redu"].category & OutputVariableOperation.DERV_R, 0 + md["energy_derv_c_redu"].category & OVO.REDU, + OVO.REDU, ) + self.assertEqual(md["energy_derv_c_redu"].category & OVO.DERV_R, 0) self.assertEqual( - md["energy_derv_c_redu"].category & OutputVariableOperation.DERV_C, - OutputVariableOperation.DERV_C, + md["energy_derv_c_redu"].category & OVO.DERV_C, + OVO.DERV_C, ) - - # apply_operation + # flag: energy2 + kk = "energy2_redu" + self.assertEqual(md[kk].category & OVO.REDU, OVO.REDU) + self.assertEqual(md[kk].category & OVO.DERV_R, 0) + self.assertEqual(md[kk].category & OVO.DERV_C, 0) + self.assertEqual(md[kk].category & OVO._SEC_DERV_R, 0) + kk = "energy2_derv_r" + self.assertEqual(md[kk].category & OVO.REDU, 0) + self.assertEqual(md[kk].category & OVO.DERV_R, OVO.DERV_R) + self.assertEqual(md[kk].category & OVO.DERV_C, 0) + self.assertEqual(md[kk].category & OVO._SEC_DERV_R, 0) + kk = "energy2_derv_c" + self.assertEqual(md[kk].category & OVO.REDU, 0) + self.assertEqual(md[kk].category & OVO.DERV_R, 0) + self.assertEqual(md[kk].category & OVO.DERV_C, OVO.DERV_C) + self.assertEqual(md[kk].category & OVO._SEC_DERV_R, 0) + kk = "energy2_derv_c_redu" + self.assertEqual(md[kk].category & OVO.REDU, OVO.REDU) + self.assertEqual(md[kk].category & OVO.DERV_R, 0) + self.assertEqual(md[kk].category & OVO.DERV_C, OVO.DERV_C) + self.assertEqual(md[kk].category & OVO._SEC_DERV_R, 0) + kk = "energy2_derv_r_derv_r" + self.assertEqual(md[kk].category & OVO.REDU, 0) + self.assertEqual(md[kk].category & OVO.DERV_R, OVO.DERV_R) + self.assertEqual(md[kk].category & OVO.DERV_C, 0) + self.assertEqual(md[kk].category & OVO._SEC_DERV_R, OVO._SEC_DERV_R) + # apply_operation: energy self.assertEqual( - apply_operation(md["energy"], OutputVariableOperation.REDU), + apply_operation(md["energy"], OVO.REDU), md["energy_redu"].category, ) self.assertEqual( - apply_operation(md["energy"], OutputVariableOperation.DERV_R), + apply_operation(md["energy"], OVO.DERV_R), md["energy_derv_r"].category, ) self.assertEqual( - apply_operation(md["energy"], OutputVariableOperation.DERV_C), + apply_operation(md["energy"], OVO.DERV_C), md["energy_derv_c"].category, ) self.assertEqual( - apply_operation(md["energy_derv_c"], OutputVariableOperation.REDU), + apply_operation(md["energy_derv_c"], OVO.REDU), md["energy_derv_c_redu"].category, ) + # apply_operation: energy2 + self.assertEqual( + apply_operation(md["energy2"], OVO.REDU), + md["energy2_redu"].category, + ) + self.assertEqual( + apply_operation(md["energy2"], OVO.DERV_R), + md["energy2_derv_r"].category, + ) + self.assertEqual( + apply_operation(md["energy2"], OVO.DERV_C), + md["energy2_derv_c"].category, + ) + self.assertEqual( + apply_operation(md["energy2_derv_c"], OVO.REDU), + md["energy2_derv_c_redu"].category, + ) + self.assertEqual( + apply_operation(md["energy2_derv_r"], OVO.DERV_R), + md["energy2_derv_r_derv_r"].category, + ) + # raise ValueError + with self.assertRaises(ValueError): + apply_operation(md["energy_redu"], OVO.REDU) + with self.assertRaises(ValueError): + apply_operation(md["energy_derv_c"], OVO.DERV_C) + with self.assertRaises(ValueError): + apply_operation(md["energy_derv_c_redu"], OVO.REDU) # raise ValueError with self.assertRaises(ValueError): - apply_operation(md["energy_redu"], OutputVariableOperation.REDU) + apply_operation(md["energy2_redu"], OVO.REDU) with self.assertRaises(ValueError): - apply_operation(md["energy_derv_c"], OutputVariableOperation.DERV_C) + apply_operation(md["energy2_derv_c"], OVO.DERV_C) with self.assertRaises(ValueError): - apply_operation(md["energy_derv_c_redu"], OutputVariableOperation.REDU) + apply_operation(md["energy2_derv_c_redu"], OVO.REDU) + with self.assertRaises(ValueError): + apply_operation(md["energy2_derv_r_derv_r"], OVO.DERV_R) # hession - hession_cat = apply_operation( - md["energy_derv_r"], OutputVariableOperation.DERV_R - ) - self.assertEqual( - hession_cat & OutputVariableOperation.DERV_R, OutputVariableOperation.DERV_R - ) + hession_cat = apply_operation(md["energy_derv_r"], OVO.DERV_R) + self.assertEqual(hession_cat & OVO.DERV_R, OVO.DERV_R) self.assertEqual( - hession_cat & OutputVariableOperation._SEC_DERV_R, - OutputVariableOperation._SEC_DERV_R, + hession_cat & OVO._SEC_DERV_R, + OVO._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) + apply_operation(hession_vardef, OVO.DERV_R) - def test_raise_no_redu_deriv(self): - with self.assertRaises(ValueError) as context: - OutputVariableDef( - "energy", - [1], - reduciable=False, - r_differentiable=True, - c_differentiable=False, - ) + def test_no_raise_no_redu_deriv(self): + OutputVariableDef( + "energy", + [1], + reduciable=False, + r_differentiable=True, + c_differentiable=False, + ) def test_raise_requires_r_deriv(self): with self.assertRaises(ValueError) as context: @@ -253,6 +352,32 @@ def test_raise_redu_not_atomic(self): with self.assertRaises(ValueError) as context: (OutputVariableDef("energy", [1], reduciable=True, atomic=False),) + def test_hessian_not_reducible(self): + with self.assertRaises(ValueError) as context: + ( + OutputVariableDef( + "energy", + [1], + reduciable=False, + atomic=False, + r_differentiable=True, + r_hessian=True, + ), + ) + + def test_hessian_not_r_differentiable(self): + with self.assertRaises(ValueError) as context: + ( + OutputVariableDef( + "energy", + [1], + reduciable=True, + atomic=False, + r_differentiable=False, + r_hessian=True, + ), + ) + def test_model_decorator(self): nf = 2 nloc = 3