type-checking investigation (#68)

* print mypy and python versions in workflow

* 'if save_path' used to fix 10 mypy errors

* automatic fix using https://github.com/hauntsaninja/no_implicit_optional

* post-fix linting

* refactor to address C901: too complex

* refactor to address C901: too complex

* ignore C901

* fix typo

* sort imports

* if save_path is not None

* no more 3.12 bc numpy 1.24

* fix all mypy issues with 1.8.0

.github/workflows/build.yml

@@ -25,7 +25,7 @@ jobs:
     name: Test build process
     strategy:
       matrix:
-        python-version: ["3.7", "3.8", "3.9", "3.10", "3.11", "3.12"]
+        python-version: ["3.7", "3.8", "3.9", "3.10", "3.11"]
     runs-on: ubuntu-latest
     steps:
       - name: Checkout
@@ -96,4 +96,6 @@ jobs:
 
       - name: type checks
         run: |
+          python --version
+          mypy --version
           mypy src/mud/ tests/

.github/workflows/main.yml

@@ -25,7 +25,7 @@ jobs:
     name: Run unit tests
     strategy:
       matrix:
-        python-version: ["3.7", "3.10", "3.12"]
+        python-version: ["3.7", "3.10", "3.11"]
     runs-on: ubuntu-latest
     steps:
       - name: Checkout

src/mud/base.py

@@ -103,8 +103,8 @@ def __init__(
         self,
         X: ArrayLike,
         y: ArrayLike,
-        domain: ArrayLike = None,
-        weights: ArrayLike = None,
+        domain: Optional[ArrayLike] = None,
+        weights: Optional[ArrayLike] = None,
         normalize: bool = False,
         pad_domain: float = 0.1,
     ):
@@ -146,7 +146,7 @@ def n_features(self):
     def n_samples(self):
         return self.y.shape[0]
 
-    def set_weights(self, weights: ArrayLike = None, normalize: bool = False):
+    def set_weights(self, weights: Optional[ArrayLike] = None, normalize: bool = False):
         """Set Sample Weights
 
         Sets the weights to use for each sample. Note weights can be one or two
@@ -249,9 +249,9 @@ def set_initial(self, distribution: Optional[rv_continuous] = None):
 
     def set_predicted(
         self,
-        distribution: rv_continuous = None,
-        bw_method: Union[str, Callable, np.generic] = None,
-        weights: ArrayLike = None,
+        distribution: Optional[rv_continuous] = None,
+        bw_method: Optional[Union[str, Callable, np.generic]] = None,
+        weights: Optional[ArrayLike] = None,
         **kwargs,
     ):
         """
@@ -425,10 +425,10 @@ def expected_ratio(self):
     def plot_param_space(
         self,
         param_idx: int = 0,
-        true_val: ArrayLike = None,
-        ax: plt.Axes = None,
-        x_range: Union[list, np.ndarray] = None,
-        ylim: float = None,
+        true_val: Optional[ArrayLike] = None,
+        ax: Optional[plt.Axes] = None,
+        x_range: Optional[Union[list, np.ndarray]] = None,
+        ylim: Optional[float] = None,
         pad_ratio: float = 0.05,
         aff: int = 100,
         in_opts={"color": "b", "linestyle": "-", "label": r"$\pi_\mathrm{init}$"},
@@ -534,8 +534,8 @@ def plot_param_space(
     def plot_obs_space(
         self,
         obs_idx: int = 0,
-        ax: plt.Axes = None,
-        y_range: ArrayLike = None,
+        ax: Optional[plt.Axes] = None,
+        y_range: Optional[ArrayLike] = None,
         aff=100,
         ob_opts={"color": "r", "linestyle": "-", "label": r"$\pi_\mathrm{obs}$"},
         pr_opts={"color": "b", "linestyle": "-", "label": r"$\pi_\mathrm{pred}$"},
@@ -618,7 +618,9 @@ def plot_obs_space(
 
         return ax
 
-    def plot_qoi(self, idx_x: int = 0, idx_y: int = 1, ax: plt.Axes = None, **kwargs):
+    def plot_qoi(
+        self, idx_x: int = 0, idx_y: int = 1, ax: Optional[plt.Axes] = None, **kwargs
+    ):
         """
         Plot 2D plot over two indices of y space.
 
@@ -652,7 +654,7 @@ def plot_params_2d(
         y: int = 0,
         contours: bool = False,
         colorbar: bool = True,
-        ax: plt.Axes = None,
+        ax: Optional[plt.Axes] = None,
         label=True,
         **kwargs,
     ):
@@ -748,7 +750,7 @@ def __init__(
         self,
         X: Union[np.ndarray, List],
         y: Union[np.ndarray, List],
-        domain: Union[np.ndarray, List] = None,
+        domain: Optional[Union[np.ndarray, List]] = None,
     ):
         # Set and validate inputs. Note we reshape inputs as necessary
         def set_shape(x, y):

src/mud/examples/adcirc.py

@@ -4,7 +4,7 @@
 Using SpatioTemporalProblem class to aggregate temporal data from ADCIRC for solving
 a two parameter estimation problem.
 """
-from typing import List, Tuple
+from typing import List, Optional, Tuple
 
 import matplotlib.colors as colors  # type: ignore
 import matplotlib.dates as mdates  # type: ignore
@@ -56,7 +56,7 @@ def tri_mesh_plot(
     stations=[0],
     zoom=None,
     colorbar_cutoff=-10,
-    save_path: str = None,
+    save_path: Optional[str] = None,
     close_fig: bool = False,
     dpi: int = 500,
 ):
@@ -70,9 +70,8 @@ def tri_mesh_plot(
     # Plot values and grid on top of it
     if value == "DP":
         name = "jet_terrain"
-        new_map = colors.LinearSegmentedColormap.from_list(
-            name, plt.cm.gist_rainbow_r(np.linspace(0.3, 0.9, 256))
-        )
+        _cmap = plt.cm.gist_rainbow_r(np.linspace(0.3, 0.9, 256))  # type: ignore
+        new_map = colors.LinearSegmentedColormap.from_list(name, _cmap)
         cutoff_val = colorbar_cutoff
         # make the norm:  Note the center is offset so that the land has more
         divnorm = colors.SymLogNorm(
@@ -97,8 +96,8 @@ def tri_mesh_plot(
     ax.set_aspect("equal")
 
     if zoom is not None:
-        ax.set_xlim([zoom[0][0] - zoom[0][1], zoom[0][0] + zoom[0][1]])
-        ax.set_ylim([zoom[1][0] - zoom[1][1], zoom[1][0] + zoom[1][1]])
+        ax.set_xlim(zoom[0][0] - zoom[0][1], zoom[0][0] + zoom[0][1])
+        ax.set_ylim(zoom[1][0] - zoom[1][1], zoom[1][0] + zoom[1][1])
         ax.set_aspect("equal")
 
     if stations is not None:
@@ -111,10 +110,10 @@ def tri_mesh_plot(
                 label=f"Recording Station {i}",
             )
     ax.legend()
-
-    save_figure(
-        f"si-{value}", save_path, close_fig=close_fig, dpi=dpi, bbox_inches="tight"
-    )
+    if save_path is not None:
+        save_figure(
+            f"si-{value}", save_path, close_fig=close_fig, dpi=dpi, bbox_inches="tight"
+        )
 
 
 def adcirc_ts_plot(
@@ -170,22 +169,26 @@ def adcirc_ts_plot(
     ax.set_ylabel("Water Elevation (m)")
     ax.set_xlabel("Time")
     _ = ax.set_title("")
-
-    save_figure(
-        "adcirc_full_ts", save_path, close_fig=close_fig, dpi=dpi, bbox_inches="tight"
-    )
+    if save_path is not None:
+        save_figure(
+            "adcirc_full_ts",
+            save_path,
+            close_fig=close_fig,
+            dpi=dpi,
+            bbox_inches="tight",
+        )
 
 
-def adcirc_time_window(
+def adcirc_time_window(  # noqa: C901
     adcirc_prob,
     time_window: Tuple[str, str],
     method="pca",
     num_components: int = 1,
     max_plot: int = 50,
     msize: int = 10,
-    ylims: List[float] = None,
-    title: str = None,
-    save_path: str = None,
+    ylims: Optional[List[float]] = None,
+    title: Optional[str] = None,
+    save_path: Optional[str] = None,
     plot_figs: List[str] = ["all"],
     close_fig: bool = False,
     dpi: int = 500,
@@ -217,22 +220,24 @@ def adcirc_time_window(
         pca_vector_plot(
             adcirc_prob, t_mask, msize=msize, max_plot=max_plot, title=title
         )
-        save_figure(
-            f"pca_vecs_{num_components}_{ndata}",
-            save_path,
-            close_fig=close_fig,
-            dpi=dpi,
-            bbox_inches="tight",
-        )
+        if save_path is not None:
+            save_figure(
+                f"pca_vecs_{num_components}_{ndata}",
+                save_path,
+                close_fig=close_fig,
+                dpi=dpi,
+                bbox_inches="tight",
+            )
     if "updated_dist" in plot_figs or "all" in plot_figs:
         updated_dist_plot(prob, lam_ref=adcirc_prob.lam_ref, title=title, ylims=ylims)
-        save_figure(
-            f"updated_dist_{num_components}_{ndata}",
-            save_path,
-            close_fig=close_fig,
-            dpi=dpi,
-            bbox_inches="tight",
-        )
+        if save_path is not None:
+            save_figure(
+                f"updated_dist_{num_components}_{ndata}",
+                save_path,
+                close_fig=close_fig,
+                dpi=dpi,
+                bbox_inches="tight",
+            )
 
     if "learned_qoi" in plot_figs or "all" in plot_figs:
         fig = plt.figure(figsize=(12, 5))
@@ -245,13 +250,14 @@ def adcirc_time_window(
 
         if title is not None:
             _ = fig.suptitle(f"{title}", fontsize=20)
-        save_figure(
-            f"qoi_{num_components}_{ndata}",
-            save_path,
-            close_fig=close_fig,
-            dpi=dpi,
-            bbox_inches="tight",
-        )
+        if save_path is not None:
+            save_figure(
+                "qoi_{num_components}_{ndata}",
+                save_path,
+                close_fig=close_fig,
+                dpi=dpi,
+                bbox_inches="tight",
+            )
 
     return prob
 

src/mud/examples/comparison.py

@@ -4,7 +4,7 @@
 Functions for running 1-dimensional polynomial inversion problem.
 """
 import logging
-from typing import List
+from typing import List, Optional
 
 import matplotlib.pyplot as plt  # type: ignore
 import numpy as np
@@ -23,9 +23,7 @@ def comparison_plot(
     d_prob: DensityProblem,
     b_prob: BayesProblem,
     space: str = "param",
-    ax: plt.Axes = None,
-    save_path: str = None,
-    dpi: int = 500,
+    ax: Optional[plt.Axes] = None,
 ):
     """
     Generate plot comparing MUD vs MAP solution
@@ -47,20 +45,16 @@ def comparison_plot(
     ax : matplotlib.pyplot.Axes, optional
         Existing matplotlib Axes object to plot onto. If none provided
         (default), then a figure is initialized.
-    save_path : str, optional
-        Path to save figure to.
-    dpi : int
-        If set to `save_path` is specified, then the resolution of the saved
-        image to use.
 
     Returns
     -------
     ax : matplotlib.pyplot.Axes
         Axes object that was plotted onto or created.
     """
 
-    # Plot comparison plots of b_prob vs DCI solutions
-    fig, ax = plt.subplots(1, 1, figsize=(6, 6))
+    if ax is None:
+        # Plot comparison plots of b_prob vs DCI solutions
+        _, ax = plt.subplots(1, 1, figsize=(6, 6))
 
     # Parameters for initial and updated plots
     legend_fsize = 14
@@ -106,10 +100,10 @@ def comparison_plot(
         b_prob.plot_param_space(ax=ax, pr_opts=None, ps_opts=ps_opts, map_opts=None)
 
         # Format figure
-        _ = ax.set_xlim([-1, 1])
-        _ = ax.set_ylim(ylim_p)
-        _ = ax.tick_params(axis="x", labelsize=tick_fsize)
-        _ = ax.tick_params(axis="y", labelsize=tick_fsize)
+        _ = ax.set_xlim(-1, 1)
+        _ = ax.set_ylim(*ylim_p)
+        ax.tick_params(axis="x", labelsize=tick_fsize)
+        ax.tick_params(axis="y", labelsize=tick_fsize)
         _ = ax.set_xlabel("$\\Lambda$", fontsize=1.25 * tick_fsize)
         _ = ax.legend(fontsize=legend_fsize, loc="upper left")
     else:
@@ -124,9 +118,9 @@ def comparison_plot(
         # y_range=np.array([[0,1]]))
 
         # Format figure
-        _ = ax.set_xlim([-1, 1])
-        _ = ax.tick_params(axis="x", labelsize=tick_fsize)
-        _ = ax.tick_params(axis="y", labelsize=tick_fsize)
+        _ = ax.set_xlim(-1, 1)
+        ax.tick_params(axis="x", labelsize=tick_fsize)
+        ax.tick_params(axis="y", labelsize=tick_fsize)
         _ = ax.set_xlabel("$\\mathcal{D}$", fontsize=1.25 * tick_fsize)
         _ = ax.legend(fontsize=legend_fsize, loc="upper left")
 
@@ -141,7 +135,7 @@ def run_comparison_example(
     domain: List[int] = [-1, 1],
     N_vals: List[int] = [1, 5, 10, 20],
     latex_labels: bool = True,
-    save_path: str = None,
+    save_path: Optional[str] = None,
     dpi: int = 500,
     close_fig: bool = False,
 ):
@@ -211,19 +205,24 @@ def run_comparison_example(
             _ = ax.set_ylim([-0.2, 28.0])
         else:
             ax.set_ylim([-0.2, 6])
-        save_figure(
-            f"bip-vs-sip-{N}.png", save_path=save_path, dpi=dpi, close_fig=close_fig
-        )
+        if save_path is not None:
+            save_figure(
+                f"bip-vs-sip-{N}.png", save_path=save_path, dpi=dpi, close_fig=close_fig
+            )
 
         ax = comparison_plot(d_prob, b_prob, space="obs")
         if N != 1:
             _ = ax.text(-0.75, 5, f"N = {N}", fontsize=22)
             _ = ax.set_ylim([-0.2, 28.0])
         else:
             ax.set_ylim([-0.2, 4.5])
-        save_figure(
-            f"bip-vs-sip-pf-{N}.png", save_path=save_path, dpi=dpi, close_fig=close_fig
-        )
+        if save_path is not None:
+            save_figure(
+                f"bip-vs-sip-pf-{N}.png",
+                save_path=save_path,
+                dpi=dpi,
+                close_fig=close_fig,
+            )
 
         res.append([d_prob, b_prob, ax])