ur-whitelab · Jgmedina95 · May 5, 2025 · Apr 29, 2025 · Apr 29, 2025 · Apr 29, 2025
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -15,3 +15,7 @@ repos:
     rev: "0.6.0"
     hooks:
       - id: nbstripout
+  - repo: https://github.com/pre-commit/mirrors-mypy
+    rev: v1.15.0  # (or whichever version you want)
+    hooks:
+      - id: mypy
diff --git a/docs/source/conf.py b/docs/source/conf.py
@@ -46,7 +46,7 @@
 # List of patterns, relative to source directory, that match files and
 # directories to ignore when looking for source files.
 # This pattern also affects html_static_path and html_extra_path.
-exclude_patterns = []
+exclude_patterns: list = []
 
 
 # -- Options for HTML output -------------------------------------------------

diff --git a/exmol/exmol.py b/exmol/exmol.py
@@ -205,12 +205,14 @@ def _load_smarts(path, rank_cutoff=500):
     return smarts
 
 
-def name_morgan_bit(m: Any, bitInfo: Dict[Any, Any], key: int) -> str:
+def name_morgan_bit(m: Any, bitInfo: Dict[Any, Any], key: int) -> Optional[str]:
     """Get the name of a Morgan bit using a SMARTS dictionary
 
     :param m: RDKit molecule
     :param bitInfo: bitInfo dictionary from rdkit.Chem.AllChem.GetMorganFingerprint
     :param key: bit key corresponding to the fingerprint you want to have named
+
+    :return: Name of the bit, or None if no match is found
        """
     global _SMARTS
     if _SMARTS is None:
@@ -260,7 +262,7 @@ def get_functional_groups(
     :param mol: RDKit molecule
     :param return_all: If True, will return all functional groups found in the molecule
     :param cutoff: Maximum rank of functional groups to consider based on popularity (increase to include groups like methyl, ethyl, etc.)
-    :return: set of unique functional group names present in the molecule
+    :return: set of unique functional group names present in the molecule. If mol is None, returns an empty set.
     """
     global _SMARTS
     if _SMARTS is None:
@@ -273,9 +275,9 @@ def get_functional_groups(
     if isinstance(mol, str):
         mol = smi2mol(mol)
     if mol is None:
-        return []
+        return set([])
 
-    matched_atoms = set()
+    matched_atoms: set = set()
     result = set()
 
     sorted_smarts = sorted(_SMARTS.items(), key=lambda x: x[1][1])
@@ -310,7 +312,7 @@ def clear_descriptors(
 def add_descriptors(
     examples: List[Example],
     descriptor_type: str = "MACCS",
-    mols: List[Any] = None,
+    mols: Optional[List[Any]] = None,
 ) -> List[Example]:
     """Add descriptors to passed examples
 
@@ -415,7 +417,7 @@ def get_basic_alphabet() -> Set[str]:
 # is always a plain uncharged element.
 
 
-def _alphabet_to_elements(alphabet: List[str]) -> Set[str]:
+def _alphabet_to_elements(alphabet: Union[List[str], Set[str]]) -> Set[str]:
     """Converts SELFIES alphabet to element symbols"""
     symbols = []
     for s in alphabet:
@@ -426,8 +428,8 @@ def _alphabet_to_elements(alphabet: List[str]) -> Set[str]:
 
 
 def _check_alphabet_consistency(
-    smiles: str, alphabet_symbols: Set[str], check=False
-) -> True:
+    smiles: str, alphabet_symbols: Union[Set[str], List[str]], check=False
+) -> bool:
     """Checks if SMILES only contains tokens from alphabet"""
 
     alphabet_symbols = _alphabet_to_elements(set(alphabet_symbols))
@@ -448,7 +450,7 @@ def run_stoned(
     num_samples: int = 2000,
     max_mutations: int = 2,
     min_mutations: int = 1,
-    alphabet: Union[List[str], Set[str]] = None,
+    alphabet: Optional[Union[List[str], Set[str]]] = None,
     return_selfies: bool = False,
     _pbar: Any = None,
 ) -> Union[Tuple[List[str], List[float]], Tuple[List[str], List[str], List[float]]]:
