Repository: WMD-group/SMACT

Issue:
Custom oxidation state calls
Currently, if someone wants to use a custom set of oxidation states, they would need to rewrite part of the codebase in order to achieve this functionality. 
To make some SMACT functions more extensible to custom oxidation states lists, we should add a function to call custom oxidation states. An obvious area where this can be impactful is the `smact_filter` function where a user may want to use their own set of oxidation states to generate SMACT-plausible compositions. The user could either load in a set of oxidation states from a file (in which case we need to specify a file format for compatibility) or load from an object in memory (e.g. a dictionary, or a dataframe etc.). 

Discussion:

Patch:
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 9afc9ae2..3ff89c33 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -4,11 +4,12 @@ repos:
     hooks:
       - id: isort
         additional_dependencies: [toml]
-        args: ["--profile", "black", "--filter-files"]
+        args: ["--profile", "black", "--filter-files","--line-length=80"]
   - repo: https://github.com/psf/black
     rev: "23.1.0"
     hooks:
       - id: black-jupyter
+        args: [--line-length=80]
   - repo: https://github.com/asottile/pyupgrade
     rev: v3.3.1
     hooks:
diff --git a/examples/Counting/ElementCombinationsParallel.py b/examples/Counting/ElementCombinationsParallel.py
index 57856c95..d19dede1 100644
--- a/examples/Counting/ElementCombinationsParallel.py
+++ b/examples/Counting/ElementCombinationsParallel.py
@@ -165,8 +165,7 @@ def main():
         )
 
         print(
-            "Counting ({} element combinations)"
-            "...".format(combination_count)
+            "Counting ({} element combinations)" "...".format(combination_count)
         )
 
         # Combinations are counted in chunks set by count_progress_interval.
