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

emmet/builders/molecules/summary.py

@@ -1,406 +0,0 @@
-from __future__ import annotations
-
-from datetime import datetime
-from itertools import chain
-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.summary import MoleculeSummaryDoc
-from emmet.core.utils import jsanitize
-
-
-from typing import TYPE_CHECKING
-
-if TYPE_CHECKING:
-    from collections.abc import Iterable, Iterator
-    from typing import Any
-
-# from monty.serialization import loadfn, dumpfn
-
-
-__author__ = "Evan Spotte-Smith"
-
-SETTINGS = EmmetBuildSettings()
-
-
-class SummaryBuilder(Builder):
-    """
-    The SummaryBuilder collects all property documents and gathers their properties
-    into a single MoleculeSummaryDoc
-
-    The process is as follows:
-        1. Gather MoleculeDocs by formula
-        2. For each doc, grab the relevant property docs
-        3. Convert property docs to MoleculeSummaryDoc
-    """
-
-    def __init__(
-        self,
-        molecules: Store,
-        charges: Store,
-        spins: Store,
-        bonds: Store,
-        multipoles: Store,
-        metal_binding: Store,
-        orbitals: Store,
-        redox: Store,
-        thermo: Store,
-        vibes: Store,
-        summary: Store,
-        query: dict | None = None,
-        settings: EmmetBuildSettings | None = None,
-        **kwargs,
-    ):
-        self.molecules = molecules
-        self.charges = charges
-        self.spins = spins
-        self.bonds = bonds
-        self.multipoles = multipoles
-        self.metal_binding = metal_binding
-        self.orbitals = orbitals
-        self.redox = redox
-        self.thermo = thermo
-        self.vibes = vibes
-        self.summary = summary
-        self.query = query if query else dict()
-        self.settings = EmmetBuildSettings.autoload(settings)
-        self.kwargs = kwargs
-
-        super().__init__(
-            sources=[
-                molecules,
-                charges,
-                spins,
-                bonds,
-                multipoles,
-                metal_binding,
-                orbitals,
-                redox,
-                thermo,
-                vibes,
-            ],
-            targets=[summary],
-            **kwargs,
-        )
-        # Uncomment in case of issue with mrun not connecting automatically to collections
-        # for i in [
-        #     self.molecules,
-        #     self.charges,
-        #     self.spins,
-        #     self.bonds,
-        #     self.multipoles,
-        #     self.metal_binding,
-        #     self.orbitals,
-        #     self.redox,
-        #     self.thermo,
-        #     self.vibes,
-        #     self.summary
-        # ]:
-        #     try:
-        #         i.connect()
-        #     except Exception as e:
-        #         print("Could not connect,", e)
-
-    def ensure_indexes(self):
-        """
-        Ensures indices on the collections needed for building
-        """
-
-        # 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")
-        self.charges.ensure_index("method")
-        self.charges.ensure_index("task_id")
-        self.charges.ensure_index("solvent")
-        self.charges.ensure_index("lot_solvent")
-        self.charges.ensure_index("property_id")
-        self.charges.ensure_index("last_updated")
-        self.charges.ensure_index("formula_alphabetical")
-
-        # Search index for charges
-        self.spins.ensure_index("molecule_id")
-        self.spins.ensure_index("method")
-        self.spins.ensure_index("task_id")
-        self.spins.ensure_index("solvent")
-        self.spins.ensure_index("lot_solvent")
-        self.spins.ensure_index("property_id")
-        self.spins.ensure_index("last_updated")
-        self.spins.ensure_index("formula_alphabetical")
-
-        # Search index for charges
-        self.bonds.ensure_index("molecule_id")
-        self.bonds.ensure_index("method")
-        self.bonds.ensure_index("task_id")
-        self.bonds.ensure_index("solvent")
-        self.bonds.ensure_index("lot_solvent")
-        self.bonds.ensure_index("property_id")
-        self.bonds.ensure_index("last_updated")
-        self.bonds.ensure_index("formula_alphabetical")
-
-        # Search index for multipoles
-        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")
-
-        # Search index for metal_binding
-        self.metal_binding.ensure_index("molecule_id")
-        self.metal_binding.ensure_index("solvent")
-        self.metal_binding.ensure_index("lot_solvent")
-        self.metal_binding.ensure_index("property_id")
-        self.metal_binding.ensure_index("last_updated")
-        self.metal_binding.ensure_index("formula_alphabetical")
-        self.metal_binding.ensure_index("method")
-
-        # Search index for orbitals
-        self.orbitals.ensure_index("molecule_id")
-        self.orbitals.ensure_index("task_id")
-        self.orbitals.ensure_index("solvent")
-        self.orbitals.ensure_index("lot_solvent")
-        self.orbitals.ensure_index("property_id")
-        self.orbitals.ensure_index("last_updated")
-        self.orbitals.ensure_index("formula_alphabetical")
-
-        # Search index for orbitals
-        self.redox.ensure_index("molecule_id")
-        self.redox.ensure_index("task_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")
-
-        # 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 vibrational properties
-        self.vibes.ensure_index("molecule_id")
-        self.vibes.ensure_index("task_id")
-        self.vibes.ensure_index("solvent")
-        self.vibes.ensure_index("lot_solvent")
-        self.vibes.ensure_index("property_id")
-        self.vibes.ensure_index("last_updated")
-        self.vibes.ensure_index("formula_alphabetical")
-
-        # Search index for molecules
-        self.summary.ensure_index("molecule_id")
-        self.summary.ensure_index("last_updated")
-        self.summary.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.summary.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 summary documents.
-        This does no datetime checking; relying on on whether
-        task_ids are included in the summary Store
-
-        Returns:
-            generator or list relevant tasks and molecules to process into documents
-        """
-
-        self.logger.info("Summary 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.summary.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 MoleculeSummaryDoc
-
-        Args:
-            tasks list[dict] : a list of MoleculeDocs in dict form
-
-        Returns:
-            [dict] : a list of new orbital docs
-        """
-
-        def _group_docs(docs: list[dict[str, Any]], by_method: bool = False):
-            """Helper function to group docs by solvent"""
-            grouped: dict[str, Any] = dict()
-
-            for doc in docs:
-                solvent = doc.get("solvent")
-                method = doc.get("method")
-                if not solvent:
-                    # Need to group by solvent
-                    continue
-                if by_method and method is None:
-                    # Trying to group by method, but no method present
-                    continue
-
-                if not by_method:
-                    grouped[solvent] = doc
-                else:
-                    if solvent not in grouped:
-                        grouped[solvent] = {method: doc}
-                    else:
-                        grouped[solvent][method] = doc
-
-            return grouped
-
-        mols = items
-        shash = mols[0]["species_hash"]
-        mol_ids = [m["molecule_id"] for m in mols]
-        self.logger.debug(f"Processing {shash} : {mol_ids}")
-
-        summary_docs = list()
-
-        for mol in mols:
-            mol_id = mol["molecule_id"]
-
-            d = {
-                "molecules": mol,
-                "partial_charges": _group_docs(
-                    list(self.charges.query({"molecule_id": mol_id})), True
-                ),
-                "partial_spins": _group_docs(
-                    list(self.spins.query({"molecule_id": mol_id})), True
-                ),
-                "bonding": _group_docs(
-                    list(self.bonds.query({"molecule_id": mol_id})), True
-                ),
-                "metal_binding": _group_docs(
-                    list(self.metal_binding.query({"molecule_id": mol_id})), True
-                ),
-                "multipole_moments": _group_docs(
-                    list(self.multipoles.query({"molecule_id": mol_id})), False
-                ),
-                "orbitals": _group_docs(
-                    list(self.orbitals.query({"molecule_id": mol_id})), False
-                ),
-                "redox": _group_docs(
-                    list(self.redox.query({"molecule_id": mol_id})), False
-                ),
-                "thermo": _group_docs(
-                    list(self.thermo.query({"molecule_id": mol_id})), False
-                ),
-                "vibration": _group_docs(
-                    list(self.vibes.query({"molecule_id": mol_id})), False
-                ),
-            }
-
-            to_delete = list()
-
-            for k, v in d.items():
-                if isinstance(v, dict) and len(v) == 0:
-                    to_delete.append(k)
-
-            for td in to_delete:
-                del d[td]
-
-            # # For debugging; keep because it might be needed again
-            # dumpfn(d, f"{mol_id}.json.gz")
-            # break
-
-            summary_doc = MoleculeSummaryDoc.from_docs(molecule_id=mol_id, docs=d)
-            summary_docs.append(summary_doc)
-
-        self.logger.debug(f"Produced {len(summary_docs)} summary docs for {shash}")
-
-        return jsanitize([doc.model_dump() for doc in summary_docs], allow_bson=True)
-
-    def update_targets(self, items: list[list[dict]]):
-        """
-        Inserts the new documents into the summary 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)} summary documents")
-            self.summary.remove_docs({self.summary.key: {"$in": molecule_ids}})
-            self.summary.update(
-                docs=docs,
-                key=["molecule_id"],
-            )
-        else:
-            self.logger.info("No items to update")