emmet-builders 0.84.3rc2__py3-none-any.whl → 0.84.3rc3__py3-none-any.whl

emmet/builders/molecules/thermo.py

@@ -84,7 +84,7 @@ class ThermoBuilder(Builder):
     single-point energy corrections.
 
     Before any documents are constructed, the following steps are taken:
-        1. Gather MoleculeDocs by formula
+        1. Gather MoleculeDocs by species hash
         2. For each doc, identify tasks with thermodynamic information such as
             zero-point energy, enthalpy, and entropy. Collect these "documents
              including complete thermodynamics" (DICTs).
@@ -148,12 +148,14 @@ class ThermoBuilder(Builder):
         self.tasks.ensure_index("last_updated")
         self.tasks.ensure_index("state")
         self.tasks.ensure_index("formula_alphabetical")
+        self.tasks.ensure_index("species_hash")
 
         # Search index for molecules
         self.molecules.ensure_index("molecule_id")
         self.molecules.ensure_index("last_updated")
         self.molecules.ensure_index("task_ids")
         self.molecules.ensure_index("formula_alphabetical")
+        self.molecules.ensure_index("species_hash")
 
         # Search index for thermo
         self.thermo.ensure_index("molecule_id")
@@ -172,23 +174,23 @@ class ThermoBuilder(Builder):
 
         self.logger.info("Finding documents to process")
         all_mols = list(
-            self.molecules.query(
-                temp_query, [self.molecules.key, "formula_alphabetical"]
-            )
+            self.molecules.query(temp_query, [self.molecules.key, "species_hash"])
         )
 
         processed_docs = set([e for e in self.thermo.distinct("molecule_id")])
         to_process_docs = {d[self.molecules.key] for d in all_mols} - processed_docs
