fix CP2K TaskDocument (#900)

* bugfix for task document schema

* lint

src/atomate2/cp2k/schemas/calc_types/utils.py

@@ -4,7 +4,7 @@
 from pathlib import Path
 
 from monty.serialization import loadfn
-from pymatgen.io.cp2k.inputs import Cp2kInput, Keyword, KeywordList
+from pymatgen.io.cp2k.inputs import Keyword, KeywordList
 
 from atomate2.cp2k.schemas.calc_types import CalcType, RunType, TaskType
 
@@ -71,7 +71,7 @@ def task_type(inputs: dict) -> TaskType:
     """
     calc_type = []
     cp2k_run_type = inputs.get("cp2k_global", {}).get("Run_type", "")
-    ci = Cp2kInput.from_dict(inputs["cp2k_input"])
+    ci = inputs["cp2k_input"]
 
     if cp2k_run_type.upper() in (
         "ENERGY",

src/atomate2/cp2k/schemas/calculation.py

@@ -16,7 +16,7 @@
 from pymatgen.electronic_structure.bandstructure import BandStructure
 from pymatgen.electronic_structure.dos import Dos
 from pymatgen.io.common import VolumetricData
-from pymatgen.io.cp2k.inputs import BasisFile, DataFile, PotentialFile
+from pymatgen.io.cp2k.inputs import BasisFile, Cp2kInput, DataFile, PotentialFile
 from pymatgen.io.cp2k.outputs import Cp2kOutput, parse_energy_file
 from typing_extensions import Self
 
@@ -71,7 +71,9 @@ class CalculationInput(BaseModel):
         None, description="Description of parameters used for each atomic kind"
     )
 
-    cp2k_input: dict = Field(None, description="The cp2k input used for this task")
+    cp2k_input: Union[dict, Cp2kInput] = Field(
+        None, description="The cp2k input used for this task"
+    )
 
     dft: dict = Field(
         None,
@@ -170,7 +172,7 @@ class CalculationOutput(BaseModel):
     ionic_steps: list[dict[str, Any]] = Field(
         None, description="Energy, forces, and structure for each ionic step"
     )
-    locpot: dict[int, list[float]] = Field(
+    locpot: Union[dict[int, list[float]], None] = Field(
         None, description="Average of the local potential along the crystal axes"
     )
     run_stats: RunStatistics = Field(