protocol for retaining DSL files (#275)

* Update results.py * Update results.py * Update qcelemental/models/results.py * Update results.py * Update test_model_results.py * Update test_model_results.py * Update results.py
MolSSI · Nov 18, 2021 · f7dfea9 · f7dfea9
1 parent ebca0bb
commit f7dfea9
Show file tree

Hide file tree

Showing 2 changed files with 104 additions and 0 deletions.
diff --git a/qcelemental/models/results.py b/qcelemental/models/results.py
@@ -527,6 +527,14 @@ 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 AtomicResultProtocols(ProtoModel):
     r"""Protocols regarding the manipulation of computational result data."""
 
@@ -537,6 +545,10 @@ class AtomicResultProtocols(ProtoModel):
     error_correction: ErrorCorrectionProtocol = Field(
         default_factory=ErrorCorrectionProtocol, description="Policies for error correction"
     )
+    native_files: NativeFilesProtocolEnum = Field(
+        NativeFilesProtocolEnum.none,
+        description="Policies for keeping processed files from the computation",
+    )
 
     class Config:
         force_skip_defaults = True
@@ -608,6 +620,7 @@ class AtomicResult(AtomicInput):
         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: Optional[Dict[str, Any]] = Field(None, description="DSL files.")
 
     success: bool = Field(..., description="The success of program execution. If False, other fields may be blank.")
     error: Optional[ComputeError] = Field(None, description=str(ComputeError.__base_doc__))
@@ -718,6 +731,28 @@ def _stdout_protocol(cls, value, values):
         else:
             raise ValueError(f"Protocol `stdout:{outp}` is not understood")
 
+    @validator("native_files")
+    def _native_file_protocol(cls, value, values):
+
+        ancp = values["protocols"].native_files
+        if ancp == "all":
+            return value
+        elif ancp == "none":
+            return None
+        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
+
 
 class ResultProperties(AtomicResultProperties):
     def __init__(self, *args, **kwargs):

diff --git a/qcelemental/tests/test_model_results.py b/qcelemental/tests/test_model_results.py
@@ -127,6 +127,41 @@ def wavefunction_data_fixture(result_data_fixture):
     return result_data_fixture
 
 
+@pytest.fixture(scope="function")
+def native_data_fixture(result_data_fixture):
+    result_data_fixture["protocols"] = {"native_files": "all"}
+    result_data_fixture["native_files"] = {
+        "input": """
+echo
+geometry units bohr
+H1                    0.000000000000     0.000000000000    -0.650000000000
+H1                    0.000000000000     0.000000000000     0.650000000000
+end
+charge 0.0
+memory 1669668536 double
+task scf energy
+""",
+        "DIPOL": """
+        0.0000000000        0.0000000000        0.0000000000
+        0.0000000000        0.0000000000       -0.7855589278
+""",
+        "gms.dat": """
+ $DATA
+Methylene...3-B-1 state...ROHF/STO-2G
+CNV      2
+
+HYDROGEN    1.0      0.8288400000      0.0000000000      0.6060939022
+   STO     2
+
+CARBON      6.0      0.0000000000      0.0000000000     -0.1018060978
+   STO     2
+
+ $END
+""",
+    }
+    return result_data_fixture
+
+
 @pytest.fixture(scope="function")
 def optimization_data_fixture(result_data_fixture):
 
@@ -308,6 +343,40 @@ def test_wavefunction_protocols(protocol, restricted, provided, expected, wavefu
         assert wfn.wavefunction.dict().keys() == expected_keys
 
 
+@pytest.mark.parametrize(
+    "protocol, provided, expected",
+    [
+        ("none", ["input", "gms.dat", "DIPOL"], []),
+        (None, ["input", "gms.dat", "DIPOL"], []),
+        ("input", ["input", "gms.dat", "DIPOL"], ["input"]),
+        ("all", ["input", "gms.dat", "DIPOL"], ["input", "gms.dat", "DIPOL"]),
+        ("all", ["DIPOL"], ["DIPOL"]),
+        ("input", ["gms.dat"], ["input"]),
+    ],
+)
+def test_native_protocols(protocol, provided, expected, native_data_fixture, request):
+
+    native_data = native_data_fixture["native_files"]
+
+    if protocol is None:
+        native_data_fixture.pop("protocols")
+    else:
+        native_data_fixture["protocols"]["native_files"] = protocol
+
+    for name in list(native_data.keys()):
+        if name not in provided:
+            native_data.pop(name)
+
+    wfn = qcel.models.AtomicResult(**native_data_fixture)
+    drop_qcsk(wfn, request.node.name)
+
+    if len(expected) == 0:
+        assert wfn.native_files is None
+    else:
+        expected_keys = set(expected)
+        assert wfn.native_files.keys() == expected_keys
+
+
 @pytest.mark.parametrize(
     "keep, indices",
     [(None, [0, 1, 2, 3, 4]), ("all", [0, 1, 2, 3, 4]), ("initial_and_final", [0, 4]), ("final", [4]), ("none", [])],