@@ -237,8 +236,7 @@ def main():
 
             time_elapsed = time.time() - start_time
             time_remaining = (
-                combination_count * (time_elapsed / data_pointer)
-                - time_elapsed
+                combination_count * (time_elapsed / data_pointer) - time_elapsed
             )
 
             print_status(
diff --git a/examples/Oxidation_states/oxidation_states.ipynb b/examples/Oxidation_states/oxidation_states.ipynb
index ca2bfdf6..588d1fdf 100644
--- a/examples/Oxidation_states/oxidation_states.ipynb
+++ b/examples/Oxidation_states/oxidation_states.ipynb
@@ -282,12 +282,10 @@
     "# Here we grab the species string for each composition generated by smact\n",
     "list_of_species = [species[3] for species in flat_list]\n",
     "A_species = [\n",
-    "    f\"{species[0].symbol}{species[0].oxidation}+\"\n",
-    "    for species in list_of_species\n",
+    "    f\"{species[0].symbol}{species[0].oxidation}+\" for species in list_of_species\n",
     "]\n",
     "B_species = [\n",
-    "    f\"{species[1].symbol}{species[1].oxidation}+\"\n",
-    "    for species in list_of_species\n",
+    "    f\"{species[1].symbol}{species[1].oxidation}+\" for species in list_of_species\n",
     "]\n",
     "X_species = [f\"{species[2].symbol}1-\" for species in list_of_species]\n",
     "\n",
diff --git a/examples/Structure_Prediction/Li-Garnet_Generator.ipynb b/examples/Structure_Prediction/Li-Garnet_Generator.ipynb
index 0b64e0a2..36ad2cf5 100644
--- a/examples/Structure_Prediction/Li-Garnet_Generator.ipynb
+++ b/examples/Structure_Prediction/Li-Garnet_Generator.ipynb
@@ -285,8 +285,7 @@
     "    sus_factor = 0\n",
     "    for i in Composition(comp).elements:\n",
     "        sus_factor += (\n",
-    "            Composition(comp).get_wt_fraction(i)\n",
-    "            * smact.Element(i.symbol).HHI_r\n",
+    "            Composition(comp).get_wt_fraction(i) * smact.Element(i.symbol).HHI_r\n",
     "        )\n",
     "    return sus_factor\n",
     "\n",
diff --git a/smact/__init__.py b/smact/__init__.py
index ea8eb597..e19dc51b 100644
--- a/smact/__init__.py
+++ b/smact/__init__.py
@@ -59,6 +59,8 @@ class Element:
 
         Element.oxidation_states_wiki (list): List of oxidation states that appear wikipedia (https://en.wikipedia.org/wiki/Template:List_of_oxidation_states_of_the_elements) Data retrieved: 2022-09-22
 
+        Element.oxidation_states_custom (list | None ): List of oxidation states that appear in the custom data file supplied (if any)
+
         Element.coord_envs (list): The allowed coordination enviroments for the ion
 
         Element.covalent_radius (float) : Covalent radius of the element
@@ -77,15 +79,35 @@ class Element:
 
     """
 
-    def __init__(self, symbol: str):
+    def __init__(
+        self, symbol: str, oxi_states_custom_filepath: Optional[str] = None
+    ):
         """Initialise Element class
 
         Args:
             symbol (str): Chemical element symbol (e.g. 'Fe')
+            oxi_states_custom_filepath (str): Path to custom oxidation states file
 
         """
+        # Get the oxidation states from the custom file if it exists
+        if oxi_states_custom_filepath:
+            try:
+                self._oxidation_states_custom = (
+                    data_loader.lookup_element_oxidation_states_custom(
+                        symbol, oxi_states_custom_filepath
+                    )
+                )
+                self.oxidation_states_custom = self._oxidation_states_custom
+            except TypeError:
+                warnings.warn(
+                    "Custom oxidation states file not found. Please check the file path."
+                )
+                self.oxidation_states_custom = None
+        else:
+            self.oxidation_states_custom = None
+        self.symbol = symbol
 
-        dataset = data_loader.lookup_element_data(symbol, copy=False)
+        dataset = data_loader.lookup_element_data(self.symbol, copy=False)
 
         if dataset == None:
             raise NameError(f"Elemental data for {symbol} not found.")
@@ -307,7 +329,10 @@ def ordered_elements(x: int, y: int) -> List[str]:
     return ordered_elements
 
 
-def element_dictionary(elements: Optional[Iterable[str]] = None):
+def element_dictionary(
+    elements: Optional[Iterable[str]] = None,
+    oxi_states_custom_filepath: Optional[str] = None,
+):
     """
     Create a dictionary of initialised smact.Element objects
 
@@ -317,13 +342,22 @@ def element_dictionary(elements: Optional[Iterable[str]] = None):
     Args:
         elements (iterable of strings) : Elements to include. If None,
             use all elements up to 103.
+        oxi_states_custom_filepath (str): Path to custom oxidation states file
+
+
     Returns:
         dict: Dictionary with element symbols as keys and smact.Element
             objects as data
     """
     if elements == None:
         elements = ordered_elements(1, 103)
-    return {symbol: Element(symbol) for symbol in elements}
+    if oxi_states_custom_filepath:
+        return {
+            symbol: Element(symbol, oxi_states_custom_filepath)
+            for symbol in elements
+        }
+    else:
+        return {symbol: Element(symbol) for symbol in elements}
 
 
 def are_eq(A: list, B: list, tolerance: float = 1e-4):
diff --git a/smact/data_loader.py b/smact/data_loader.py
index 79631a8e..79e0dd8a 100644
--- a/smact/data_loader.py
+++ b/smact/data_loader.py
@@ -269,6 +269,59 @@ def lookup_element_oxidation_states_wiki(symbol, copy=True):
         return None
 
 
+_el_ox_states_custom = None
+
+
+def lookup_element_oxidation_states_custom(symbol, filepath, copy=True):
+    """
+    Retrieve a list of known oxidation states for an element.
+    The oxidation states list is specified by the user in a text file.
+
+    Args:
+        symbol (str) : the atomic symbol of the element to look up.
+        copy (Optional(bool)): if True (default), return a copy of the
+            oxidation-state list, rather than a reference to the cached
+            data -- only use copy=False in performance-sensitive code
+            and where the list will not be modified!
+
+    Returns:
+        list: List of known oxidation states for the element.
+
+            Return None if oxidation states for the Element were not
+            found in the external data.
+    """
+
+    global _el_ox_states_custom
+
+    if _el_ox_states_custom is None:
+        _el_ox_states_custom = {}
+
+        for items in _get_data_rows(filepath):
+            _el_ox_states_custom[items[0]] = [
+                int(oxidationState) for oxidationState in items[1:]
+            ]
+
+    if symbol in _el_ox_states_custom:
+        if copy:
+            # _el_ox_states_custom stores lists -> if copy is set, make an implicit
+            # deep copy.  The elements of the lists are integers, which are
+            # "value types" in Python.
+
+            return [
+                oxidationState
+                for oxidationState in _el_ox_states_custom[symbol]
+            ]
+        else:
+            return _el_ox_states_custom[symbol]
+    else:
+        if _print_warnings:
+            print(
+                "WARNING: Oxidation states for element {} "
+                "not found.".format(symbol)
+            )
+        return None
+
+
 # Loader and cache for the element HHI scores.
 
 _element_hhis = None
@@ -367,9 +420,7 @@ def lookup_element_data(symbol, copy=True):
             # or, if not clearly a number, to None
             clean_items = items[0:2] + list(map(float_or_None, items[2:]))
 
-            _element_data.update(
-                {items[0]: dict(list(zip(keys, clean_items)))}
-            )
+            _element_data.update({items[0]: dict(list(zip(keys, clean_items)))})
 
     if symbol in _element_data:
         if copy:
@@ -384,9 +435,7 @@ def lookup_element_data(symbol, copy=True):
             return _element_data[symbol]
     else:
         if _print_warnings:
-            print(
-                "WARNING: Elemental data for {}" " not found.".format(symbol)
-            )
+            print("WARNING: Elemental data for {}" " not found.".format(symbol))
             print(_element_data)
         return None
 
@@ -471,9 +520,7 @@ def lookup_element_shannon_radius_data(symbol, copy=True):
             # function on each element.
             # The dictionary values are all Python "value types", so
             # nothing further is required to make a deep copy.
-            return [
-                item.copy() for item in _element_shannon_radii_data[symbol]
-            ]
+            return [item.copy() for item in _element_shannon_radii_data[symbol]]
         else:
             return _element_shannon_radii_data[symbol]
     else:
diff --git a/smact/screening.py b/smact/screening.py
index 13db4930..1931e0ab 100644
--- a/smact/screening.py
+++ b/smact/screening.py
@@ -1,10 +1,14 @@
 import itertools
+import os
 import warnings
 from collections import namedtuple
 from itertools import combinations
 from typing import Iterable, List, Optional, Tuple, Union
 
 from smact import Element, neutral_ratios
+from smact.data_loader import (
+    lookup_element_oxidation_states_custom as oxi_custom,
+)
 
 # Use named tuple to improve readability of smact_filter outputs
 _allowed_compositions = namedtuple(
@@ -324,7 +328,7 @@ def smact_filter(
         threshold (int): Threshold for stoichiometry limit, default = 8
         stoichs (list[int]): A selection of valid stoichiometric ratios for each site.
         species_unique (bool): Whether or not to consider elements in different oxidation states as unique in the results.
-        oxidation_states_set (string): A string to choose which set of oxidation states should be chosen. Options are 'default', 'icsd', 'pymatgen' and 'wiki' for the default, icsd, pymatgen structure predictor and Wikipedia (https://en.wikipedia.org/wiki/Template:List_of_oxidation_states_of_the_elements) oxidation states respectively.
+        oxidation_states_set (string): A string to choose which set of oxidation states should be chosen. Options are 'default', 'icsd', 'pymatgen' and 'wiki' for the default, icsd, pymatgen structure predictor and Wikipedia (https://en.wikipedia.org/wiki/Template:List_of_oxidation_states_of_the_elements) oxidation states respectively. A filepath to an oxidation states text file can also be supplied as well.
         comp_tuple (bool): Whether or not to return the results as a named tuple of elements and stoichiometries (True) or as a normal tuple of elements and stoichiometries (False).
     Returns:
         allowed_comps (list): Allowed compositions for that chemical system
@@ -371,10 +375,12 @@ def smact_filter(
     }
     if oxidation_states_set in oxi_set:
         ox_combos = oxi_set[oxidation_states_set]
+    elif os.path.exists(oxidation_states_set):
+        ox_combos = [oxi_custom(e.symbol, oxidation_states_set) for e in els]
     else:
         raise (
             Exception(
-                f'{oxidation_states_set} is not valid. Enter either "default", "icsd", "pymatgen" or "wiki" for oxidation_states_set.'
+                f'{oxidation_states_set} is not valid. Enter either "default", "icsd", "pymatgen","wiki" or a filepath to a textfile of oxidation states.'
             )
         )
     if oxidation_states_set == "wiki":
diff --git a/smact/structure_prediction/database.py b/smact/structure_prediction/database.py
index 72001c03..42f0d764 100644
--- a/smact/structure_prediction/database.py
+++ b/smact/structure_prediction/database.py
@@ -202,9 +202,7 @@ class StructureDB:
 
         return num
 
-    def get_structs(
-        self, composition: str, table: str
-    ) -> List[SmactStructure]:
+    def get_structs(self, composition: str, table: str) -> List[SmactStructure]:
         """Get SmactStructures for a given composition.
 
         Args:
diff --git a/smact/structure_prediction/mutation.py b/smact/structure_prediction/mutation.py
index 9ca2e055..41667506 100644
--- a/smact/structure_prediction/mutation.py
+++ b/smact/structure_prediction/mutation.py
@@ -373,8 +373,7 @@ class CationMutator:
     def cond_sub_prob(self, s1: str, s2: str) -> float:
         """Calculate the probability of substitution of one species with another."""
         return (
-            np.exp(self.get_lambda(s1, s2))
-            / np.exp(self.get_lambdas(s2)).sum()
+            np.exp(self.get_lambda(s1, s2)) / np.exp(self.get_lambdas(s2)).sum()
         )
 
     def cond_sub_probs(self, s1: str) -> pd.Series:
diff --git a/smact/structure_prediction/prediction.py b/smact/structure_prediction/prediction.py
index 85741ef5..5fedfd28 100644
--- a/smact/structure_prediction/prediction.py
+++ b/smact/structure_prediction/prediction.py
@@ -183,8 +183,7 @@ class StructurePredictor:
         sub_species = list(map(list, sub_species))
 
         potential_nary_parents: List[List[SmactStructure]] = list(
-            self.db.get_with_species(specs, self.table)
-            for specs in sub_species
+            self.db.get_with_species(specs, self.table) for specs in sub_species
         )
 
         for spec_idx, parents in enumerate(potential_nary_parents):
diff --git a/smact/structure_prediction/structure.py b/smact/structure_prediction/structure.py
index e9993774..0f072b35 100644
--- a/smact/structure_prediction/structure.py
+++ b/smact/structure_prediction/structure.py
@@ -436,10 +436,7 @@ class SmactStructure:
         lattice_param = float(lines[1])
 
         lattice = np.array(
-            [
-                [float(point) for point in line.split(" ")]
-                for line in lines[2:5]
-            ]
+            [[float(point) for point in line.split(" ")] for line in lines[2:5]]
         )
 
         sites = defaultdict(list)


Test patch:
diff --git a/smact/tests/files/test_oxidation_states.txt b/smact/tests/files/test_oxidation_states.txt
new file mode 100644
index 00000000..122f9ee0
--- /dev/null
+++ b/smact/tests/files/test_oxidation_states.txt
@@ -0,0 +1,108 @@
+H -1 +1
+He
+Li +1
+Be +1 +2
+B +1 +2 +3
+C -4 -3 -2 -1 +1 +2 +3 +4
+N -3 -2 -1 +1 +2 +3 +4 +5
+O -2 -1 +1 +2
+F -1
+Ne
+Na -1 +1
+Mg +1 +2
+Al +1 +2 +3
+Si -4 -3 -2 -1 +1 +2 +3 +4
+P -3 -2 -1 +1 +2 +3 +4 +5
+S -2 -1	+1 +2 +3 +4 +5 +6
+Cl -1 +1 +2 +3 +4 +5 +6 +7
+Ar
+K -1 +1
+Ca +1 +2
+Sc +1 +2 +3
+Ti -1 +1 +2 +3 +4
+V -1 +1 +2 +3 +4 +5
+Cr -2 -1 +1 +2 +3 +4 +5 +6
+Mn -3 -2 -1 +1 +2 +3 +4 +5 +6 +7
+Fe -2 -1 +1 +2 +3 +4 +5 +6
+Co -1 +1 +2 +3 +4 +5
+Ni -1 +1 +2 +3 +4
+Cu +1 +2 +3 +4
+Zn +1 +2
+Ga +1 +2 +3
+Ge -4 -3 -2 -1 +1 +2 +3 +4
+As -3 +1 +2 +3 +5
+Se -2 +1 +2 +4 +6
+Br -1 +1 +2 +3 +4 +5 +7
+Kr +2
+Rb -1 +1
+Sr +1 +2
+Y +1 +2 +3
+Zr +1 +2 +3 +4
+Nb -1 +1 +2 +3 +4 +5
+Mo -2 -1 +1 +2 +3 +4 +5 +6
+Tc -3 -1 +1 +2 +3 +4 +5 +6 +7
+Ru -2 +1 +2 +3 +4 +5 +6 +7 +8
+Rh -1 +1 +2 +3 +4 +5 +6
+Pd +1 +2 +4 +6
+Ag +1 +2 +3 +4
+Cd +1 +2
+In +1 +2 +3
+Sn -4 +2 +4
+Sb -3 +3 +5
+Te -2 +2 +4 +5 +6
+I -1 +1 +3 +4 +5 +7
+Xe +1 +2 +4 +6 +8
+Cs -1 +1
+Ba +2
+La +2 +3
+Ce +2 +3 +4
+Pr +2 +3 +4
+Nd +2 +3 +4
+Pm +2 +3
+Sm +2 +3
+Eu +2 +3
+Gd +1 +2 +3
+Tb +1 +2 +3 +4
+Dy +2 +3 +4
+Ho +2 +3
+Er +2 +3
+Tm +2 +3
+Yb +2 +3
+Lu +3
+Hf +2 +3 +4
+Ta -1 +2 +3 +4 +5
+W -2 -1 +1 +2 +3 +4 +5 +6
+Re -3 -1 +1 +2 +3 +4 +5 +6 +7
+Os -2 -1 +1 +2 +3 +4 +5 +6 +7 +8
+Ir -3 -1 +1 +2 +3 +4 +5 +6 +7 +8
+Pt -2 -1 +1 +2 +3 +4 +5 +6
+Au -1 +1 +2 +3 +5
+Hg +1 +2 +4
+Tl -1 +1 +3
+Pb -4 +2 +4
+Bi -3 +1 +3 +5 +7
+Po -2 +2 +4 +5 +6
+At -1 +1 +3 +5 +7
+Rn +2 +6
+Fr +1
+Ra +2
+Ac +2 +3
+Th +2 +3 +4
+Pa +2 +3 +4 +5
+U +2 +3 +4 +5 +6
+Np +3 +4 +5 +6 +7
+Pu +2 +3 +4 +5 +6 +7 +8
+Am +2 +3 +4 +5 +6 +7
+Cm +2 +3 +4 +6 +8
+Bk +2 +3 +4
+Cf +2 +3 +4
+Es +2 +3 +4
+Fm +2 +3
+Md +2 +3
+No +2 +3
+Lr +3
+Rf +4
+Db +5
+Sg +6
+Bh +7
+Hs +8
diff --git a/smact/tests/test_core.py b/smact/tests/test_core.py
index ccfb2768..b895bc53 100755
--- a/smact/tests/test_core.py
+++ b/smact/tests/test_core.py
@@ -15,12 +15,17 @@ from smact import Species
 from smact.builder import wurtzite
 from smact.properties import band_gap_Harrison, compound_electroneg
 
+files_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), "files")
+TEST_OX_STATES = os.path.join(files_dir, "test_oxidation_states.txt")
+
 
 class TestSequenceFunctions(unittest.TestCase):
     # ---------------- TOP-LEVEL ----------------
 
     def test_Element_class_Pt(self):
