make edits

compressed the cifs into fstar.tar.gz. Implemented type hints. reworked to simplify and run better.

Author: jonathanjdenney

pymatgen/analysis/fstar/fstar.py

@@ -5,15 +5,15 @@
 atomic form factors. Rev. Sci. Instrum. 89, 093002 (2018). 10.1063/1.5044555
 Yin, L. et al.Thermodynamics of Antisite Defects in Layered NMC Cathodes: Systematic Insights from High-Precision Powder
 Diffraction Analyses Chem. Mater 2020 32 (3), 1002-1010. 10.1021/acs.chemmater.9b03646
-
 """
-from __future__ import annotations
 
 import os
-
 import numpy as np
 import pandas as pd
 import plotly.express as px
+from typing import Callable
+from pymatgen.core.periodic_table import Element
+from pymatgen.core.structure import Structure
 
 # Load in the neutron form factors
 with open(os.path.join(os.path.dirname(__file__), "neutron_factors.csv")) as f:
@@ -23,64 +23,84 @@
 
 class FStarDiagram:
     """
-    Take a list of structure objects and/or cifs and use them to generate an f* phase diagram.
+    Take a list of symmetrized structure objects and use them to generate an f* phase diagram.
     """
-
-    def __init__(self, structures, scattering_type="X-ray", custom_scatter=None):
+    def __init__(self, structures: list[Structure], scattering_type: str='X-ray'):
         """
         Initialize the f* diagram generator with the list of structures and scattering type.
 
         Args:
             structures(list): List of structure objects to use in the diagram. These MUST be symmetrized structure
-                objects.
+                objects. All structures must only vary in the occupancy of the sites.
             scattering_type(str): Type of scattering to use in the f* calculation. Defaults to 'X-ray'
                 which uses the atomic number as the scattering factor. 'Neutron' is a built in scattering
-                type which uses neutron scattering factors. 'Custom' allows the user to supplement their
-                own calculation with any set of scattering factors.
-            custom_scatter(function): when using custom scattering set this equal to a global varialble that is equal
-                to the custom scattering function.
+                type which uses neutron scattering factors.
         """
-
+        # check if the input structures list is valid
+        for ind, struct in enumerate(structures):
+            try:
+                if struct.equivalent_indices == structures[ind-1].equivalent_indices:
+                    continue
+                else:
+                    raise ValueError(
+                        "All structues must only vary in occupancy."
+                    )
+            except AttributeError:
+                raise AttributeError(
+                    "Must use symmeteized structure objects"
+                )        
         self._structures = structures
         self._scatter = scattering_type
-        self._custscat = custom_scatter
-        self._equiv_inds = [struct.equivalent_indices for struct in self._structures]
-        self.site_labels = self.get_site_labels()
-        self.coords = self.get_fstar_coords()
-        self.plot = px.scatter_ternary(
-            data_frame=self.coords, a=self.site_labels[0], b=self.site_labels[1], c=self.site_labels[2]
-        )
+        self._scatter_dict = {'X-ray':self.xray_scatter,'Neutron':self.neutron_scatter}
+        self.site_labels = self._get_site_labels()
+        self.fstar_coords = self._get_fstar_coords()
+        self.set_plot_list([self.site_labels[0], self.site_labels[1], self.site_labels[2]])
+        self.make_plot()
         print("The labels for this structure's unique sites are")
         print(self.site_labels)
-
-    def edit_fstar_diagram(self, combine_list=False, plot_list=False, **kwargs):
+        
+    def combine_sites(self, site_lists: list[list[str]]) -> None:
         """
-        Edit the plot of the f* diagram using plotly express.
-
+        Many structures have more than three sites. If this is the case you may want to
+            add some sites together to make a psudo-site.
         Args:
