Merge branch 'main' into reconcile

emmet-core/emmet/core/oxidation_states.py

@@ -0,0 +1,87 @@
+import logging
+from collections import defaultdict
+from itertools import groupby
+from typing import Dict, List
+
+import numpy as np
+from pydantic import BaseModel
+from pymatgen.analysis.bond_valence import BVAnalyzer
+from pymatgen.core import Structure
+from pymatgen.core.periodic_table import Specie
+from typing_extensions import Literal
+
+
+class OxidationStateDoc(BaseModel):
+
+    possible_species: List[str]
+    possible_valences: List[float]
+    average_oxidation_states: Dict[str, float]
+    method: Literal["BVAnalyzer", "oxi_state_guesses"]
+    structure: Structure
+
+    @classmethod
+    def from_structure(cls, structure: Structure):
+        structure.remove_oxidation_states()
+        try:
+            bva = BVAnalyzer()
+            valences = bva.get_valences(structure)
+            possible_species = {
+                str(Specie(structure[idx].specie, oxidation_state=valence))
+                for idx, valence in enumerate(valences)
+            }
+
+            structure.add_oxidation_state_by_site(valences)
+
+            # construct a dict of average oxi_states for use
+            # by MP2020 corrections. The format should mirror
+            # the output of the first element from Composition.oxi_state_guesses()
+            # e.g. {'Li': 1.0, 'O': -2.0}
+
+            site_oxidation_list = defaultdict(list)
+            for site in structure:
+                site_oxidation_list[site.specie.element].append(site.specie.oxi_state)
+
+            oxi_state_dict = {
+                str(el): np.mean(oxi_states)
+                for el, oxi_states in site_oxidation_list.items()
+            }
+
+            d = {
+                "possible_species": list(possible_species),
+                "possible_valences": valences,
+                "average_oxidation_states": oxi_state_dict,
+                "method": "BVAnalyzer",
+                "structure": structure,
+            }
+
+            return cls(**d)
+
+        except Exception as e:
+            logging.error("BVAnalyzer failed with: {}".format(e))
+
+            try:
+                first_oxi_state_guess = structure.composition.oxi_state_guesses(
+                    max_sites=-50
+                )[0]
+                valences = [
+                    first_oxi_state_guess[site.species_string] for site in structure
+                ]
+                possible_species = {
+                    str(Specie(el, oxidation_state=valence))
+                    for el, valence in first_oxi_state_guess.items()
+                }
+
+                structure.add_oxidation_state_by_site(valences)
+
+                d = {
+                    "possible_species": list(possible_species),
+                    "possible_valences": valences,
+                    "average_oxidation_states": first_oxi_state_guess,
+                    "method": "oxi_state_guesses",
+                    "structure": structure,
+                }
+
+                return cls(**d)
+
+            except Exception as e:
+                logging.error("Oxidation state guess failed with: {}".format(e))

tests/emmet-core/test_oxidation_states.py

@@ -0,0 +1,53 @@
+import pytest
+from pymatgen.core import Structure
+from pymatgen.util.testing import PymatgenTest
+
+from emmet.core.oxidation_states import OxidationStateDoc
+
+test_structures = {
+    name: struc
+    for name, struc in PymatgenTest.TEST_STRUCTURES.items()
+    if name
+    in [
+        "SiO2",
+        "Li2O",
+        "LiFePO4",
+        "TlBiSe2",
+        "K2O2",
+        "Li3V2(PO4)3",
+        "CsCl",
+        "Li2O2",
+        "NaFePO4",
+        "Pb2TiZrO6",
+        "SrTiO3",
+    ]
+}
+
+fail_structures = {
+    name: struc
+    for name, struc in PymatgenTest.TEST_STRUCTURES.items()
+    if name
+    in [
+        "TiO2",
+        "BaNiO3",
+        "VO2",
+    ]
+}
+
+
+@pytest.mark.parametrize("structure", test_structures.values())
+def test_oxidation_state(structure: Structure):
+    """Very simple test to make sure this actually works"""
+
+    doc = OxidationStateDoc.from_structure(structure)
+    print(structure.composition)
+    assert doc is not None
+
+
+@pytest.mark.parametrize("structure", fail_structures.values())
+def test_oxidation_state_failures(structure: Structure):
+    """Very simple test to make sure this actually fails"""
+
+    doc = OxidationStateDoc.from_structure(structure)
+    print(structure.composition)
+    assert doc is None