import logging import operator from datetime import datetime from itertools import groupby from typing import Iterable, List, Optional, Union from pydantic import field_validator, BaseModel, Field from pymatgen.analysis.structure_matcher import ElementComparator, StructureMatcher from pymatgen.core.composition import Composition from pymatgen.entries.computed_entries import ComputedEntry, ComputedStructureEntry from pymatgen.symmetry.analyzer import SpacegroupAnalyzer from emmet.core.mpid import MPID from emmet.core.common import convert_datetime logger = logging.getLogger(__name__) def generic_groupby(list_in, comp=operator.eq) -> List[int]: """ Group a list of unsortable objects Args: list_in: A list of generic objects comp: (Default value = operator.eq) The comparator Returns: [int] list of labels for the input list """ list_out = [None] * len(list_in) # type: List[Union[int, None]] label_num = 0 for i1, ls1 in enumerate(list_out): if ls1 is not None: continue list_out[i1] = label_num for i2, ls2 in list(enumerate(list_out))[i1 + 1 :]: if comp(list_in[i1], list_in[i2]): if list_out[i2] is None: list_out[i2] = list_out[i1] else: list_out[i1] = list_out[i2] label_num -= 1 label_num += 1 return list_out # type: ignore def s_hash(el): return el.data["comp_delith"] class StructureGroupDoc(BaseModel): """ Group of structure """ group_id: Optional[str] = Field( None, description="The combined material_id of the grouped document is given by the numerically smallest " "material_id, you can also append the followed by the ignored species at the end.", ) has_distinct_compositions: Optional[bool] = Field( None, description="True if multiple compositions are present in the group." ) material_ids: Optional[list] = Field( None, description="A list of materials ids for all of the materials that were grouped together.", ) host_material_ids: Optional[list] = Field( None, description="Material id(s) that correspond(s) to the host structure(s), which has/have the lowest" "concentration of ignored specie.", ) insertion_material_ids: Optional[list] = Field( None, description="Material ids that correspond to the non-host structures identified in a given structure group.", ) framework_formula: Optional[str] = Field( None, description="The chemical formula for the framework (the materials system without the ignored species).", ) ignored_specie: Optional[str] = Field( None, description="Ignored atomic specie which is usually the working ion (ex: Li, Mg, or Ca).", ) chemsys: Optional[str] = Field( None, description="The chemical system this group belongs to, if the atoms for the ignored species is " "present the chemsys will also include the ignored species.", ) last_updated: Optional[datetime] = Field( None, description="Timestamp when this document was built." ) # Make sure that the datetime field is properly formatted @field_validator("last_updated", mode="before") @classmethod def handle_datetime(cls, v): return convert_datetime(cls, v) @classmethod def from_grouped_entries( cls, entries: List[Union[ComputedEntry, ComputedStructureEntry]], ignored_specie: str, ) -> "StructureGroupDoc": """ " Assuming a list of entries are already grouped together, create a StructureGroupDoc Args: entries: A list of entries that is already grouped together. ignored_specie: The specie that is ignored during structure matching """ all_atoms = set() all_comps = set() for ient in entries: all_atoms |= set(ient.composition.as_dict().keys()) all_comps.add(ient.composition.reduced_formula) common_atoms = all_atoms - set([ignored_specie]) if len(common_atoms) == 0: framework_str = "ignored" else: comp_d = {k: entries[0].composition.as_dict()[k] for k in common_atoms} framework_comp = Composition.from_dict(comp_d) framework_str = framework_comp.reduced_formula ids = [ient.data["material_id"] for ient in entries] lowest_id = min(ids, key=_get_id_lexi) sub_script = "_".join([ignored_specie]) host_and_insertion_ids = StructureGroupDoc.get_host_and_insertion_ids( entries=entries, ignored_specie=ignored_specie ) fields = { "group_id": f"{lowest_id}_{sub_script}", "material_ids": ids, "host_material_ids": host_and_insertion_ids["host_ids"], "insertion_material_ids": host_and_insertion_ids["insertion_ids"], "framework_formula": framework_str, "ignored_specie": ignored_specie, "chemsys": "-".join(sorted(all_atoms | set([ignored_specie]))), "has_distinct_compositions": len(all_comps) > 1, } return cls(**fields) @classmethod def from_ungrouped_structure_entries( cls, entries: List[Union[ComputedEntry, ComputedStructureEntry]], ignored_specie: str, ltol: float = 0.2, stol: float = 0.3, angle_tol: float = 5.0, ) -> List["StructureGroupDoc"]: """ Create a list of StructureGroupDocs from a list of ungrouped entries. Args: entries: the list of ComputedStructureEntries to process. ignored_specie: the ignored specie for the structure matcher ltol: length tolerance for the structure matcher stol: site position tolerance for the structure matcher angle_tol: angel tolerance for the structure matcher """ results = [] sm = StructureMatcher( comparator=ElementComparator(), primitive_cell=True, ignored_species=[ignored_specie], ltol=ltol, stol=stol, angle_tol=angle_tol, ) # Add a framework field to each entry's data attribute for ient in entries: ient.data["framework"] = _get_framework( ient.composition.reduced_formula, ignored_specie ) # split into groups for each framework, must sort before grouping entries.sort(key=lambda x: x.data["framework"]) framework_groups = groupby(entries, key=lambda x: x.data["framework"]) cnt_ = 0 for framework, f_group in framework_groups: # if you only have ignored atoms put them into one "ignored" group f_group_l = list(f_group) if framework == "ignored": struct_group = cls.from_grouped_entries( f_group_l, ignored_specie=ignored_specie ) cnt_ += len(struct_group.material_ids) # type: ignore continue logger.debug( f"Performing structure matching for {framework} with {len(f_group_l)} documents." ) for g in group_entries_with_structure_matcher(f_group_l, sm): struct_group = cls.from_grouped_entries( g, ignored_specie=ignored_specie ) cnt_ += len(struct_group.material_ids) # type: ignore results.append(struct_group) if cnt_ != len(entries): raise RuntimeError( "The number of entries in all groups the end does not match the number of supplied entries documents." "Something is seriously wrong, please rebuild the entire database and see if the problem persists." ) return results @staticmethod def get_host_and_insertion_ids( entries: List[Union[ComputedEntry, ComputedStructureEntry]], ignored_specie: str, ) -> dict: host_and_insertion_ids = { "host_id": None, "host_ids": [], "host_entries": [], "insertion_ids": [], } # type:dict ignored_specie_min_fraction = min( [e.composition.get_atomic_fraction(ignored_specie) for e in entries] ) for e in entries: if ( e.composition.get_atomic_fraction(ignored_specie) == ignored_specie_min_fraction ): host_and_insertion_ids["host_entries"].append(e) host_and_insertion_ids["host_ids"].append(e.data["material_id"]) else: host_and_insertion_ids["insertion_ids"].append(e.data["material_id"]) host_and_insertion_ids["host_id"] = min( host_and_insertion_ids["host_entries"], key=lambda x: x.energy_per_atom ).data["material_id"] return host_and_insertion_ids def group_entries_with_structure_matcher( g, struct_matcher: StructureMatcher, working_ion: Optional[str] = None, ) -> Iterable[List[Union[ComputedStructureEntry, ComputedEntry]]]: """ Group the entries together based on similarity of the primitive cells Args: g: a list of entries struct_matcher: the StructureMatcher object used to aggregate structures working_ion: the name of the working ion, if none use the first ignored species from the structure_matcher Returns: subgroups: subgroups that are grouped together based on structure similarity """ if working_ion is None: wion: str = struct_matcher.as_dict()["ignored_species"][0] # Sort the entries by symmetry and by working ion fraction def get_num_sym_ops(ent): sga = SpacegroupAnalyzer(ent.structure) return len(sga.get_space_group_operations()) g.sort(key=get_num_sym_ops, reverse=True) g.sort(key=lambda x: x.composition.get_atomic_fraction(wion)) labs = generic_groupby( g, comp=lambda x, y: struct_matcher.fit(x.structure, y.structure, symmetric=True), ) for ilab in set(labs): sub_g = [g[itr] for itr, jlab in enumerate(labs) if jlab == ilab] yield [el for el in sub_g] def _get_id_lexi(task_id) -> Union[int, str]: """Get a lexicographic representation for a task ID""" # matches "mp-1234" or "1234" followed by and optional "-(Alphanumeric)" mpid = MPID(task_id) return mpid.parts def _get_framework(formula, ignored_specie) -> str: """ Return the reduced formula of the entry without any of the ignored species Return 'ignored' if the all the atoms are ignored """ dd_ = Composition(formula).as_dict() if dd_.keys() == set([ignored_specie]): return "ignored" if ignored_specie in dd_: dd_.pop(ignored_specie) return Composition.from_dict(dd_).reduced_formula