-        Pt = smact.Element("Pt")
+        Pt = smact.Element(
+            "Pt",
+        )
         self.assertEqual(Pt.name, "Platinum")
         self.assertEqual(Pt.ionpot, 8.95883)
         self.assertEqual(Pt.number, 78)
@@ -34,9 +39,10 @@ class TestSequenceFunctions(unittest.TestCase):
 
     def test_element_dictionary(self):
         newlist = ["O", "Rb", "W"]
-        dictionary = smact.element_dictionary(newlist)
+        dictionary = smact.element_dictionary(newlist, TEST_OX_STATES)
         self.assertEqual(dictionary["O"].crustal_abundance, 461000.0)
         self.assertEqual(dictionary["Rb"].oxidation_states, [-1, 1])
+        self.assertEqual(dictionary["Rb"].oxidation_states_custom, [-1, 1])
         self.assertEqual(dictionary["W"].name, "Tungsten")
         self.assertTrue("Rn" in smact.element_dictionary())
 
@@ -46,9 +52,7 @@ class TestSequenceFunctions(unittest.TestCase):
                 [1.00, 2.00, 3.00], [1.001, 1.999, 3.00], tolerance=1e-2
             )
         )
-        self.assertFalse(
-            smact.are_eq([1.00, 2.00, 3.00], [1.001, 1.999, 3.00])
-        )
+        self.assertFalse(smact.are_eq([1.00, 2.00, 3.00], [1.001, 1.999, 3.00]))
 
     def test_gcd_recursive(self):
         self.assertEqual(smact._gcd_recursive(4, 12, 10, 32), 2)
