Merge pull request #174 from EFuem/v1-cf-changes

ColabFit Changes
openkim · Apr 1, 2024 · 9ee3ce8 · 9ee3ce8
2 parents 54c4179 + 4138e57
commit 9ee3ce8
Showing 1 changed file with 18 additions and 11 deletions.
diff --git a/kliff/dataset/dataset.py b/kliff/dataset/dataset.py
@@ -179,16 +179,12 @@ def from_colabfit(
         """
         try:
             configuration_id = data_object["relationships"][0]["configuration"]
-            fetched_configuration = database_client.configurations.find_one(
-                {"colabfit-id": data_object["relationships"][0]["configuration"]}
+            fetched_configuration = database_client.get_cleaned_configuration(
+                configuration_id
             )
             fetched_properties = list(
-                database_client.property_instances.find(
-                    {
-                        "colabfit-id": {
-                            "$in": data_object["relationships"][0]["property_instance"]
-                        }
-                    }
+                database_client.get_cleaned_property_instances(
+                    data_object["relationships"][0]["property_instance"]
                 )
             )
         except:
@@ -228,8 +224,17 @@ def from_colabfit(
             weight=weight,
         )
         self.metadata = {
-            "data_object": data_object,
+            "do-id": data_object["colabfit-id"],
+            "co-id": fetched_configuration["colabfit-id"],
+            "pi-ids": [pi["colabfit-id"] for pi in fetched_properties],
+            "names": fetched_configuration["names"],
         }
+        # Update self.metadata with information from metadata collection
+        md_dict = database_client.get_metadata_from_do_doc(data_object)
+        if md_dict:
+            md_dict["md-id"] = md_dict["colabfit-id"]
+            md_dict.pop("colabfit-id")
+            self.metadata.update(md_dict)
 
         return self
 
@@ -581,6 +586,7 @@ def from_colabfit(
         colabfit_dataset: str,
         colabfit_uri: str = "mongodb://localhost:27017",
         weight: Optional[Weight] = None,
+        **kwargs,
     ) -> "Dataset":
         """
         Read configurations from colabfit database and initialize a dataset.
@@ -598,7 +604,7 @@ def from_colabfit(
         """
         instance = cls()
         instance.add_from_colabfit(
-            colabfit_database, colabfit_dataset, colabfit_uri, weight
+            colabfit_database, colabfit_dataset, colabfit_uri, weight, **kwargs
         )
         return instance
 
@@ -650,6 +656,7 @@ def add_from_colabfit(
         colabfit_dataset: str,
         colabfit_uri: str = "mongodb://localhost:27017",
         weight: Optional[Weight] = None,
+        **kwargs,
     ):
         """
         Read configurations from colabfit database and add them to the dataset.
@@ -664,7 +671,7 @@ def add_from_colabfit(
 
         """
         # open link to the mongo
-        mongo_client = MongoDatabase(colabfit_database, uri=colabfit_uri)
+        mongo_client = MongoDatabase(colabfit_database, uri=colabfit_uri, **kwargs)
         configs = Dataset._read_from_colabfit(mongo_client, colabfit_dataset, weight)
         self.configs.extend(configs)