emmet-builders 0.84.10rc2__py3-none-any.whl → 0.85.0__py3-none-any.whl

emmet/builders/molecules/redox.py

@@ -1,502 +0,0 @@
-from __future__ import annotations
-
-import copy
-from collections import defaultdict
-from datetime import datetime
-from itertools import chain, groupby
-from math import ceil
-
-from maggma.builders import Builder
-from maggma.core import Store
-from maggma.utils import grouper
-
-from emmet.builders.settings import EmmetBuildSettings
-from emmet.core.molecules.bonds import metals
-from emmet.core.molecules.redox import RedoxDoc
-from emmet.core.molecules.thermo import MoleculeThermoDoc
-from emmet.core.qchem.molecule import MoleculeDoc
-from emmet.core.qchem.task import TaskDocument
-from emmet.core.utils import confirm_molecule, get_graph_hash, jsanitize
-
-from typing import TYPE_CHECKING
-
-if TYPE_CHECKING:
-    from collections.abc import Iterable, Iterator
-    from typing import Any
-
-__author__ = "Evan Spotte-Smith"
-
-SETTINGS = EmmetBuildSettings()
-
-
-class RedoxBuilder(Builder):
-    """
-    The RedoxBuilder extracts the highest-quality redox data (vertical and
-    adiabatic reduction and oxidation potentials, etc.)
-    from a MoleculeDoc (lowest electronic energy, highest level of theory).
-
-    The process is as follows:
-        1. Gather MoleculeDocs by species hash
-        2. Further group based on (covalent) isomorphism and charge
-        3. For each MoleculeDoc:
-            3a. Identify relevant MoleculeThermoDocs
-            3b. Look for single-point energy calculations conducted at the
-            molecule's charge +- 1. These will be used to calculation
-            vertical electron affinities and ionization energies
-            3c. Group MoleculeThermoDocs and single-point calculations based on solvent
-            and level of theory
-        4. Construct RedoxDocs by looking for molecules (with associated
-            calculations) that:
-            - Have charges that differ by +- 1
-            - Use the same solvent and level of theory
-    """
-
-    def __init__(
-        self,
-        tasks: Store,
-        molecules: Store,
-        thermo: Store,
-        redox: Store,
-        query: dict | None = None,
-        settings: EmmetBuildSettings | None = None,
-        **kwargs,
-    ):
-        self.tasks = tasks
-        self.molecules = molecules
-        self.thermo = thermo
-        self.redox = redox
-        self.query = query if query else dict()
-        self.settings = EmmetBuildSettings.autoload(settings)
-        self.kwargs = kwargs
-
-        super().__init__(sources=[tasks, molecules, thermo], targets=[redox], **kwargs)
-        # Uncomment in case of issue with mrun not connecting automatically to collections
-        # for i in [self.tasks, self.molecules, self.thermo, self.redox]:
-        #     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 thermo
-        self.thermo.ensure_index("molecule_id")
-        self.thermo.ensure_index("task_id")
-        self.thermo.ensure_index("solvent")
-        self.thermo.ensure_index("lot_solvent")
-        self.thermo.ensure_index("property_id")
-        self.thermo.ensure_index("last_updated")
-        self.thermo.ensure_index("formula_alphabetical")
-
-        # Search index for redox
-        self.redox.ensure_index("molecule_id")
-        self.redox.ensure_index("solvent")
-        self.redox.ensure_index("lot_solvent")
-        self.redox.ensure_index("property_id")
-        self.redox.ensure_index("last_updated")
-        self.redox.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.redox.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 redox documents.
-        This does no datetime checking; relying on on whether
-        task_ids are included in the orbitals Store
-
-        Returns:
-            generator or list relevant tasks and molecules to process into documents
-        """
-
-        self.logger.info("Redox 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.redox.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 a RedoxDoc
-
-        Args:
-            tasks list[dict] : a list of MoleculeDocs in dict form
-
-        Returns:
-            [dict] : a list of new redox 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}")
-
-        redox_docs = list()
-
-        # Group by (covalent) molecular graph connectivity
-        group_by_graph = self._group_by_graph(mols)
-
-        for graph_group in group_by_graph.values():
-            # Molecule docs will be grouped by charge
-            charges: dict[int, Any] = dict()
-
-            for gg in graph_group:
-                # First, grab relevant MoleculeThermoDocs and identify possible IE/EA single-points
-                thermo_docs = [
-                    MoleculeThermoDoc(**e)
-                    for e in self.thermo.query({"molecule_id": gg.molecule_id})
-                ]
-
-                if len(thermo_docs) == 0:
-                    # Current building scheme requires a MoleculeThermoDoc
-                    continue
-
-                ie_sp_task_ids = [
-                    e["task_id"]
-                    for e in gg.entries
-                    if e["charge"] == gg.charge + 1
-                    and e["task_type"] in ["Single Point", "Force"]
-                    and e["output"].get("final_energy")
-                ]
-                ie_tasks = list()
-                for i in ie_sp_task_ids:
-                    tdoc = self.tasks.query_one(
-                        {
-                            "task_id": i,
-                            "species_hash": shash,
-                            "orig": {"$exists": True},
-                        }
-                    )
-
-                    if tdoc is None:
-                        try:
-                            tdoc = self.tasks.query_one(
-                                {
-                                    "task_id": int(i),
-                                    "species_hash": shash,
-                                    "orig": {"$exists": True},
-                                }
-                            )
-                        except ValueError:
-                            tdoc = None
-
-                    if tdoc is None:
-                        continue
-
-                    ie_tasks.append(TaskDocument(**tdoc))
-
-                ea_sp_task_ids = [
-                    e["task_id"]
-                    for e in gg.entries
-                    if e["charge"] == gg.charge - 1
-                    and e["task_type"] in ["Single Point", "Force"]
-                    and e["output"].get("final_energy")
-                ]
-                ea_tasks = list()
-                for i in ea_sp_task_ids:
-                    tdoc = self.tasks.query_one(
-                        {
-                            "task_id": i,
-                            "species_hash": shash,
-                            "orig": {"$exists": True},
-                        }
-                    )
-
-                    if tdoc is None:
-                        try:
-                            tdoc = self.tasks.query_one(
-                                {
-                                    "task_id": int(i),
-                                    "species_hash": shash,
-                                    "orig": {"$exists": True},
-                                }
-                            )
-                        except ValueError:
-                            tdoc = None
-
-                    if tdoc is None:
-                        continue
-
-                    ea_tasks.append(TaskDocument(**tdoc))
-
-                grouped_docs = self._collect_by_lot_solvent(
-                    thermo_docs, ie_tasks, ea_tasks
-                )
-                if gg.charge in charges:
-                    charges[gg.charge].append((gg, grouped_docs))
-                else:
-                    charges[gg.charge] = [(gg, grouped_docs)]
-
-            for charge, collection in charges.items():
-                for mol, docs in collection:
-                    # Get all possible molecules for adiabatic oxidation and reduction
-                    red_coll = charges.get(charge - 1, list())
-                    ox_coll = charges.get(charge + 1, list())
-
-                    for lot_solv, docset in docs.items():
-                        # Collect other molecules that have MoleculeThermoDocs at the
-                        # exact same level of theory
-
-                        combined = docset["thermo_doc"].combined_lot_solvent
-
-                        relevant_red = list()
-                        relevant_ox = list()
-
-                        for rmol, rdocs in red_coll:
-                            if lot_solv in rdocs:
-                                if (
-                                    rdocs[lot_solv]["thermo_doc"].combined_lot_solvent
-                                    == combined
-                                ):
-                                    relevant_red.append(rdocs[lot_solv])
-
-                        for omol, odocs in ox_coll:
-                            if lot_solv in odocs:
-                                if (
-                                    odocs[lot_solv]["thermo_doc"].combined_lot_solvent
-                                    == combined
-                                ):
-                                    relevant_ox.append(odocs[lot_solv])
-
-                        # Take best options (based on electronic energy), where available
-                        if len(relevant_red) == 0:
-                            red_doc = None
-                        else:
-                            red_doc = sorted(
-                                relevant_red,
-                                key=lambda x: x["thermo_doc"].electronic_energy,
-                            )[0]["thermo_doc"]
-
-                        if len(relevant_ox) == 0:
-                            ox_doc = None
-                        else:
-                            ox_doc = sorted(
-                                relevant_ox,
-                                key=lambda x: x["thermo_doc"].electronic_energy,
-                            )[0]["thermo_doc"]
-
-                        ea_doc = docset.get("ea_doc")
-                        ie_doc = docset.get("ie_doc")
-
-                        redox_docs.append(
-                            RedoxDoc.from_docs(
-                                base_molecule_doc=mol,
-                                base_thermo_doc=docset["thermo_doc"],
-                                red_doc=red_doc,
-                                ox_doc=ox_doc,
-                                ea_doc=ea_doc,
-                                ie_doc=ie_doc,
-                            )
-                        )
-
-        self.logger.debug(f"Produced {len(redox_docs)} redox docs for {shash}")
-
-        return jsanitize(
-            [doc.model_dump() for doc in redox_docs if doc is not None], allow_bson=True
-        )
-
-    def update_targets(self, items: list[list[dict]]):
-        """
-        Inserts the new documents into the orbitals 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)} redox documents")
-            self.redox.remove_docs({self.redox.key: {"$in": molecule_ids}})
-            self.redox.update(
-                docs=docs,
-                key=["molecule_id", "solvent"],
-            )
-        else:
-            self.logger.info("No items to update")
-
-    @staticmethod
-    def _group_by_graph(mol_docs: list[MoleculeDoc]) -> dict[int, list[MoleculeDoc]]:
-        """
-        Group molecule docs by molecular graph connectivity
-
-        :param entries: List of entries (dicts derived from TaskDocuments)
-        :return: Grouped molecule entries
-        """
-
-        graph_hashes_nometal: list[str] = list()
-        results = defaultdict(list)
-
-        # Within each group, group by the covalent molecular graph
-        for t in mol_docs:
-            mol = confirm_molecule(t.molecule)
-
-            mol_nometal = copy.deepcopy(mol)
-
-            if mol.composition.alphabetical_formula not in [m + "1" for m in metals]:
-                mol_nometal.remove_species(metals)
-
-            mol_nometal.set_charge_and_spin(0)
-            gh_nometal = get_graph_hash(mol_nometal, node_attr="specie")
-
-            match = None
-            for i, gh in enumerate(graph_hashes_nometal):
-                if gh_nometal == gh:
-                    match = i
-                    break
-
-            if match is None:
-                results[len(graph_hashes_nometal)].append(t)
-                graph_hashes_nometal.append(gh_nometal)
-            else:
-                results[match].append(t)
-
-        return results
-
-    @staticmethod
-    def _collect_by_lot_solvent(
-        thermo_docs: list[MoleculeThermoDoc],
-        ie_docs: list[TaskDocument],
-        ea_docs: list[TaskDocument],
-    ) -> dict[str, Any]:
-        """
-        For a given MoleculeDoc, group potential MoleculeThermoDocs and TaskDocs for
-        IE/EA calculations based on level of theory and solvent.
-
-        Args:
-            thermo_docs (list of MoleculeThermoDocs): List of MoleculeThermoDocs for this MoleculeDoc
-            ie_docs (list of TaskDocuments): List of TaskDocs which could be used
-                to calculate vertical ionization energies for this MoleculeDoc
-            ea_docs (list of TaskDocuments): List of TaskDocs which could be used
-                to calculate vertical electron affinities for this MoleculeDoc:
-
-        Returns:
-            dict {<lot_solvent>: {
-                        "thermo_doc": MoleculeThermoDoc, "ie_doc": TaskDocument, "ea_doc": TaskDocument
-                    }
-                 }
-        """
-
-        def _lot_solv(doc: MoleculeThermoDoc | TaskDocument):
-            if isinstance(doc, MoleculeThermoDoc):
-                if doc.correction:
-                    return doc.correction_lot_solvent
-            return doc.lot_solvent
-
-        thermo_grouped = groupby(sorted(thermo_docs, key=_lot_solv), key=_lot_solv)
-        ie_grouped = groupby(sorted(ie_docs, key=_lot_solv), key=_lot_solv)
-        ea_grouped = groupby(sorted(ea_docs, key=_lot_solv), key=_lot_solv)
-
-        groups = dict()
-
-        for k, g in thermo_grouped:
-            g_list = list(g)
-
-            # Should never be more than one MoleculeThermoDoc per MoleculeDoc
-            # Just for safety...
-            if len(g_list) > 1:
-                g_list_sorted = sorted(g_list, key=lambda x: x.electronic_energy)
-                this_thermo_doc = g_list_sorted[0]
-            else:
-                this_thermo_doc = g_list[0]
-
-            groups[k] = {"thermo_doc": this_thermo_doc}
-
-        for k, g in ie_grouped:
-            # Must be a MoleculeThermoDoc to make a RedoxDoc
-            if k not in groups:
-                continue
-
-            this_ie_doc = sorted(list(g), key=lambda x: x.output.final_energy)[0]
-            groups[k]["ie_doc"] = this_ie_doc
-
-        for k, g in ea_grouped:
-            # Must be a MoleculeThermoDoc to make a RedoxDoc
-            if k not in groups:
-                continue
-
-            this_ea_doc = sorted(list(g), key=lambda x: x.output.final_energy)[0]
-            groups[k]["ea_doc"] = this_ea_doc
-
-        return groups