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

emmet/builders/molecules/atomic.py

@@ -41,7 +41,7 @@ class PartialChargesBuilder(Builder):
     energy) will be used.
 
     The process is as follows:
-        1. Gather MoleculeDocs by formula
+        1. Gather MoleculeDocs by species hash
         2. For each molecule, group all tasks by solvent.
         3. For each solvent, sort tasks by level of theory and electronic energy
         4. For each method:
@@ -86,12 +86,14 @@ class PartialChargesBuilder(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 charges
         self.charges.ensure_index("molecule_id")
@@ -111,23 +113,23 @@ class PartialChargesBuilder(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.charges.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]]:
         """
@@ -152,28 +154,26 @@ class PartialChargesBuilder(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.charges.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
@@ -190,9 +190,9 @@ class PartialChargesBuilder(Builder):
         """
 
         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}")
 
         charges_docs = list()
 
@@ -237,7 +237,7 @@ class PartialChargesBuilder(Builder):
                     tdoc = self.tasks.query_one(
                         {
                             "task_id": task,
-                            "formula_alphabetical": formula,
+                            "species_hash": shash,
                             "orig": {"$exists": True},
                         }
                     )
@@ -247,7 +247,7 @@ class PartialChargesBuilder(Builder):
                             tdoc = self.tasks.query_one(
                                 {
                                     "task_id": int(task),
-                                    "formula_alphabetical": formula,
+                                    "species_hash": shash,
                                     "orig": {"$exists": True},
                                 }
                             )
@@ -271,7 +271,7 @@ class PartialChargesBuilder(Builder):
 
                     charges_docs.append(doc)
 
-        self.logger.debug(f"Produced {len(charges_docs)} charges docs for {formula}")
+        self.logger.debug(f"Produced {len(charges_docs)} charges docs for {shash}")
 
         return jsanitize([doc.model_dump() for doc in charges_docs], allow_bson=True)
 
@@ -320,7 +320,7 @@ class PartialSpinsBuilder(Builder):
     data available (based on level of theory and electronic energy) will be used.
 
     The process is as follows:
-        1. Gather MoleculeDocs by formula
+        1. Gather MoleculeDocs by species_hash
         2. For each molecule, group all tasks by solvent.
         3. For each solvent, sort tasks by level of theory and electronic energy
         4. For each method:
@@ -365,12 +365,14 @@ class PartialSpinsBuilder(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 spins
         self.spins.ensure_index("molecule_id")
@@ -390,23 +392,23 @@ class PartialSpinsBuilder(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.spins.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]]:
         """
@@ -431,28 +433,26 @@ class PartialSpinsBuilder(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.spins.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
@@ -469,9 +469,9 @@ class PartialSpinsBuilder(Builder):
         """
 
         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}")
 
         spins_docs = list()
 
@@ -520,7 +520,7 @@ class PartialSpinsBuilder(Builder):
                     tdoc = self.tasks.query_one(
                         {
                             "task_id": task,
-                            "formula_alphabetical": formula,
+                            "species_hash": shash,
                             "orig": {"$exists": True},
                         }
                     )
@@ -530,7 +530,7 @@ class PartialSpinsBuilder(Builder):
                             tdoc = self.tasks.query_one(
                                 {
                                     "task_id": int(task),
-                                    "formula_alphabetical": formula,
+                                    "species_hash": shash,
                                     "orig": {"$exists": True},
                                 }
                             )
@@ -551,9 +551,7 @@ class PartialSpinsBuilder(Builder):
 
                     spins_docs.append(doc)
 
-        self.logger.debug(
-            f"Produced {len(spins_docs)} partial spins docs for {formula}"
-        )
+        self.logger.debug(f"Produced {len(spins_docs)} partial spins docs for {shash}")
 
         return jsanitize([doc.model_dump() for doc in spins_docs], allow_bson=True)
 

emmet/builders/molecules/bonds.py

@@ -40,7 +40,7 @@ class BondingBuilder(Builder):
     data available (based on level of theory and electronic energy) will be used.
 
     The process is as follows:
-        1. Gather MoleculeDocs by formula
+        1. Gather MoleculeDocs by species hash
         2. For each molecule, group all tasks by solvent.
         3. For each solvent, sort tasks by level of theory and electronic energy
         4. For each method:
@@ -85,12 +85,14 @@ class BondingBuilder(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 bonds
         self.bonds.ensure_index("molecule_id")
@@ -110,23 +112,23 @@ class BondingBuilder(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.bonds.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]]:
         """
@@ -151,28 +153,26 @@ class BondingBuilder(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.bonds.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
@@ -189,9 +189,9 @@ class BondingBuilder(Builder):
         """
 
         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}")
 
         bonding_docs = list()
 
@@ -255,7 +255,7 @@ class BondingBuilder(Builder):
                     tdoc = self.tasks.query_one(
                         {
                             "task_id": task,
-                            "formula_alphabetical": formula,
+                            "species_hash": shash,
                             "orig": {"$exists": True},
                         }
                     )