@@ -346,6 +350,12 @@ class TestSequenceFunctions(unittest.TestCase):
                 (("Na", "Fe", "Cl"), (1, 1, -1), (1, 1, 2)),
             ],
         )
+        self.assertEqual(
+            result,
+            smact.screening.smact_filter(
+                [Na, Fe, Cl], threshold=2, oxidation_states_set=TEST_OX_STATES
+            ),
+        )
         result_comp_tuple = smact.screening.smact_filter(
             [Na, Fe, Cl], threshold=2, comp_tuple=True
         )
diff --git a/smact/tests/test_structure.py b/smact/tests/test_structure.py
index 8ebc32dc..d92253bc 100644
--- a/smact/tests/test_structure.py
+++ b/smact/tests/test_structure.py
@@ -363,9 +363,7 @@ class CationMutatorTest(unittest.TestCase):
         """Set up the test initial structure and mutator."""
         cls.test_struct = SmactStructure.from_file(TEST_POSCAR)
 
-        cls.test_mutator = CationMutator.from_json(
-            lambda_json=TEST_LAMBDA_JSON
-        )
+        cls.test_mutator = CationMutator.from_json(lambda_json=TEST_LAMBDA_JSON)
         cls.test_pymatgen_mutator = CationMutator.from_json(
             lambda_json=None, alpha=lambda x, y: -5
         )