Implement native files protocols for optimization and torsiondrive

qcelemental/models/v2/atomic.py

@@ -8,7 +8,7 @@
 from ...util import provenance_stamp
 from .basemodels import ExtendedConfigDict, ProtoModel, check_convertible_version, qcschema_draft
 from .basis_set import BasisSet
-from .common_models import DriverEnum, Model, Provenance
+from .common_models import DriverEnum, Model, Provenance, NativeFilesProtocolEnum
 from .molecule import Molecule
 from .types import Array
 
@@ -640,14 +640,6 @@ def allows(self, policy: str):
         return self.policies.get(policy, self.default_policy)
 
 
-class NativeFilesProtocolEnum(str, Enum):
-    r"""CMS program files to keep from a computation."""
-
-    all = "all"
-    input = "input"
-    none = "none"
-
-
 class AtomicProtocols(ProtoModel):
     r"""Protocols regarding the manipulation of computational result data."""
 
@@ -822,7 +814,7 @@ class AtomicResult(ProtoModel):
         description="The primary logging output of the program, whether natively standard output or a file. Presence vs. absence (or null-ness?) configurable by protocol.",
     )
     stderr: Optional[str] = Field(None, description="The standard error of the program execution.")
-    native_files: Dict[str, Any] = Field({}, description="DSL files.")
+    native_files: Dict[str, Any] = Field({}, description="Other program-specific files returned from the computation.")
 
     success: Literal[True] = Field(
         True, description="The success of program execution. If False, other fields may be blank."

qcelemental/models/v2/common_models.py

@@ -63,3 +63,11 @@ def derivative_int(self):
             return 0
         else:
             return egh.index(self)
+
+
+class NativeFilesProtocolEnum(str, Enum):
+    r"""Program-specific to keep from a computation."""
+
+    all = "all"
+    input = "input"
+    none = "none"

qcelemental/models/v2/optimization.py

@@ -12,7 +12,7 @@
 from ...util import provenance_stamp
 from .atomic import AtomicProperties, AtomicResult, AtomicSpecification
 from .basemodels import ExtendedConfigDict, ProtoModel, check_convertible_version
-from .common_models import Provenance
+from .common_models import Provenance, NativeFilesProtocolEnum
 from .molecule import Molecule
 from .types import Array
 
@@ -38,20 +38,24 @@ class TrajectoryProtocolEnum(str, Enum):
 
 class OptimizationProtocols(ProtoModel):
     """
-    Protocols regarding the manipulation of a Optimization output data.
+    Protocols regarding the manipulation of Optimization output data.
     """
 
     schema_name: Literal["qcschema_optimization_protocols"] = "qcschema_optimization_protocols"
     trajectory_results: TrajectoryProtocolEnum = Field(
         TrajectoryProtocolEnum.none, description=str(TrajectoryProtocolEnum.__doc__)
     )
+    native_files: NativeFilesProtocolEnum = Field(
+        NativeFilesProtocolEnum.none,
+        description="Policies for keeping processed files from the computation",
+    )
 
     model_config = ExtendedConfigDict(force_skip_defaults=True)
 
     def convert_v(
         self, target_version: int, /
     ) -> Union["qcelemental.models.v1.OptimizationProtocols", "qcelemental.models.v2.OptimizationProtocols"]:
-        """Convert to instance of particular QCSchema version."""
+        """Convert to an instance of a particular QCSchema version."""
         import qcelemental as qcel
 
         if check_convertible_version(target_version, error="OptimizationProtocols") == "self":
@@ -61,6 +65,7 @@ def convert_v(
         if target_version == 1:
             # serialization is compact, so use model to assure value
             dself.pop("trajectory_results", None)
+            dself.pop("native_files", None)
             dself["trajectory"] = self.trajectory_results.value
 
             self_vN = qcel.models.v1.OptimizationProtocols(**dself)
@@ -277,7 +282,7 @@ class OptimizationResult(ProtoModel):
     provenance: Provenance = Field(..., description=str(Provenance.__doc__))
 
     # native_files placeholder for when any opt programs supply extra files or need an input file. no protocol at present
-    native_files: Dict[str, Any] = Field({}, description="DSL files.")
+    native_files: Dict[str, Any] = Field({}, description="Other program-specific files returned from the computation.")
 
     properties: OptimizationProperties = Field(..., description=str(OptimizationProperties.__doc__))
 
@@ -309,6 +314,32 @@ def _trajectory_protocol(cls, v, info):
 
         return v
 
+    @field_validator("native_files")
+    @classmethod
+    def _native_file_protocol(cls, value, info):
+        # Do not propagate validation errors
+        if "input_data" not in info.data:
+            raise ValueError("Input_data was not properly formed.")
+
+        ancp = info.data["input_data"].specification.protocols.native_files
+        if ancp == "all":
+            return value
+        elif ancp == "none":
+            return {}
+        elif ancp == "input":
+            return_keep = ["input"]
+            if value is None:
+                files = {}
+            else:
+                files = value.copy()
+        else:
+            raise ValueError(f"Protocol `native_files:{ancp}` is not understood")
+
+        ret = {}
+        for rk in return_keep:
+            ret[rk] = files.get(rk, None)
+        return ret
+
     def convert_v(
         self, target_version: int, /
     ) -> Union["qcelemental.models.v1.OptimizationResult", "qcelemental.models.v2.OptimizationResult"]:

qcelemental/models/v2/torsion_drive.py

@@ -3,18 +3,15 @@
 
 from pydantic import Field, conlist, field_validator
 
-from ...util import provenance_stamp
 from .basemodels import ExtendedConfigDict, ProtoModel, check_convertible_version
-from .common_models import DriverEnum, Provenance
+from .common_models import DriverEnum, Provenance, NativeFilesProtocolEnum
 from .molecule import Molecule
 from .optimization import OptimizationProperties, OptimizationResult, OptimizationSpecification
-from .types import Array
+from ...util import provenance_stamp
 
 if TYPE_CHECKING:
     import qcelemental
 
-    from .common_models import ReprArgs
-
 
 # ====  Protocols  ==============================================================
 
@@ -39,6 +36,11 @@ class TorsionDriveProtocols(ProtoModel):
         ScanResultsProtocolEnum.none, description=str(ScanResultsProtocolEnum.__doc__)
     )
 
+    native_files: NativeFilesProtocolEnum = Field(
+        NativeFilesProtocolEnum.none,
+        description="Policies for keeping processed files from the computation",
+    )
+
     model_config = ExtendedConfigDict(force_skip_defaults=True)
 
 
@@ -229,8 +231,8 @@ class TorsionDriveResult(ProtoModel):
     stdout: Optional[str] = Field(None, description="The standard output of the program.")
     stderr: Optional[str] = Field(None, description="The standard error of the program.")
 
-    # native_files placeholder for when any td programs supply extra files or need an input file. no protocol at present
-    native_files: Dict[str, Any] = Field({}, description="DSL files.")
+    # native_files placeholder for when any opt programs supply extra files or need an input file. no protocol at present
+    native_files: Dict[str, Any] = Field({}, description="Other program-specific files returned from the computation.")
 
     properties: TorsionDriveProperties = Field(..., description=str(TorsionDriveProperties.__doc__))
 
@@ -269,6 +271,32 @@ def _scan_protocol(cls, v, info):
 
         return v
 
+    @field_validator("native_files")
+    @classmethod
+    def _native_file_protocol(cls, value, info):
+        # Do not propagate validation errors
+        if "input_data" not in info.data:
+            raise ValueError("Input_data was not properly formed.")
+
+        ancp = info.data["input_data"].specification.protocols.native_files
+        if ancp == "all":
+            return value
+        elif ancp == "none":
+            return {}
+        elif ancp == "input":
+            return_keep = ["input"]
+            if value is None:
+                files = {}
+            else:
+                files = value.copy()
+        else:
+            raise ValueError(f"Protocol `native_files:{ancp}` is not understood")
+
+        ret = {}
+        for rk in return_keep:
+            ret[rk] = files.get(rk, None)
+        return ret
+
     def convert_v(
         self, target_version: int, /
     ) -> Union["qcelemental.models.v1.TorsionDriveResult", "qcelemental.models.v2.TorsionDriveResult"]: