Merge pull request #67 from dgasmith/openff

OpenFF Workflow Expansion

qcfractal/interface/client.py

@@ -246,12 +246,25 @@ def add_options(self, opt_list, full_return=False):
     ### Collections section
 
     def list_collections(self, collection_type=None):
+        """Lists the available collections currently on the server.
+
+        Parameters
+        ----------
+        collection_type : None, optional
+            If `None` all collection types will be returned, otherwise only the
+            specified collection type will be returned
+
+        Returns
+        -------
+        dict
+            A dictionary containing the available collection types.
+        """
 
         query = {}
         if collection_type is not None:
             query = {"collection": collection_type.lower()}
 
-        payload = {"meta": {"projection": {"name": True, "collection":True}}, "data": query}
+        payload = {"meta": {"projection": {"name": True, "collection": True}}, "data": query}
         r = self._request("get", "collection", payload)
 
         if collection_type is None:
@@ -262,8 +275,23 @@ def list_collections(self, collection_type=None):
         else:
             return [x["name"] for x in r.json()["data"]]
 
-
     def get_collection(self, collection_type, collection_name, full_return=False):
+        """Aquires a given collection from the server
+
+        Parameters
+        ----------
+        collection_type : str
+            The collection type to be accessed
+        collection_name : str
+            The name of the collection to be accssed
+        full_return : bool, optional
+            If False, returns a Collection object otherwise returns raw JSON
+
+        Returns
+        -------
+        Collection
+            A Collection object if the given collection was found otherwise returns `None`.
+        """
 
         payload = {"meta": {}, "data": [(collection_type.lower(), collection_name)]}
         r = self._request("get", "collection", payload)
@@ -322,7 +350,7 @@ def get_procedures(self, procedure_id, return_objects=True):
         if return_objects:
             ret = []
             for packet in r.json()["data"]:
-                tmp = orm.build_orm(packet)
+                tmp = orm.build_orm(packet, client=self)
                 ret.append(tmp)
             return ret
         else:

qcfractal/interface/collections/biofragment.py

@@ -115,4 +115,5 @@ def submit_torsion_drives(self, options_set, torsions):
                 submissions.append(ret)
         return submissions
 
-collection_utils.register_collection(BioFragment)
+collection_utils.register_collection(BioFragment)
+

qcfractal/interface/collections/dataset.py

@@ -703,4 +703,5 @@ def __getitem__(self, args):
         return self.df[args]
 
 
-register_collection(Dataset)
+register_collection(Dataset)
+

qcfractal/interface/collections/openffworkflow.py

@@ -1,7 +1,6 @@
 """Mongo QCDB Fragment object and helpers
 """
 
-import json
 import copy
 
 from .collection import Collection
@@ -37,38 +36,92 @@ def __init__(self, name, options=None, client=None, **kwargs):
             A Portal client to connect to a server
 
         """
+
+        if client is None:
+            raise KeyError("OpenFFWorkflow must have a client.")
+
         super().__init__(name, client=client, options=options, **kwargs)
 
+        self._torsiondrive_cache = {}
+
+        # First workflow is saved
+        if "id" not in self.data:
+            ret = self.save()
+            if len(ret) == 0:
+                raise ValueError("Attempted to insert duplicate Workflow with name '{}'".format(name))
+            self.data["id"] = ret[0][1]
+
     def _init_collection_data(self, additional_args):
         options = additional_args.get("options", None)
         if options is None:
             raise KeyError("No record of OpenFFWorkflow {} found and no initial options passed in.".format(name))
 
         ret = copy.deepcopy(options)
-        ret["fragments"] = {}  # No known fragments
-        ret["molecules"] = []
-
-        ret["fragment_cache"] = {}  # Caches pulled fragment data
+        ret["fragments"] = {}
 
         return ret
 
     def _pre_save_prep(self, client):
         pass
 
     def get_options(self, key):
+        """
+        Obtains "base" workflow options that do not change.
+
+        Parameters
+        ----------
+        key : str
+            The original workflow options.
+
+        Returns
+        -------
+        dict
+            The requested options dictionary.
+        """
         if key not in self.__required_fields:
             raise KeyError("Key `{}` not understood.".format(key))
 
         return copy.deepcopy(self.data[key])
 
     def list_fragments(self):
-        return copy.deepcopy(list(self.data["fragments"]))
+        """
+        List all fragments associated with this workflow.
 
-    def list_initial_molecules(self):
-        return copy.deepcopy(self.data["molecules"])
+        Returns
+        -------
+        list of str
+            A list of fragment id's.
+        """
+        return list(self.data["fragments"])
 
     def add_fragment(self, fragment_id, data, provenance={}):
+        """
+        Adds a new fragment to the workflow along with the associated torsiondrives required.
 
+        Parameters
+        ----------
+        fragment_id : str
+            The tag associated with fragment. In general this should be the canonical isomeric
+            explicit hydrogen mapped SMILES tag for this fragment.
+        data : dict
+            A dictionary of label : {intial_molecule, grid_spacing, dihedrals} keys.
+
+        provenance : dict, optional
+            The provenance of the fragments creation
+
+        Example
+        -------
+
+        data = {
+           "label1": {
+                "initial_molecule": ptl.data.get_molecule("butane.json"),
+                "grid_spacing": [60],
+                "dihedrals": [[0, 2, 3, 1]],
+            },
+            ...
+        }
+        wf.add_fragment("CCCC", data=)
+        """
         if fragment_id not in self.data["fragments"]:
             self.data["fragments"][fragment_id] = {}
 
@@ -78,13 +131,15 @@ def add_fragment(self, fragment_id, data, provenance={}):
                 print("Already found label {} for fragment_ID {}, skipping.".format(name, fragment_id))
                 continue
 
+            # Build out a new service
             torsion_meta = copy.deepcopy(
                 {k: self.data[k]
                  for k in ("torsiondrive_meta", "optimization_meta", "qc_meta")})
 
             for k in ["grid_spacing", "dihedrals"]:
                 torsion_meta["torsiondrive_meta"][k] = packet[k]
 
+            # Get hash of torsion
             ret = self.client.add_service("torsiondrive", [packet["initial_molecule"]], torsion_meta)
 
             hash_lists = []
@@ -93,30 +148,116 @@ def add_fragment(self, fragment_id, data, provenance={}):
             if len(hash_lists) != 1:
                 raise KeyError("Something went very wrong.")
 
-            hash_index = hash_lists[0]
-            frag_data[name] = hash_index
+            # add back to fragment data
+            packet["hash_index"] = hash_lists[0]
+            frag_data[name] = packet
+
+        # Push collection data back to server
+        self.save(overwrite=True)
 
-        self.data["fragments"][fragment_id] = frag_data
+    def get_fragment_data(self, fragments=None, refresh_cache=False):
+        """Obtains fragment torsiondrives from server to local data.
 
