1
"""
2
Utility functions for on-node procedures.
3
"""
4

5 4
import json
6

7 4
from qcelemental.models import ResultInput
8

9 4
from ..interface.models import Molecule, KVStore
10

11

12 4
def unpack_single_task_spec(storage, meta, molecules):
13
    """Transforms a metadata compute packet into an expanded
14
    QC Schema for multiple runs.
15

16
    Parameters
17
    ----------
18
    storage : DBSocket
19
        A live connection to the current database.
20
    meta : dict
21
        A JSON description of the metadata involved with the computation
22
    molecules : list of str, dict
23
        A list of molecule ID's or full JSON molecules associated with the run.
24

25
    Returns
26
    -------
27
    ret : tuple(dict, list)
28
        A dictionary of JSON representations with keys built in.
29
        The list is an array of any errors occurred
30

31
    Examples
32
    --------
33

34
    >>> meta = {
35
        "procedure": "single",
36
        "driver": "energy",
37
        "method": "HF",
38
        "basis": "sto-3g",
39
        "keywords": "default",
40
        "program": "psi4",
41
    }
42

43
    >>> molecules = [{"geometry": [0, 0, 0], "symbols" : ["He"]}]
44

45
    >>> unpack_single_task_spec(storage, meta, molecules)
46

47
    """
48

49
    # Get the required molecules
50 0
    raw_molecules_query = storage.get_add_molecules_mixed(molecules)
51

52
    # Pull out the needed keywords
53 0
    if meta["keywords"] is None:
54 0
        keyword_set = {}
55
    else:
56 0
        keyword_set = storage.get_add_keywords_mixed([meta["keywords"]])["data"][0]
57 0
        keyword_set = keyword_set["values"]
58

59
    # Create the "universal header"
60 0
    task_meta = json.dumps(
61
        {
62
            "driver": meta["driver"],
63
            "keywords": keyword_set,
64
            "model": {"method": meta["method"], "basis": meta["basis"]},
65
            "extras": {"_qcfractal_tags": {"program": meta["program"], "keywords": meta["keywords"]}},
66
        }
67
    )
68

69 0
    tasks = []
70 0
    for mol in raw_molecules_query["data"]:
71 0
        if mol is None:
72 0
            tasks.append(None)
73 0
            continue
74

75 0
        data = json.loads(task_meta)
76 0
        data["molecule"] = mol
77

78 0
        tasks.append(ResultInput(**data))
79

80 0
    return tasks, []
81

82

83 4
def parse_single_tasks(storage, results):
84
    """Summary
85

86
    Parameters
87
    ----------
88
    storage : DBSocket
89
        A live connection to the current database.
90
    results : dict
91
        A (key, result) dictionary of the single return results.
92

93
    Returns
94
    -------
95

96
    Examples
97
    --------
98

99
    """
100

101 1
    for k, v in results.items():
102 1
        outputs = [v["stdout"], v["stderr"], v["error"]]
103 1
        kvstores = [KVStore(data=x) if x is not None else None for x in outputs]
104 1
        stdout, stderr, error = storage.add_kvstore(kvstores)["data"]
105 1
        v["stdout"] = stdout
106 1
        v["stderr"] = stderr
107 1
        v["error"] = error
108

109
        # Flatten data back out
110 1
        v["method"] = v["model"]["method"]
111 1
        v["basis"] = v["model"]["basis"]
112 1
        del v["model"]
113

114
        # Molecule should be by ID
115 1
        v["molecule"] = storage.add_molecules([Molecule(**v["molecule"])])["data"][0]
116

117 1
        v["keywords"] = v["extras"]["_qcfractal_tags"]["keywords"]
118 1
        v["program"] = v["extras"]["_qcfractal_tags"]["program"]
119 1
        del v["extras"]["_qcfractal_tags"]
120 1
        del v["schema_name"]
121 1
        del v["schema_version"]
122

123 1
        if v.pop("success"):
124 1
            v["status"] = "COMPLETE"
125
        else:
126 0
            v["status"] = "ERROR"
127

128 1
    return results

Read our documentation on viewing source code .

Loading