[WIP] add supercell fields to MigrationGraphDoc (#459)

* add new attributes and static method skeleton

* delete redundant imports

* linting

* add warning field to MGDoc + sc field method skeleton

* fix test

* mypy

* added warning + deprecated, static method

* fix flake8

* fix mypy

* linting

* change deprecated default

* add test for sc method

* linting

* change required fields + kwargs

* test file fix

emmet-core/emmet/core/mobility/migrationgraph.py

@@ -1,10 +1,14 @@
 from datetime import datetime
-from typing import List, Union, Dict
+from typing import List, Type, Union, Dict, Tuple, Sequence
 
 from pydantic import BaseModel, Field, validator
+import numpy as np
 from emmet.core.base import EmmetBaseModel
+from pymatgen.core import Structure
+from pymatgen.analysis.structure_matcher import StructureMatcher
 from pymatgen.entries.computed_entries import ComputedEntry, ComputedStructureEntry
 from pymatgen.analysis.diffusion.neb.full_path_mapper import MigrationGraph
+from pymatgen.analysis.diffusion.utils.supercells import get_sc_fromstruct
 
 
 class MigrationGraphDoc(EmmetBaseModel):
@@ -14,13 +18,23 @@ class MigrationGraphDoc(EmmetBaseModel):
     Note: this doc is not self-contained within pymatgen, as it has dependence on pymatgen.analysis.diffusion, a namespace package aka pymatgen-diffusion.
     """
 
-    battery_id: str = Field(None, description="The battery id for this MigrationGraphDoc")
+    battery_id: str = Field(..., description="The battery id for this MigrationGraphDoc")
 
     last_updated: datetime = Field(
         None,
         description="Timestamp for the most recent calculation for this MigrationGraph document.",
     )
 
+    warnings: Sequence[str] = Field(
+        [],
+        description="Any warnings related to this property."
+    )
+
+    deprecated: bool = Field(
+        False,
+        description="Indicates whether a migration graph fails to be constructed from the provided entries. Defaults to False, indicating mg can be constructed from entries."
+    )
+
     hop_cutoff: float = Field(
         None,
         description="The numerical value in angstroms used to cap the maximum length of a hop."
@@ -37,21 +51,62 @@ class MigrationGraphDoc(EmmetBaseModel):
     )
 
     migration_graph: MigrationGraph = Field(
-        None,
+        ...,
         description="The MigrationGraph object as defined in pymatgen.analysis.diffusion."
     )
 
+    populate_sc_fields: bool = Field(
+        True,
+        description="Flag indicating whether this document has populated the supercell fields"
+    )
+
+    min_length_sc: float = Field(
+        None,
+        description="The minimum length used to generate supercell using pymatgen."
+    )
+
+    minmax_num_atoms: Tuple[int, int] = Field(
+        None,
+        description="The min/max number of atoms used to genreate supercell using pymatgen."
+    )
+
+    matrix_supercell_structure: Structure = Field(
+        None,
+        description="The matrix suprcell structure that does not contain the mobile ions for the purpose of migration analysis."
+    )
+
+    conversion_matrix: List[List[Union[int, float]]] = Field(
+        None,
+        description="The conversion matrix used to convert unit cell to supercell."
+    )
+
+    inserted_ion_coords: Dict = Field(
+        None,
+        description="A dictionary containing all mobile ion fractional coordinates in terms of supercell."
+    )
+
+    insert_coords_combo: List[str] = Field(
+        None,
+        description="A list of combinations 'a+b' to designate hops in the supercell. Each combo should correspond to one unique hop in MigrationGraph."
+    )
+
     @classmethod
     def from_entries_and_distance(
         cls,
         battery_id: str,
         grouped_entries: List[ComputedStructureEntry],
         working_ion_entry: Union[ComputedEntry, ComputedStructureEntry],
-        hop_cutoff: float
+        hop_cutoff: float,
+        populate_sc_fields: bool = True,
+        ltol: float = 0.2,
+        stol: float = 0.3,
+        angle_tol: float = 5,
+        **kwargs,
     ) -> Union["MigrationGraphDoc", None]:
         """
         This classmethod takes a group of ComputedStructureEntries (can also use ComputedEntry for wi) and generates a full sites structure.
         Then a MigrationGraph object is generated with with_distance() method with a designated cutoff.
+        If populate_sc_fields set to True, this method will populate the supercell related fields. Required kwargs are min_length_sc and minmax_num_atoms.
         """
 
         ranked_structures = MigrationGraph.get_structure_from_entries(
@@ -66,10 +121,178 @@ def from_entries_and_distance(
             max_distance=hop_cutoff
         )
 
-        return cls(
-            battery_id=battery_id,
-            hop_cutoff=hop_cutoff,
-            entries_for_generation=grouped_entries,
-            working_ion_entry=working_ion_entry,
-            migration_graph=migration_graph
+        if not populate_sc_fields:
+            return cls(
+                battery_id=battery_id,
+                hop_cutoff=hop_cutoff,
+                entries_for_generation=grouped_entries,
+                working_ion_entry=working_ion_entry,
+                migration_graph=migration_graph,
+                **kwargs
+            )
+
+        else:
+
+            if all(arg in kwargs for arg in ["min_length_sc", "minmax_num_atoms"]):
+                sm = StructureMatcher(ltol, stol, angle_tol)
+                host_sc, sc_mat, min_length_sc, minmax_num_atoms, coords_dict, combo = MigrationGraphDoc.generate_sc_fields(
+                    mg=migration_graph,
+                    min_length_sc=kwargs["min_length_sc"],
+                    minmax_num_atoms=kwargs["minmax_num_atoms"],
+                    sm=sm
+                )
+
+                return cls(
+                    battery_id=battery_id,
+                    hop_cutoff=hop_cutoff,
+                    entries_for_generation=grouped_entries,
+                    working_ion_entry=working_ion_entry,
+                    migration_graph=migration_graph,
+                    matrix_supercell_structure=host_sc,
+                    conversion_matrix=sc_mat,
+                    inserted_ion_coords=coords_dict,
+                    insert_coords_combo=combo,
+                    **kwargs
+                )
+
+            else:
+                raise TypeError("Please make sure to have kwargs min_length_sc and minmax_num_atoms if populate_sc_fields is set to True.")
+
+    @staticmethod
+    def generate_sc_fields(
+        mg: MigrationGraph,
+        min_length_sc: float,
+        minmax_num_atoms: Tuple[int, int],
+        sm: StructureMatcher
+    ):
+        min_length_sc = min_length_sc
+        minmax_num_atoms = minmax_num_atoms
+
+        sc_mat = get_sc_fromstruct(
+            base_struct=mg.structure,
+            min_atoms=minmax_num_atoms[0],
+            max_atoms=minmax_num_atoms[1],
+            min_length=min_length_sc
         )
+
+        sc_mat = sc_mat.tolist()
+        host_sc = mg.host_structure * sc_mat
+
+        coords_dict = MigrationGraphDoc.ordered_sc_site_dict(mg.only_sites, sc_mat)
+        combo, coords_dict = MigrationGraphDoc.get_hop_sc_combo(mg.unique_hops, sc_mat, sm, host_sc, coords_dict)
+
+        return host_sc, sc_mat, min_length_sc, minmax_num_atoms, coords_dict, combo
+
+    def ordered_sc_site_dict(
+        uc_sites_only: Structure,
+        sc_mat: List[List[Union[int, float]]]
+    ):
+        uc_no_site = uc_sites_only.copy()
+        uc_no_site.remove_sites(range(len(uc_sites_only)))
+        working_ion = uc_sites_only[0].species_string
+        sc_site_dict = {}  # type: dict
+
+        for i, e in enumerate(uc_sites_only):
+            uc_one_set = uc_no_site.copy()
+            uc_one_set.insert(0, working_ion, e.frac_coords)
+            sc_one_set = uc_one_set * sc_mat
+            for index in (range(len(sc_one_set))):
+                sc_site_dict[len(sc_site_dict) + 1] = {"uc_site_type": i, "site": sc_one_set[index]}
+
+        ordered_site_dict = {i: e for i, e in enumerate(sorted(sc_site_dict.values(), key=lambda v: float(np.linalg.norm(v["site"].frac_coords))))}
+        return ordered_site_dict
+
+    def get_hop_sc_combo(
+        unique_hops: Dict,
+        sc_mat: List[List[Union[int, float]]],
+        sm: StructureMatcher,
+        host_sc: Structure,
+        ordered_sc_site_dict: dict
+    ):
+        combo = []
+        working_ion = ordered_sc_site_dict[0]["site"].species_string
+
+        unique_hops = {k: v for k, v in sorted(unique_hops.items())}
+        for one_hop in unique_hops.values():
+            added = False
+            sc_isite_set = {k: v for k, v in ordered_sc_site_dict.items() if v["uc_site_type"] == one_hop["iindex"]}
+            sc_esite_set = {k: v for k, v in ordered_sc_site_dict.items() if v["uc_site_type"] == one_hop["eindex"]}
+            for sc_iindex, sc_isite in sc_isite_set.items():
+                for sc_eindex, sc_esite in sc_esite_set.items():
+                    sc_check = host_sc.copy()
+                    sc_check.insert(0, working_ion, sc_isite['site'].frac_coords)
+                    sc_check.insert(1, working_ion, sc_esite['site'].frac_coords)
+                    if MigrationGraphDoc.compare_sc_one_hop(one_hop, sc_mat, sm, host_sc, sc_check, working_ion, (sc_isite["uc_site_type"], sc_esite["uc_site_type"])):
+                        combo.append(f"{sc_iindex}+{sc_eindex}")
+                        added = True
+                        break
+                if added:
+                    break
+
+            if not added:
+                new_combo, ordered_sc_site_dict = MigrationGraphDoc.append_new_site(host_sc, ordered_sc_site_dict, one_hop, sc_mat)
+                combo.append(new_combo)
+
+        return combo, ordered_sc_site_dict
+
+    def compare_sc_one_hop(
+        one_hop: Dict,
+        sc_mat: List,
+        sm: StructureMatcher,
+        host_sc: Structure,
+        sc_check: Structure,
+        working_ion: str,
+        uc_site_types: Tuple[int, int]
+    ):
+        sc_mat_inv = np.linalg.inv(sc_mat)
+        convert_sc_icoords = np.dot(one_hop["ipos"], sc_mat_inv)
+        convert_sc_ecoords = np.dot(one_hop["epos"], sc_mat_inv)
+        convert_sc = host_sc.copy()
+        convert_sc.insert(0, working_ion, convert_sc_icoords)
+        convert_sc.insert(1, working_ion, convert_sc_ecoords)
+
+        if sm.fit(convert_sc, sc_check):
+            one_hop_dis = one_hop["hop"].length
+            sc_check_hop_dis = np.linalg.norm(sc_check[0].coords - sc_check[1].coords)
+            if np.isclose(one_hop_dis, sc_check_hop_dis, rtol=0.1, atol=0.1):
+                if one_hop["iindex"] == uc_site_types[0] and one_hop["eindex"] == uc_site_types[1]:
+                    return True
+
+        return False
+
+    def append_new_site(
+        host_sc: Structure,
+        ordered_sc_site_dict: Dict,
+        one_hop: Dict,
+        sc_mat: List[List[Union[int, float]]]
+    ):
+        sc_mat_inv = np.linalg.inv(sc_mat)
+        sc_ipos = np.dot(one_hop["ipos"], sc_mat_inv)
+        sc_epos = np.dot(one_hop["epos"], sc_mat_inv)
+        sc_iindex, sc_eindex = None, None
+        host_sc_insert = host_sc.copy()
+
+        for k, v in ordered_sc_site_dict.items():
+            if np.allclose(sc_ipos, v["site"].frac_coords, rtol=0.1, atol=0.1):
+                sc_iindex = k
+            if np.allclose(sc_epos, v["site"].frac_coords, rtol=0.1, atol=0.1):
+                sc_eindex = k
+
+        if sc_iindex is None:
+            host_sc_insert.insert(0, ordered_sc_site_dict[0]["site"].species_string, sc_ipos)
+            ordered_sc_site_dict[len(ordered_sc_site_dict)] = {
+                "uc_site_type": one_hop["iindex"],
+                "site": host_sc_insert[0],
+                "extra_site": True
+            }
+            sc_iindex = len(ordered_sc_site_dict) - 1
+        if sc_eindex is None:
+            host_sc_insert.insert(0, ordered_sc_site_dict[0]["site"].species_string, sc_epos)
+            ordered_sc_site_dict[len(ordered_sc_site_dict)] = {
+                "uc_site_type": one_hop["eindex"],
+                "site": host_sc_insert[0],
+                "extra_site": True
+            }
+            sc_eindex = len(ordered_sc_site_dict) - 1
+
+        return f"{sc_iindex}+{sc_eindex}", ordered_sc_site_dict

tests/emmet-core/mobility/test_migrationgraph.py

@@ -1,5 +1,7 @@
+import numpy as np
 import pytest
 from monty.serialization import loadfn
+from pymatgen.analysis.structure_matcher import StructureMatcher
 from pymatgen.entries.computed_entries import ComputedEntry
 from emmet.core.mobility.migrationgraph import MigrationGraphDoc
 
@@ -10,39 +12,77 @@ def get_entries(test_dir):
     entries = loadfn(test_dir / "mobility/LiMnP2O7_batt.json")
     return (entries, entry_Li)
 
+
 @pytest.fixture(scope="session")
-def migration_graph(test_dir):
+def migration_graph_prop():
     """
     set the expected parameters for the migrationgraph
     """
     expected_properties = {
-        "LiMnP2O7":{
+        "LiMnP2O7": {
             "max_distance": 5,
             "num_uhops": 8,
             "longest_hop": 4.92647,
-            "shortest_hop": 2.77240
+            "shortest_hop": 2.77240,
+            "min_length_sc": 7,
+            "minmax_num_atoms": (80, 160)
         }
     }
-
     return expected_properties
 
 
-def test_from_entries_and_distance(migration_graph, get_entries):
-    for expected in migration_graph.values():
+@pytest.fixture(scope="session")
+def mg_for_sc_fields(test_dir):
+    """
+    get MigrationGraph object generated with methods from pymatgen.analysis.diffusion for testing generate_sc_fields
+    """
+    mg_for_sc = loadfn(test_dir / "mobility/mg_for_sc.json")
+    return mg_for_sc
+
+
+def test_from_entries_and_distance(migration_graph_prop, get_entries):
+    for expected in migration_graph_prop.values():
         mgdoc = MigrationGraphDoc.from_entries_and_distance(
             battery_id="mp-1234",
             grouped_entries=get_entries[0],
             working_ion_entry=get_entries[1],
-            hop_cutoff=5
+            hop_cutoff=5,
+            populate_sc_fields=True,
+            min_length_sc=7,
+            minmax_num_atoms=(80, 160)
         )
 
         mg = mgdoc.migration_graph
         res_d = {
             "max_distance": mgdoc.hop_cutoff,
             "num_uhops": len(mg.unique_hops),
             "longest_hop": sorted(mg.unique_hops.items(), key=lambda x: x[1]["hop_distance"])[-1][1]["hop_distance"],
-            "shortest_hop": sorted(mg.unique_hops.items(), key=lambda x: x[1]["hop_distance"])[0][1]["hop_distance"]
+            "shortest_hop": sorted(mg.unique_hops.items(), key=lambda x: x[1]["hop_distance"])[0][1]["hop_distance"],
+            "min_length_sc": mgdoc.min_length_sc,
+            "minmax_num_atoms": mgdoc.minmax_num_atoms
         }
         for k, v in expected.items():
-            print(res_d[k], pytest.approx(v, 0.01))
-            assert res_d[k] == pytest.approx(v, 0.01)
+            assert res_d[k] == pytest.approx(v, 0.01)
+
+
+def test_generate_sc_fields(mg_for_sc_fields):
+    sm = StructureMatcher()
+    host_sc, sc_mat, min_length, min_max_num_atoms, coords_dict, combo = MigrationGraphDoc.generate_sc_fields(mg_for_sc_fields, 7, (80, 160), sm)
+    sc_mat_inv = np.linalg.inv(sc_mat)
+    expected_sc_list = []
+
+    for one_hop in mg_for_sc_fields.unique_hops.values():
+        host_sc_insert = host_sc.copy()
+        host_sc_insert.insert(0, "Li", np.dot(one_hop["ipos"], sc_mat_inv))
+        host_sc_insert.insert(0, "Li", np.dot(one_hop["epos"], sc_mat_inv))
+        expected_sc_list.append(host_sc_insert)
+
+    for one_combo in combo:
+        hop_sc = host_sc.copy()
+        sc_iindex, sc_eindex = list(map(int, one_combo.split("+")))
+        sc_isite = coords_dict[sc_iindex]["site"]
+        sc_esite = coords_dict[sc_eindex]["site"]
+        hop_sc.insert(0, "Li", sc_isite.frac_coords)
+        hop_sc.insert(0, "Li", sc_esite.frac_coords)
+        check_sc_list = [sm.fit(hop_sc, check_sc) for check_sc in expected_sc_list]
+        assert sum(check_sc_list) >= 1