-        to_process_forms = {
-            d["formula_alphabetical"]
+        to_process_hashes = {
+            d["species_hash"]
             for d in all_mols
             if d[self.molecules.key] in to_process_docs
         }
 
-        N = ceil(len(to_process_forms) / number_splits)
+        N = ceil(len(to_process_hashes) / number_splits)
 
-        for formula_chunk in grouper(to_process_forms, N):
-            yield {"query": {"formula_alphabetical": {"$in": list(formula_chunk)}}}
+        for hash_chunk in grouper(to_process_hashes, N):
+            query = dict(temp_query)
+            query["species_hash"] = {"$in": list(hash_chunk)}
+            yield {"query": query}
 
     def get_items(self) -> Iterator[List[Dict]]:
         """
@@ -213,28 +215,26 @@ class ThermoBuilder(Builder):
 
         self.logger.info("Finding documents to process")
         all_mols = list(
-            self.molecules.query(
-                temp_query, [self.molecules.key, "formula_alphabetical"]
-            )
+            self.molecules.query(temp_query, [self.molecules.key, "species_hash"])
         )
 
         processed_docs = set([e for e in self.thermo.distinct("molecule_id")])
         to_process_docs = {d[self.molecules.key] for d in all_mols} - processed_docs
-        to_process_forms = {
-            d["formula_alphabetical"]
+        to_process_hashes = {
+            d["species_hash"]
             for d in all_mols
             if d[self.molecules.key] in to_process_docs
         }
 
         self.logger.info(f"Found {len(to_process_docs)} unprocessed documents")
-        self.logger.info(f"Found {len(to_process_forms)} unprocessed formulas")
+        self.logger.info(f"Found {len(to_process_hashes)} unprocessed hashes")
 
         # Set total for builder bars to have a total
-        self.total = len(to_process_forms)
+        self.total = len(to_process_hashes)
 
-        for formula in to_process_forms:
+        for shash in to_process_hashes:
             mol_query = dict(temp_query)
-            mol_query["formula_alphabetical"] = formula
+            mol_query["species_hash"] = shash
             molecules = list(self.molecules.query(criteria=mol_query))
 
             yield molecules
@@ -273,9 +273,9 @@ class ThermoBuilder(Builder):
             return doc
 
         mols = [MoleculeDoc(**item) for item in items]
-        formula = mols[0].formula_alphabetical
+        shash = mols[0].species_hash
         mol_ids = [m.molecule_id for m in mols]
-        self.logger.debug(f"Processing {formula} : {mol_ids}")
+        self.logger.debug(f"Processing {shash} : {mol_ids}")
 
         thermo_docs = list()
 
@@ -334,7 +334,7 @@ class ThermoBuilder(Builder):
                 tdoc = self.tasks.query_one(
                     {
                         "task_id": task,
-                        "formula_alphabetical": formula,
+                        "species_hash": shash,
                         "orig": {"$exists": True},
                     }
                 )
@@ -344,7 +344,7 @@ class ThermoBuilder(Builder):
                         tdoc = self.tasks.query_one(
                             {
                                 "task_id": int(task),
-                                "formula_alphabetical": formula,
+                                "species_hash": shash,
                                 "orig": {"$exists": True},
                             }
                         )
@@ -465,7 +465,7 @@ class ThermoBuilder(Builder):
                     sorted(with_eval_e, key=lambda x: (x[1], x[2]))[0][0]
                 )
 
-        self.logger.debug(f"Produced {len(thermo_docs)} thermo docs for {formula}")
+        self.logger.debug(f"Produced {len(thermo_docs)} thermo docs for {shash}")
 
         return jsanitize([doc.model_dump() for doc in thermo_docs], allow_bson=True)
 

emmet/builders/molecules/trajectory.py

@@ -0,0 +1,525 @@
+from collections import defaultdict
+from datetime import datetime
+from itertools import chain
+from math import ceil
+from typing import Optional, Iterable, Iterator, List, Dict
+
+from maggma.builders import Builder
+from maggma.core import Store
+from maggma.utils import grouper
+
+from emmet.core.qchem.task import TaskDocument
+from emmet.core.qchem.molecule import MoleculeDoc, evaluate_lot
+from emmet.core.molecules.trajectory import ForcesDoc, TrajectoryDoc
+from emmet.core.utils import jsanitize
+from emmet.builders.settings import EmmetBuildSettings
+
+
+__author__ = "Evan Spotte-Smith"
+
+SETTINGS = EmmetBuildSettings()
+
+
+class ForcesBuilder(Builder):
+    """
+    The ForcesBuilder extracts the highest-quality force data from a
+    MoleculeDoc (lowest electronic energy, highest level of theory for
+    each solvent available).
+
+    The process is as follows:
+        1. Gather MoleculeDocs by species hash
+        2. For each doc, sort tasks by solvent
+        3. For each solvent, grab the best TaskDoc (doc with force
+            information that has the highest level of theory with lowest
+             electronic energy for the molecule)
+        4. Convert TaskDoc to ForcesDoc
+
+    """
+
+    def __init__(
+        self,
+        tasks: Store,
+        molecules: Store,
+        forces: Store,
+        query: Optional[Dict] = None,
+        settings: Optional[EmmetBuildSettings] = None,
+        **kwargs,
+    ):
+        self.tasks = tasks
+        self.molecules = molecules
+        self.forces = forces
+        self.query = query if query else dict()
+        self.settings = EmmetBuildSettings.autoload(settings)
+        self.kwargs = kwargs
+
+        super().__init__(sources=[tasks, molecules], targets=[forces], **kwargs)
+        # Uncomment in case of issue with mrun not connecting automatically to collections
+        # for i in [self.tasks, self.molecules, self.forces]:
+        #     try:
+        #         i.connect()
+        #     except Exception as e:
+        #         print("Could not connect,", e)
+
+    def ensure_indexes(self):
+        """
+        Ensures indices on the collections needed for building
+        """
+
+        # Basic search index for tasks
+        self.tasks.ensure_index("task_id")
+        self.tasks.ensure_index("last_updated")
+        self.tasks.ensure_index("state")
+        self.tasks.ensure_index("formula_alphabetical")
+        self.tasks.ensure_index("species_hash")
+
+        # Search index for molecules
+        self.molecules.ensure_index("molecule_id")
+        self.molecules.ensure_index("last_updated")
+        self.molecules.ensure_index("task_ids")
+        self.molecules.ensure_index("formula_alphabetical")
+        self.molecules.ensure_index("species_hash")
+
+        # Search index for force properties
+        self.forces.ensure_index("molecule_id")
+        self.forces.ensure_index("task_id")
+        self.forces.ensure_index("solvent")
+        self.forces.ensure_index("lot_solvent")
+        self.forces.ensure_index("property_id")
+        self.forces.ensure_index("last_updated")
+        self.forces.ensure_index("formula_alphabetical")
+
+    def prechunk(self, number_splits: int) -> Iterable[Dict]:  # pragma: no cover
+        """Prechunk the builder for distributed computation"""
+
+        temp_query = dict(self.query)
+        temp_query["deprecated"] = False
+
+        self.logger.info("Finding documents to process")
+        all_mols = list(
+            self.molecules.query(temp_query, [self.molecules.key, "species_hash"])
+        )
+
+        processed_docs = set([e for e in self.forces.distinct("molecule_id")])
+        to_process_docs = {d[self.molecules.key] for d in all_mols} - processed_docs
+        to_process_hashes = {
+            d["species_hash"]
+            for d in all_mols
+            if d[self.molecules.key] in to_process_docs
+        }
+
+        N = ceil(len(to_process_hashes) / number_splits)
+
+        for hash_chunk in grouper(to_process_hashes, N):
+            query = dict(temp_query)
+            query["species_hash"] = {"$in": list(hash_chunk)}
+            yield {"query": query}
+
+    def get_items(self) -> Iterator[List[Dict]]:
+        """
+        Gets all items to process into force documents.
+        This does no datetime checking; relying on on whether
+        task_ids are included in the forces Store
+
+        Returns:
+            generator or list relevant tasks and molecules to process into documents
+        """
+
+        self.logger.info("Forces builder started")
+        self.logger.info("Setting indexes")
+        self.ensure_indexes()
+
+        # Save timestamp to mark buildtime
+        self.timestamp = datetime.utcnow()
+
+        # Get all processed molecules
+        temp_query = dict(self.query)
+        temp_query["deprecated"] = False
+
+        self.logger.info("Finding documents to process")
+        all_mols = list(
+            self.molecules.query(temp_query, [self.molecules.key, "species_hash"])
+        )
+
+        processed_docs = set([e for e in self.forces.distinct("molecule_id")])
+        to_process_docs = {d[self.molecules.key] for d in all_mols} - processed_docs
+        to_process_hashes = {
+            d["species_hash"]
+            for d in all_mols
+            if d[self.molecules.key] in to_process_docs
+        }
+
+        self.logger.info(f"Found {len(to_process_docs)} unprocessed documents")
+        self.logger.info(f"Found {len(to_process_hashes)} unprocessed hashes")
+
+        # Set total for builder bars to have a total
+        self.total = len(to_process_hashes)
+
+        for shash in to_process_hashes:
+            mol_query = dict(temp_query)
+            mol_query["species_hash"] = shash
+            molecules = list(self.molecules.query(criteria=mol_query))
+
+            yield molecules
+
+    def process_item(self, items: List[Dict]) -> List[Dict]:
+        """
+        Process the tasks into ForcesDoc
+
+        Args:
+            items List[Dict] : a list of MoleculeDocs in dict form
+
+        Returns:
+            [dict] : a list of new forces docs
+        """
+
+        mols = [MoleculeDoc(**item) for item in items]
+        shash = mols[0].species_hash
+        mol_ids = [m.molecule_id for m in mols]
+        self.logger.debug(f"Processing {shash} : {mol_ids}")
+
+        force_docs = list()
+
+        for mol in mols:
+            force_entries = [
+                e
+                for e in mol.entries
+                if e["charge"] == mol.charge and e["task_type"] == "Force"
+            ]
+
+            # Organize by solvent environment
+            by_solvent = defaultdict(list)
+            for entry in force_entries:
+                by_solvent[entry["solvent"]].append(entry)
+
+            for solvent, entries in by_solvent.items():
+                # No force calculations
+                if len(entries) == 0:
+                    continue
+                else:
+                    best = sorted(
+                        entries,
+                        key=lambda x: (
+                            sum(evaluate_lot(x["level_of_theory"])),
+                            x["energy"],
+                        ),
+                    )[0]
+                    task = best["task_id"]
+
+                tdoc = self.tasks.query_one(
+                    {
+                        "task_id": task,
+                        "species_hash": shash,
+                        "orig": {"$exists": True},
+                    }
+                )
+
+                if tdoc is None:
+                    try:
+                        tdoc = self.tasks.query_one(
+                            {
+                                "task_id": int(task),
+                                "species_hash": shash,
+                                "orig": {"$exists": True},
+                            }
+                        )
+                    except ValueError:
+                        tdoc = None
+
+                if tdoc is None:
+                    continue
+
+                task_doc = TaskDocument(**tdoc)
+
+                if task_doc is None:
+                    continue
+
+                force_doc = ForcesDoc.from_task(
+                    task_doc, molecule_id=mol.molecule_id, deprecated=False
+                )
+                force_docs.append(force_doc)
+
+        self.logger.debug(f"Produced {len(force_docs)} force docs for {shash}")
+
+        return jsanitize([doc.model_dump() for doc in force_docs], allow_bson=True)
+
+    def update_targets(self, items: List[List[Dict]]):
+        """
+        Inserts the new force docs into the forces collection
+
+        Args:
+            items [[dict]]: A list of documents to update
+        """
+
+        docs = list(chain.from_iterable(items))  # type: ignore
+
+        # Add timestamp
+        for item in docs:
+            item.update(
+                {
+                    "_bt": self.timestamp,
+                }
+            )
+
+        molecule_ids = list({item["molecule_id"] for item in docs})
+
+        if len(items) > 0:
+            self.logger.info(f"Updating {len(docs)} force documents")
+            self.forces.remove_docs({self.forces.key: {"$in": molecule_ids}})
+            self.forces.update(
+                docs=docs,
+                key=["molecule_id", "solvent"],
+            )
+        else:
+            self.logger.info("No items to update")
+
+
+class TrajectoryBuilder(Builder):
+    """
+    The TrajectoryBuilder extracts the highest-quality optimization trajectory data from a
+    MoleculeDoc. In general, this will mean that the geometry optimization calculation(s)
+    corresponding to the structure(s) that make up this molecule will be used to extract
+    optimization trajectories.
+
+    The process is as follows:
+        1. Gather MoleculeDocs by species hash
+        2. For each doc, sort tasks by solvent
+        3. For each solvent, grab the best TaskDoc (geometry optimization calculation
+            that has the highest level of theory with lowest
+            electronic energy for the molecule)
+        4. Convert TaskDoc to TrajectoryDoc
+    """
+
+    def __init__(
+        self,
+        tasks: Store,
+        molecules: Store,
+        trajectories: Store,
+        query: Optional[Dict] = None,
+        settings: Optional[EmmetBuildSettings] = None,
+        **kwargs,
+    ):
+        self.tasks = tasks
+        self.molecules = molecules
+        self.trajectories = trajectories
+        self.query = query if query else dict()
+        self.settings = EmmetBuildSettings.autoload(settings)
+        self.kwargs = kwargs
+
+        super().__init__(sources=[tasks, molecules], targets=[trajectories], **kwargs)
+        # Uncomment in case of issue with mrun not connecting automatically to collections
+        # for i in [self.tasks, self.molecules, self.trajectories]:
+        #     try:
+        #         i.connect()
+        #     except Exception as e:
+        #         print("Could not connect,", e)
+
+    def ensure_indexes(self):
+        """
+        Ensures indices on the collections needed for building
+        """
+
+        # Basic search index for tasks
+        self.tasks.ensure_index("task_id")
+        self.tasks.ensure_index("last_updated")
+        self.tasks.ensure_index("state")
+        self.tasks.ensure_index("formula_alphabetical")
+        self.tasks.ensure_index("species_hash")
+
+        # Search index for molecules
+        self.molecules.ensure_index("molecule_id")
+        self.molecules.ensure_index("last_updated")
+        self.molecules.ensure_index("task_ids")
+        self.molecules.ensure_index("formula_alphabetical")
+        self.molecules.ensure_index("species_hash")
+
+        # Search index for geometry optimization trajectory properties
+        self.trajectories.ensure_index("molecule_id")
+        self.trajectories.ensure_index("task_id")
+        self.trajectories.ensure_index("solvent")
+        self.trajectories.ensure_index("lot_solvent")
+        self.trajectories.ensure_index("property_id")
+        self.trajectories.ensure_index("last_updated")
+        self.trajectories.ensure_index("formula_alphabetical")
+
+    def prechunk(self, number_splits: int) -> Iterable[Dict]:  # pragma: no cover
+        """Prechunk the builder for distributed computation"""
+
+        temp_query = dict(self.query)
+        temp_query["deprecated"] = False
+
+        self.logger.info("Finding documents to process")
+        all_mols = list(
+            self.molecules.query(temp_query, [self.molecules.key, "species_hash"])
+        )
+
+        processed_docs = set([e for e in self.trajectories.distinct("molecule_id")])
+        to_process_docs = {d[self.molecules.key] for d in all_mols} - processed_docs
+        to_process_hashes = {
+            d["species_hash"]
+            for d in all_mols
+            if d[self.molecules.key] in to_process_docs
+        }
+
+        N = ceil(len(to_process_hashes) / number_splits)
+
+        for hash_chunk in grouper(to_process_hashes, N):
+            yield {"query": {"species_hash": {"$in": list(hash_chunk)}}}
+
+    def get_items(self) -> Iterator[List[Dict]]:
+        """
+        Gets all items to process into trajectory documents.
+        This does no datetime checking; relying on on whether
+        task_ids are included in the forces Store
+
+        Returns:
+            generator or list relevant tasks and molecules to process into documents
+        """
+
+        self.logger.info("Trajectories builder started")
+        self.logger.info("Setting indexes")
+        self.ensure_indexes()
+
+        # Save timestamp to mark buildtime
+        self.timestamp = datetime.utcnow()
+
+        # Get all processed molecules
+        temp_query = dict(self.query)
+        temp_query["deprecated"] = False
+
+        self.logger.info("Finding documents to process")
+        all_mols = list(
+            self.molecules.query(temp_query, [self.molecules.key, "species_hash"])
+        )
+
+        processed_docs = set([e for e in self.trajectories.distinct("molecule_id")])
+        to_process_docs = {d[self.molecules.key] for d in all_mols} - processed_docs
+        to_process_hashes = {
+            d["species_hash"]
+            for d in all_mols
+            if d[self.molecules.key] in to_process_docs
+        }
+
+        self.logger.info(f"Found {len(to_process_docs)} unprocessed documents")
+        self.logger.info(f"Found {len(to_process_hashes)} unprocessed hashes")
+
+        # Set total for builder bars to have a total
+        self.total = len(to_process_hashes)
+
+        for shash in to_process_hashes:
+            mol_query = dict(temp_query)
+            mol_query["species_hash"] = shash
+            molecules = list(self.molecules.query(criteria=mol_query))
+
+            yield molecules
+
+    def process_item(self, items: List[Dict]) -> List[Dict]:
+        """
+        Process the tasks into TrajectoryDocs
+
+        Args:
+            items List[Dict] : a list of MoleculeDocs in dict form
+
+        Returns:
+            [dict] : a list of new trajectory docs
+        """
+
+        mols = [MoleculeDoc(**item) for item in items]
+        shash = mols[0].species_hash
+        mol_ids = [m.molecule_id for m in mols]
+        self.logger.debug(f"Processing {shash} : {mol_ids}")
+
+        trajectory_docs = list()
+
+        for mol in mols:
+            entries = mol.best_entries
+
+            # Organize by solvent environment
+            by_solvent = defaultdict(list)
+            for entry in entries.values():
+                by_solvent[entry["solvent"]].append(entry)
+
+            for solvent, entries in by_solvent.items():
+                # No "best" entry - shouldn't happen, but just in case
+                if len(entries) == 0:
+                    continue
+                else:
+                    # In case there are multiple optimized structures with the same solvent but different LOT
+                    best = sorted(
+                        entries,
+                        key=lambda x: (
+                            sum(evaluate_lot(x["level_of_theory"])),
+                            x["energy"],
+                        ),
+                    )[0]
+                    task = best["task_id"]
+
+                tdoc = self.tasks.query_one(
+                    {
+                        "task_id": task,
+                        "species_hash": shash,
+                        "orig": {"$exists": True},
+                    }
+                )
+
+                if tdoc is None:
+                    try:
+                        tdoc = self.tasks.query_one(
+                            {
+                                "task_id": int(task),
+                                "species_hash": shash,
+                                "orig": {"$exists": True},
+                            }
+                        )
+                    except ValueError:
+                        tdoc = None
+
+                if tdoc is None:
+                    continue
+
+                task_doc = TaskDocument(**tdoc)
+
+                if task_doc is None:
+                    continue
+
+                trajectory_doc = TrajectoryDoc.from_task(
+                    task_doc, molecule_id=mol.molecule_id, deprecated=False
+                )
+                trajectory_docs.append(trajectory_doc)
+
+        self.logger.debug(
+            f"Produced {len(trajectory_docs)} trajectory docs for {shash}"
+        )
+
+        return jsanitize([doc.model_dump() for doc in trajectory_docs], allow_bson=True)
+
+    def update_targets(self, items: List[List[Dict]]):
+        """
+        Inserts the new force docs into the trajectories collection
+
+        Args:
+            items [[dict]]: A list of documents to update
+        """
+
+        docs = list(chain.from_iterable(items))  # type: ignore
+
+        # Add timestamp
+        for item in docs:
+            item.update(
+                {
+                    "_bt": self.timestamp,
+                }
+            )
+
+        molecule_ids = list({item["molecule_id"] for item in docs})
+
+        if len(items) > 0:
+            self.logger.info(f"Updating {len(docs)} trajectory documents")
+            self.trajectories.remove_docs(
+                {self.trajectories.key: {"$in": molecule_ids}}
+            )
+            self.trajectories.update(
+                docs=docs,
+                key=["molecule_id", "solvent"],
+            )
+        else:
+            self.logger.info("No items to update")