@@ -491,11 +493,11 @@ def run_stoned(
         # Mutate the SELFIES:
         if _pbar:
             _pbar.set_description(f"🥌STONED🥌 Mutations: {num_mutations}")
-        selfies_mut = stoned.get_mutated_SELFIES(
+        selfies_mut: list | tuple = stoned.get_mutated_SELFIES(
             selfies_ls.copy(), num_mutations=num_mutations, alphabet=alphabet
         )
         # Convert back to SMILES:
-        smiles_back = [sf.decoder(x) for x in selfies_mut]
+        smiles_back: list | tuple = [sf.decoder(x) for x in selfies_mut]
         # check if smiles are consistent with alphabet and downslect
         selfies_mut, smiles_back = zip(
             *[
@@ -664,10 +666,10 @@ def sample_space(
     ],
     batched: bool = True,
     preset: str = "medium",
-    data: List[Union[str, rdchem.Mol]] = None,
-    method_kwargs: Dict = None,
-    num_samples: int = None,
-    stoned_kwargs: Dict = None,
+    data: Optional[List[Union[str, rdchem.Mol]]] = None,
+    method_kwargs: Optional[Dict] = None,
+    num_samples: Optional[int] = None,
+    stoned_kwargs: Optional[Dict] = None,
     quiet: bool = False,
     use_selfies: bool = False,
     sanitize_smiles: bool = True,
@@ -1017,7 +1019,7 @@ def is_low(e):
 def plot_space(
     examples: List[Example],
     exps: List[Example],
-    figure_kwargs: Dict = None,
+    figure_kwargs: Optional[Dict] = None,
     mol_size: Tuple[int, int] = (200, 200),
     highlight_clusters: bool = False,
     mol_fontsize: int = 8,
@@ -1111,11 +1113,11 @@ def normalizer(x):
 def plot_cf(
     exps: List[Example],
     fig: Any = None,
-    figure_kwargs: Dict = None,
+    figure_kwargs: Optional[Dict] = None,
     mol_size: Tuple[int, int] = (200, 200),
     mol_fontsize: int = 10,
-    nrows: int = None,
-    ncols: int = None,
+    nrows: Optional[int] = None,
+    ncols: Optional[int] = None,
 ):
     """Draw the given set of Examples in a grid
 
@@ -1159,10 +1161,10 @@ def plot_cf(
 
 def plot_descriptors(
     examples: List[Example],
-    output_file: str = None,
+    output_file: Optional[str] = None,
     fig: Any = None,
-    figure_kwargs: Dict = None,
-    title: str = None,
+    figure_kwargs: Optional[Dict] = None,
+    title: Optional[str] = None,
     return_svg: bool = False,
 ):
     """Plot descriptor attributions from given set of Examples.

diff --git a/exmol/plot_utils.py b/exmol/plot_utils.py
@@ -290,6 +290,7 @@ def similarity_map_using_tstats(
     if return_svg:
         return text
     _imgtext2mpl(text)
+    return None
 
 
 def plot_space_by_fit(
@@ -300,7 +301,7 @@ def plot_space_by_fit(
     mol_fontsize: int = 8,
     offset: int = 0,
     ax: Any = None,
-    figure_kwargs: Dict = None,
+    figure_kwargs: Optional[Dict] = None,
     cartoon: bool = False,
     rasterized: bool = False,
 ):

diff --git a/exmol/py.typed b/exmol/py.typed
diff --git a/exmol/version.py b/exmol/version.py
@@ -1 +1 @@
-__version__ = "3.2.2"
+__version__ = "3.2.3"
diff --git a/paper1_CFs/RF.ipynb b/paper1_CFs/RF.ipynb
@@ -385,9 +385,9 @@
  "metadata": {
   "celltoolbar": "Tags",
   "kernelspec": {
-   "display_name": "exmol312",
+   "display_name": "Python 3 (ipykernel)",
    "language": "python",
-   "name": "exmol312"
+   "name": "python3"
   },
   "language_info": {
    "codemirror_mode": {

diff --git a/paper3_Scents/GNNModelTrainingAndEvaluation.ipynb b/paper3_Scents/GNNModelTrainingAndEvaluation.ipynb
@@ -511,7 +511,7 @@
     "# Code to compute L2 regularization based on that in the \"MLP on MNIST\" Example on the Haiku Github repository (https://github.com/deepmind/dm-haiku/blob/main/examples/mnist.py)\n",
     "def loss_fn_logits_reg(params, x, y):\n",
     "    l2_lossTerm = regularizationStrength * sum(\n",
-    "        jnp.sum(jnp.square(p)) for p in jax.tree_leaves(params)\n",
+    "        jnp.sum(jnp.square(p)) for p in jax.tree_util.tree_leaves(params)\n",
     "    )\n",
     "    logits = loss_fn_logits(params, x, y)\n",
     "    return logits + l2_lossTerm"

diff --git a/setup.py b/setup.py
@@ -9,14 +9,14 @@
 
 setup(
     name="exmol",
-    version=__version__,
+    version=__version__,  # type: ignore
     description="Counterfactual generation with STONED SELFIES",
     author="Aditi Seshadri, Geemi Wellawatte, Andrew White",
     author_email="andrew.white@rochester.edu",
     url="https://ur-whitelab.github.io/exmol/",
     license="MIT",
     packages=["exmol", "exmol.stoned"],
-    package_data={"exmol": ["lime_data/*.txt", "lime_data/*.pb"]},
+    package_data={"exmol": ["lime_data/*.txt", "lime_data/*.pb", "exmol/py.typed"]},
     install_requires=[
         "selfies >= 2.0.0",
         "numpy",