-            combine_list(list): This is a list of lists which indicates what unique sites need to be combined to make
-                the plot ternary.
-            plot_list(list): This is a list that indicates what unique sites or combined sites to plot and what order to
-            plot them in.
-            kwargs: use this to add any other arguments from scatter_ternary .
-        """
-        if combine_list:
-            for combo in combine_list:
-                self.coords[str(combo)] = sum([self.coords[site] for site in combo])
-                if str(combo) not in self.site_labels:
-                    self.site_labels.append(str(combo))
-        if plot_list:
-            self.plot = px.scatter_ternary(
-                data_frame=self.coords, a=plot_list[0], b=plot_list[1], c=plot_list[2], **kwargs
-            )
-        else:
-            self.plot = px.scatter_ternary(
-                data_frame=self.coords, a=self.site_labels[0], b=self.site_labels[1], c=self.site_labels[2], **kwargs
-            )
-
-    def get_site_labels(self):
+            site_lists(list): A list of lists of site lable strings. This allows you to combine
+                more than one set of sites at once.
+        """
+        for combo in site_lists:
+            for site in combo:
+                if site not in self.site_labels:
+                    raise ValueError(
+                    "All sites must be in the site_labels list"
+                    )
+            self.fstar_coords[str(combo)] = sum([self.fstar_coords[site] for site in combo])
+            if str(combo) not in self.site_labels:
+                self.site_labels.append(str(combo))
+    def set_plot_list(self, site_list: list[str]) -> None:
+        """
+        set the list of sites to plot and the order to plot them in.
+        Args:
+            site_list(list): A list of site lable strings. Index 0 goes on the top of the
+                plot, index 1 goes on the bottom left, and index 2 goes on the bottom right.
+        """
+        for site in site_list:
+            if site not in self.site_labels:
+                raise ValueError(
+                    "All sites must be in the site_labels list"
+                )
+        self.plot_list = site_list
+    def make_plot(self, **kwargs):
         """
-        Generates unique site labels based on composition, order, and symmetry equivalence in the structure object.
+        Makes a plotly express scatter_ternary plot useing the fstar_coords dataframe and the
+            sites in plot list.
+        Args:
+            **kwargs: this can be any argument that the scatter_ternary fucntion can use.
+        """
+        self.plot = px.scatter_ternary(data_frame=self.fstar_coords, a=self.plot_list[0], b=self.plot_list[1],
+                                       c=self.plot_list[2], **kwargs)
+    def _get_site_labels(self):
+        """
+        Generates unique site labels based on composition, order, and symetry equivalence in the structure object.
         Ex:
         Structure Summary
         Lattice
@@ -110,58 +130,53 @@ def get_site_labels(self):
             (0.0000, 0.0000, 7.1375) [0.0000, 0.0000, 0.5000]
         '[0.   0.   0.25]O' - PeriodicSite: O (0.0000, 0.0000, 3.5688) [0.0000, 0.0000, 0.2500]
         """
+        site_labels = []
+        for site in self._structures[0].equivalent_indices:
+            site_labels.append(str(self._structures[0][site[0]].frac_coords) + \
+                    [str(sp) for sp, _ in self._structures[0][site[0]].species.items()][0])
+        return site_labels
 
-        site_labels_fin = []
-        for ind1, struct in enumerate(self._equiv_inds):
-            site_labels = []
-            for _ind2, site in enumerate(struct):
-                label = str(self._structures[ind1][site[0]].frac_coords) + next(
-                    str(sp) for sp, occ in self._structures[ind1][site[0]].species.items()
-                )
-                if label not in site_labels:
-                    site_labels.append(label)
-            if len(site_labels) > len(site_labels_fin):
-                site_labels_fin = site_labels
-        return site_labels_fin
-
-    def get_fstar_coords(self):
+    def _get_fstar_coords(self):
         """
         Calculate the f* coordinates for the list of structures.
         """
-
         fstar_df_full = pd.DataFrame(columns=self.site_labels)
-
-        for ind1, struct in enumerate(self._equiv_inds):
+        for ind1, struct in enumerate(self._structures):
             fstar_df = pd.DataFrame(columns=self.site_labels, data=[[0.0 for i in self.site_labels]])
-            for ind2, site in enumerate(struct):
+            for ind2, site in enumerate(struct.equivalent_indices):
                 occ_f_list = []
                 mult = len(site)
                 site_frac_coord = str(self._structures[ind1][site[0]].frac_coords)
                 column = [label for label in self.site_labels if site_frac_coord in label]
                 elements_and_occupancies = self._structures[ind1][site[0]].species.items()
                 for sp, occ in elements_and_occupancies:
-                    if self._scatter == "X-ray":
-                        f_occ = sp.Z * occ
-                    if self._scatter == "Neutron":
-                        for i, n in enumerate(NEUTRON_SCATTER_DF["Isotope"].values):
-                            if hasattr(sp, "element"):
-                                if n == str(sp.element):
-                                    f_occ = float(NEUTRON_SCATTER_DF.loc[i]["Coh b"]) * occ
-                                    break
-                            else:
-                                if n == str(sp):
-                                    f_occ = float(NEUTRON_SCATTER_DF.loc[i]["Coh b"]) * occ
-                                    break
-                    if self._scatter == "Custom":
-                        if hasattr(sp, "element"):
-                            f_occ = self._custscat(str(sp.element), occ, ind1, ind2)
-                        else:
-                            f_occ = self._custscat(str(sp), occ, ind1, ind2)
+                    # ind1 and ind2 are added in case someone wants to make a custom scatter function
+                    # that uses information in the structure object
+                    f_occ = self._scatter_dict[self._scatter](sp, occ, ind1, ind2) 
                     occ_f_list.append(f_occ)