@@ -265,7 +265,7 @@ class BondingBuilder(Builder):
                             tdoc = self.tasks.query_one(
                                 {
                                     "task_id": int(task),
-                                    "formula_alphabetical": formula,
+                                    "species_hash": shash,
                                     "orig": {"$exists": True},
                                 }
                             )
@@ -288,13 +288,13 @@ class BondingBuilder(Builder):
                     )
                     bonding_docs.append(doc)
 
-        self.logger.debug(f"Produced {len(bonding_docs)} bonding docs for {formula}")
+        self.logger.debug(f"Produced {len(bonding_docs)} bonding docs for {shash}")
 
         return jsanitize([doc.model_dump() for doc in bonding_docs], allow_bson=True)
 
     def update_targets(self, items: List[List[Dict]]):
         """
-        Inserts the new documents into the charges collection
+        Inserts the new documents into the bonds collection
 
         Args:
             items [[dict]]: A list of documents to update

emmet/builders/molecules/electric.py

@@ -0,0 +1,282 @@
+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.electric import ElectricMultipoleDoc
+from emmet.core.utils import jsanitize
+from emmet.builders.settings import EmmetBuildSettings
+
+
+__author__ = "Evan Spotte-Smith"
+
+SETTINGS = EmmetBuildSettings()
+
+
+class ElectricMultipoleBuilder(Builder):
+    """
+    The ElectricMultipoleBuilder defines the electric multipole properties of a MoleculeDoc.
+
+    This builder will attempt to build documents for each molecule, in each solvent.
+    For each molecule-solvent combination, the highest-quality
+    data available (based on level of theory and electronic energy) will be used.
+
+    The process is as follows:
+        1. Gather MoleculeDocs by species hash
+        2. For each molecule, group all tasks by solvent.
+        3. For each solvent, grab the best TaskDoc (doc with elecrtric dipole/multipole information
+        that has the highest level of theory with the lowest electronic energy) for the molecule
+        4. Convert TaskDoc to ElectricMultipoleDoc
+    """
+
+    def __init__(
+        self,
+        tasks: Store,
+        molecules: Store,
+        multipoles: Store,
+        query: Optional[Dict] = None,
+        settings: Optional[EmmetBuildSettings] = None,
+        **kwargs,
+    ):
+        self.tasks = tasks
+        self.molecules = molecules
+        self.multipoles = multipoles
+        self.query = query if query else dict()
+        self.settings = EmmetBuildSettings.autoload(settings)
+        self.kwargs = kwargs
+
+        super().__init__(sources=[tasks, molecules], targets=[multipoles], **kwargs)
+        # Uncomment in case of issue with mrun not connecting automatically to collections
+        # for i in [self.tasks, self.molecules, self.multipoles]:
+        #     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 electric
+        self.multipoles.ensure_index("method")
+        self.multipoles.ensure_index("molecule_id")
+        self.multipoles.ensure_index("task_id")
+        self.multipoles.ensure_index("solvent")
+        self.multipoles.ensure_index("lot_solvent")
+        self.multipoles.ensure_index("property_id")
+        self.multipoles.ensure_index("last_updated")
+        self.multipoles.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.multipoles.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 multipole documents.
+        This does no datetime checking; relying on on whether
+        task_ids are included in the multipoles Store
+
+        Returns:
+            generator or list relevant tasks and molecules to process into documents
+        """
+
+        self.logger.info("Electric multipoles 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.multipoles.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 ElectricMultipoleDocs
+
+        Args:
+            tasks List[Dict] : a list of MoleculeDocs in dict form
+
+        Returns:
+            [dict] : a list of new electric multipole 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}")
+
+        multipole_docs = list()
+
+        for mol in mols:
+            # Relevant tasks are those with the correct charge and spin
+            # for which there are AT LEAST electric dipoles present
+            # (ideally, multipole information would also be present)
+            multip_entries = [
+                e
+                for e in mol.entries
+                if e["charge"] == mol.charge
+                and e["spin_multiplicity"] == mol.spin_multiplicity
+                and (e["output"].get("dipoles") is not None)
+            ]
+
+            # Organize by solvent environment
+            by_solvent = defaultdict(list)
+            for entry in multip_entries:
+                by_solvent[entry["solvent"]].append(entry)
+
+            for solvent, entries in by_solvent.items():
+                # No documents with enthalpy and entropy
+                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
+
+                multipole_doc = ElectricMultipoleDoc.from_task(
+                    task_doc, molecule_id=mol.molecule_id, deprecated=False
+                )
+                multipole_docs.append(multipole_doc)
+
+        self.logger.debug(
+            f"Produced {len(multipole_docs)} electric multipole docs for {shash}"
+        )
+
+        return jsanitize([doc.model_dump() for doc in multipole_docs], allow_bson=True)
+
+    def update_targets(self, items: List[List[Dict]]):
+        """
+        Inserts the new documents into the multipoles 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)} electric multipole documents")
+            self.multipoles.remove_docs({self.multipoles.key: {"$in": molecule_ids}})
+            # Neither molecule_id nor method need to be unique, but the combination must be
+            self.multipoles.update(
+                docs=docs,
+                key=["molecule_id", "solvent"],
+            )
+        else:
+            self.logger.info("No items to update")