-    def get_data(self):
+        Parameters
+        ----------
+        fragments : None, optional
+            A list of fragment ID's to query upon
+        refresh_cache : bool, optional
+            If True requery everything, otherwise use the cache to prevent extra lookups.
+        """
 
+        # If no fragments explicitly shown, grab all
+        if fragments is None:
+            fragments = self.data["fragments"].keys()
+
+        # Figure out the lookup
         lookup = []
-        for k, v in self.data["fragments"].items():
-            lookup.extend(list(v.values()))
+        for frag in fragments:
+            lookup.extend([v["hash_index"] for v in self.data["fragments"][frag].values()])
+
+        if refresh_cache is False:
+            lookup = list(set(lookup) - self._torsiondrive_cache.keys())
 
+        # Grab the data and update cache
         data = self.client.get_procedures({"hash_index": lookup})
-        data = {x._hash_index: x for x in data}
+        self._torsiondrive_cache.update({x._hash_index: x for x in data})
+
+
+    def list_final_energies(self, fragments=None, refresh_cache=False):
+        """
+        Returns the final energies for the requested fragments.
+
+        Parameters
+        ----------
+        fragments : None, optional
+            A list of fragment ID's to query upon
+        refresh_cache : bool, optional
+            If True requery everything, otherwise use the cache to prevent extra lookups.
+
+        Returns
+        -------
+        dict
+            A dictionary structure with fragment and label fields available for access.
+        """
+
+        # If no fragments explicitly shown, grab all
+        if fragments is None:
+            fragments = self.data["fragments"].keys()
+
+        # Get the data if available
+        self.get_fragment_data(fragments=fragments, refresh_cache=refresh_cache)
+
+        ret = {}
+        for frag in fragments:
+            tmp = {}
+            for k, v in self.data["fragments"][frag].items():
+                if v["hash_index"] in self._torsiondrive_cache:
+                    tmp[k] = self._torsiondrive_cache[v["hash_index"]].final_energies()
+                else:
+                    tmp[k] = None
+
+            ret[frag] = tmp
+
+        return ret
+
+
+    def list_final_molecules(self, fragments=None, refresh_cache=False):
+        """
+        Returns the final molecules for the requested fragments.
+
+        Parameters
+        ----------
+        fragments : None, optional
+            A list of fragment ID's to query upon
+        refresh_cache : bool, optional
+            If True requery everything, otherwise use the cache to prevent extra lookups.
+
+        Returns
+        -------
+        dict
+            A dictionary structure with fragment and label fields available for access.
+        """
+
+        # If no fragments explicitly shown, grab all
+        if fragments is None:
+            fragments = self.data["fragments"].keys()
+
+        # Get the data if available
+        self.get_fragment_data(fragments=fragments, refresh_cache=refresh_cache)
 
         ret = {}
-        for frag, reqs in self.data["fragments"].items():
-            ret[frag] = {}
-            for label, hash_index in reqs.items():
-                try:
-                    ret[frag][label] = {json.dumps(k):v for k, v in data[hash_index].final_energies().items()}
-                except KeyError:
-                    ret[frag][label] = None
+        for frag in fragments:
+            tmp = {}
+            for k, v in self.data["fragments"][frag].items():
+                if v["hash_index"] in self._torsiondrive_cache:
+                    tmp[k] = self._torsiondrive_cache[v["hash_index"]].final_molecules()
+                else:
+                    tmp[k] = None
 
+            ret[frag] = tmp
 
         return ret
 
-collection_utils.register_collection(OpenFFWorkflow)
+
+collection_utils.register_collection(OpenFFWorkflow)
+

qcfractal/interface/orm/build_orm.py

@@ -6,7 +6,7 @@
 from .optimization_orm import OptimizationORM
 
 
-def build_orm(data, procedure=None):
+def build_orm(data, procedure=None, client=None):
     """
     Constructs a Service ORM from incoming JSON data.
 
@@ -16,6 +16,8 @@ def build_orm(data, procedure=None):
         A JSON representation of the procedure.
     procedure : None, optional
         The name of the procedure. If blank the procedure name is pulled from the `data["procedure"]` field.
+    client : FractalClient, optional
+        A activate server connection.
 
     Returns
     -------
@@ -44,8 +46,8 @@ def build_orm(data, procedure=None):
     # import json
     # print(json.dumps(data, indent=2))
     if data["procedure"].lower() == "torsiondrive":
-        return TorsionDriveORM.from_json(data)
+        return TorsionDriveORM.from_json(data, client=client)
     elif data["procedure"].lower() == "optimization":
-        return OptimizationORM.from_json(data)
+        return OptimizationORM.from_json(data, client=client)
     else:
         raise KeyError("Service names {} not recognized.".format(data["procedure"]))