-
                 fstar = np.absolute(mult * sum(occ_f_list))
                 fstar_df.loc[0][column[0]] = round(float(fstar), 4)
             tot = sum(sum(list(fstar_df.values)))
             fstar_df = pd.DataFrame(columns=self.site_labels, data=[fs / tot for fs in list(fstar_df.values)])
             fstar_df_full = pd.concat([fstar_df_full, fstar_df], ignore_index=True)
         return fstar_df_full
+    def xray_scatter(self, el: Element, occ: float, i1: int, i2: int) -> float:
+        """
+        X-ray scattering function. i2 and i2 are unused.
+        """
+        f_occ = el.Z * occ
+        return f_occ
+    def neutron_scatter(self, el: Element, occ: float, i1: int, i2: int) -> float:
+        """
+        Neutron scattering function. i2 and i2 are unused.
+        """
+        for i, n in enumerate(NEUTRON_SCATTER_DF['Isotope'].values):
+            if hasattr(el, "element"):
+                if n == str(el.element):
+                    f_occ = float(NEUTRON_SCATTER_DF.loc[i]['Coh b']) * occ
+                    break
+            else:
+                if n == str(el):
+                    f_occ = float(NEUTRON_SCATTER_DF.loc[i]['Coh b']) * occ
+                    break
+        return f_occ

tests/analysis/fstar/test_fstar.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-import os
+import tarfile
 
 from pymatgen.analysis.fstar.fstar import FStarDiagram
 from pymatgen.io.cif import CifParser
@@ -9,29 +9,28 @@
 
 class Test_FStarDiagram(PymatgenTest):
     def setUp(self):
-        self.cif_list = [file for file in os.listdir(f"{TEST_FILES_DIR}/analysis/fstar") if file.endswith(".cif")]
+        self.cif_gz = tarfile.open(f"{TEST_FILES_DIR}/fstar/fstar.tar.gz","r")
         self.struct_list = [
-            CifParser(f"{TEST_FILES_DIR}/analysis/fstar/" + file).get_structures(
-                primitive=False, symmetrized=True, check_occu=False
-            )[0]
-            for file in self.cif_list
+            CifParser.from_str(self.cif_gz.extractfile(file).read().decode('utf-8')).get_structures(
+                primitive=False, symmetrized=True, check_occu=False)[0] for file in self.cif_gz.getnames()
         ]
         self.fstar = FStarDiagram(structures=self.struct_list)
 
     def test_edit_fstar_diagram(self):
         assert self.fstar.site_labels == ["[0. 0. 0.]Li", "[0.  0.  0.5]Co", "[0.   0.   0.25]O"]
         new = FStarDiagram(structures=self.struct_list)
         assert self.fstar.plot == new.plot
-        new.edit_fstar_diagram(combine_list=[["[0.  0.  0.5]Co", "[0. 0. 0.]Li"]])
+        new.combine_sites(site_lists=[["[0.  0.  0.5]Co", "[0. 0. 0.]Li"]])
         assert new.site_labels == [
             "[0. 0. 0.]Li",
             "[0.  0.  0.5]Co",
             "[0.   0.   0.25]O",
             "['[0.  0.  0.5]Co', '[0. 0. 0.]Li']",
         ]
-        assert list(new.coords["['[0.  0.  0.5]Co', '[0. 0. 0.]Li']"].to_numpy()) == list(
-            self.fstar.coords["[0. 0. 0.]Li"].to_numpy() + self.fstar.coords["[0.  0.  0.5]Co"].to_numpy()
+        assert list(new.fstar_coords["['[0.  0.  0.5]Co', '[0. 0. 0.]Li']"].to_numpy()) == list(
+            self.fstar.fstar_coords["[0. 0. 0.]Li"].to_numpy() + self.fstar.fstar_coords["[0.  0.  0.5]Co"].to_numpy()
         )
-        assert self.fstar.plot == new.plot
-        new.edit_fstar_diagram(plot_list=["[0.  0.  0.5]Co", "[0.   0.   0.25]O", "[0. 0. 0.]Li"])
+        new.set_plot_list(site_list=["[0.  0.  0.5]Co", "[0.   0.   0.25]O", "[0. 0. 0.]Li"])
+        assert self.fstar.plot_list != new.plot_list
+        new.make_plot()
         assert self.fstar.plot != new.plot