diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml
index d9748e1..cef4b53 100644
--- a/.github/workflows/build.yml
+++ b/.github/workflows/build.yml
@@ -2,37 +2,39 @@ name: builds
 
 on:
   push:
-    branches: [ main ]
+    branches: [main]
     paths:
       - "src/**.py"
+      - "tests/**.py"
       - "setup.py"
       - "setup.cfg"
       - ".github/workflows/build.yml"
   pull_request:
-    branches-ignore: '**docker**'
+    branches-ignore: ["**docker**"]
     paths:
       - "src/**.py"
+      - "tests/**.py"
       - "setup.py"
       - "setup.cfg"
       - ".github/workflows/build.yml"
   schedule:
-      - cron: "0 0 1 * *"
+    - cron: "0 0 1 * *"
 
 jobs:
   build:
-    name: Test build process
+    name: packaging
     strategy:
       matrix:
-        python-version: ["3.7", "3.8", "3.9", "3.10"]
+        python-version: ["3.7", "3.8", "3.9", "3.10", "3.11", "3.12"]
     runs-on: ubuntu-latest
     steps:
       - name: Checkout
-        uses: actions/checkout@v2
+        uses: actions/checkout@v3
         with:
           fetch-depth: 0
 
       - name: Set up Python ${{ matrix.python-version }}
-        uses: actions/setup-python@v1
+        uses: actions/setup-python@v2
         with:
           python-version: ${{ matrix.python-version }}
 
@@ -42,7 +44,7 @@ jobs:
           pip install --upgrade wheel setuptools setuptools_scm
 
       - name: Inspect version info
-        run: | 
+        run: |
           python setup.py --version
           git describe --dirty --tags --long --match "*[0-9]*"
 
@@ -55,12 +57,15 @@ jobs:
         run: |
           pip install -e .
           pip uninstall -y mud
- 
+
       - name: Test build
         run: |
           python setup.py sdist bdist_wheel
-          pip uninstall -y mud
 
+      - name: New build
+        run: |
+          pip install build
+          python -m build --sdist --wheel
 
   style:
     name: Enforce style
@@ -70,12 +75,12 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - name: checkout
-        uses: actions/checkout@v2
+        uses: actions/checkout@v3
         with:
           fetch-depth: 1
 
       - name: setup
-        uses: actions/setup-python@v1
+        uses: actions/setup-python@v2
         with:
           python-version: ${{ matrix.python-version }}
 
@@ -94,4 +99,6 @@ jobs:
 
       - name: type checks
         run: |
+          python --version
+          mypy --version
           mypy src/mud/ tests/
diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml
index adafb97..153963c 100644
--- a/.github/workflows/main.yml
+++ b/.github/workflows/main.yml
@@ -5,13 +5,15 @@ on:
     branches: [ main ]
     paths:
       - "src/**.py"
+      - "tests/**.py"
       - "setup.py"
       - "setup.cfg"
       - ".github/workflows/main.yml"
   pull_request:
-    branches-ignore: "**docker**"
+    branches-ignore: ["**docker**"]
     paths:
       - "src/**.py"
+      - "tests/**.py"
       - "setup.py"
       - "setup.cfg"
       - ".github/workflows/main.yml"
@@ -20,19 +22,19 @@ on:
 
 jobs:
   unit-tests:
-    name: Run unit tests
+    name: pytest
     strategy:
       matrix:
-        python-version: ["3.7", "3.10"]
+        python-version: ["3.7", "3.10", "3.11", "3.12"]
     runs-on: ubuntu-latest
     steps:
       - name: Checkout
-        uses: actions/checkout@v2
+        uses: actions/checkout@v3
         with:
           fetch-depth: 0
 
       - name: Set up Python ${{ matrix.python-version }}
-        uses: actions/setup-python@v1
+        uses: actions/setup-python@v2
         with:
           python-version: ${{ matrix.python-version }}
 
@@ -45,25 +47,25 @@ jobs:
         run: pytest
 
       - name: Upload coverage to coveralls.io
-        if: ${{ matrix.python-version }} == 3.9
+        if: ${{ matrix.python-version }} == 3.10
         run: coveralls --service=github
         env:
           GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
 
   integration-tests:
-    name: Run integration tests
+    name: cli examples
     strategy:
       matrix:
-        python-version: ["3.9"]
+        python-version: ["3.10", "3.12"]
     runs-on: ubuntu-latest
     steps:
       - name: Checkout
-        uses: actions/checkout@v2
+        uses: actions/checkout@v3
         with:
           fetch-depth: 1
 
       - name: Set up Python ${{ matrix.python-version }}
-        uses: actions/setup-python@v1
+        uses: actions/setup-python@v2
         with:
           python-version: ${{ matrix.python-version }}
 
@@ -77,6 +79,7 @@ jobs:
             texlive-fonts-extra \
             texlive-science \
             latexmk \
+            dvipng \
             cm-super
 
       - name: Install Python dependencies
@@ -86,7 +89,7 @@ jobs:
           pip install .[examples]
 
       - name: Old CLI
-        continue-on-error: true
+        continue-on-error: false
         run: mud_run_all -v
 
       - name: New CLI
diff --git a/.readthedocs.yml b/.readthedocs.yml
index 299280c..3b79cad 100644
--- a/.readthedocs.yml
+++ b/.readthedocs.yml
@@ -4,6 +4,11 @@
 # Required
 version: 2
 
+build:
+  os: "ubuntu-20.04"
+  tools:
+    python: "3.8"
+
 # Build documentation in the docs/ directory with Sphinx
 sphinx:
   configuration: docs/conf.py
@@ -18,7 +23,6 @@ formats:
   - epub
 
 python:
-  version: 3.8
   install:
     - requirements: docs/requirements.txt
     - method: pip
diff --git a/notebooks/FenicsSolve.ipynb b/notebooks/FenicsSolve.ipynb
index f808f30..caceed8 100644
--- a/notebooks/FenicsSolve.ipynb
+++ b/notebooks/FenicsSolve.ipynb
@@ -190,7 +190,6 @@
     "    gamma=-3,\n",
     "    save_data=True,\n",
     "):\n",
-    "\n",
     "    if len(mins) != len(maxs):\n",
     "        raise ValueError(\"min/max arrays must be of same length\")\n",
     "\n",
diff --git a/notebooks/MUD-ADCIRC-Figs.ipynb b/notebooks/MUD-ADCIRC-Figs.ipynb
index 7998ee8..587bec5 100644
--- a/notebooks/MUD-ADCIRC-Figs.ipynb
+++ b/notebooks/MUD-ADCIRC-Figs.ipynb
@@ -3516,7 +3516,6 @@
     "    close_fig: bool = False,\n",
     "    dpi: int = 500,\n",
     "):\n",
-    "\n",
     "    triangles = mtri.Triangulation(\n",
     "        adcirc_grid_data[\"X\"], adcirc_grid_data[\"Y\"], adcirc_grid_data[\"triangles\"]\n",
     "    )\n",
diff --git a/notebooks/MUD_Butler_2021.ipynb b/notebooks/MUD_Butler_2021.ipynb
index 1da23b0..23d24bf 100644
--- a/notebooks/MUD_Butler_2021.ipynb
+++ b/notebooks/MUD_Butler_2021.ipynb
@@ -1915,7 +1915,6 @@
     "def plot_conv(\n",
     "    info, method=\"pca\", field=\"err\", agg=\"mean\", ax=None, fit=False, **kwargs\n",
     "):\n",
-    "\n",
     "    stat = info.groupby([\"type\", \"n\"])[field].aggregate([agg])\n",
     "    stat = stat.unstack(level=0)[(agg, method)]\n",
     "\n",
diff --git a/notebooks/Nonlinear_Examples.ipynb b/notebooks/Nonlinear_Examples.ipynb
index 7f2f23b..a5f13fa 100644
--- a/notebooks/Nonlinear_Examples.ipynb
+++ b/notebooks/Nonlinear_Examples.ipynb
@@ -200,7 +200,6 @@
     "    t_start=1.0,\n",
     "    sampling_freq=100.0,\n",
     "):\n",
-    "\n",
     "    time = np.linspace(time_range[0], time_range[1], 1000)\n",
     "    times = np.arange(t_start, time_range[1], 1 / sampling_freq)[0:N]\n",
     "\n",
diff --git a/notebooks/ODE_ExpDecay.ipynb b/notebooks/ODE_ExpDecay.ipynb
index f841c92..ab54e71 100644
--- a/notebooks/ODE_ExpDecay.ipynb
+++ b/notebooks/ODE_ExpDecay.ipynb
@@ -122,7 +122,6 @@
     "    sampling_freq=100.0,\n",
     "    std_dev=0.05,\n",
     "):\n",
-    "\n",
     "    u_t_lambda = lambda t, l: u_0 * np.exp(-np.outer(l, t))\n",
     "\n",
     "    # Build initial samples\n",
@@ -190,7 +189,6 @@
     "    sampling_freq=10.0,\n",
     "    std_dev=0.05,\n",
     "):\n",
-    "\n",
     "    u_t_lambda = lambda t, l1, l2: (l1 * np.exp(-np.outer(t, l2))).T\n",
     "\n",
     "    # Build initial samples\n",
@@ -619,7 +617,6 @@
     "    sampling_freq=10.0,\n",
     "    std_dev=0.05,\n",
     "):\n",
-    "\n",
     "    u_t_lambda = lambda t, l1, l2: (l1 * np.exp(-np.outer(t, l2))).T\n",
     "\n",
     "    # Build initial samples\n",
diff --git a/setup.cfg b/setup.cfg
index fa03ab4..7efdeb6 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -32,10 +32,10 @@ package_dir =
 # Add here dependencies of your project (semicolon/line-separated), e.g.
 install_requires =
     importlib-metadata; python_version<"3.10"
-    numpy>=1.20
-    scipy
-    matplotlib
-    scikit-learn
+    numpy>=1.20,<1.27
+    scipy>=1.2,<2.0
+    matplotlib>=1.0,<4.0
+    scikit-learn>=1.0,<2.0
 
 # Require a specific Python version, e.g. Python 2.7 or >= 3.4
 python_requires = >=3.7
@@ -67,7 +67,7 @@ examples =
     click
     tqdm
     prettytable
-    pandas
+    pandas>=1.2,<3.0
     wget
 
 
diff --git a/src/mud/base.py b/src/mud/base.py
index 10ec5b8..bc9b719 100644
--- a/src/mud/base.py
+++ b/src/mud/base.py
@@ -103,12 +103,11 @@ 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,
     ):
-
         self.X = set_shape(np.array(X), (1, -1))
         self.y = set_shape(np.array(y), (-1, 1))
 
@@ -147,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
@@ -250,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,
     ):
         """
@@ -426,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}$"},
@@ -535,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}$"},
@@ -619,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.
 
@@ -653,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,
     ):
@@ -749,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):
@@ -1042,7 +1043,6 @@ def __init__(
         cov_o=None,
         alpha=1.0,
     ):
-
         # Make sure A is 2D array
         A = np.array(A)
         self.A = A if A.ndim == 2 else A.reshape(1, -1)
@@ -1351,7 +1351,6 @@ def __init__(
         cov_o=None,
         alpha=1.0,
     ):
-
         if isinstance(sigma, (float, int)):
             sigma = [sigma] * len(data)
 
@@ -1425,7 +1424,6 @@ class IterativeLinearProblem(LinearGaussianProblem):
     def __init__(
         self, A, b, y=None, mu_i=None, cov=None, data_cov=None, idx_order=None
     ):
-
         # Make sure A is 2D array
         self.A = A if A.ndim == 2 else A.reshape(1, -1)
 
@@ -1545,7 +1543,6 @@ class SpatioTemporalProblem(object):
     """
 
     def __init__(self, df=None):
-
         self._domain = None
         self._lam = None
         self._data = None
@@ -1790,7 +1787,7 @@ def load(
             Dictionary containing data from file for PDE problem class
 
         """
-        if type(df) == str:
+        if isinstance(df, str):
             try:
                 if df.endswith("nc") and xr_avail:
                     ds = xr.load_dataset(df)
@@ -1805,7 +1802,7 @@ def load(
         def get_set_val(f, v):
             if f in ds.keys():
                 self.__setattr__(f, ds[v])
-            elif v is not None and type(v) != str:
+            elif v is not None and not isinstance(v, str):
                 self.__setattr__(f, v)
 
         field_names = {
diff --git a/src/mud/examples/adcirc.py b/src/mud/examples/adcirc.py
index 55c0aaf..6960805 100644
--- a/src/mud/examples/adcirc.py
+++ b/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,11 +56,10 @@ 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,
 ):
-
     triangles = mtri.Triangulation(
         adcirc_grid_data["X"], adcirc_grid_data["Y"], adcirc_grid_data["triangles"]
     )
@@ -71,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(
@@ -98,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:
@@ -112,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(
@@ -171,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,
@@ -218,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))
@@ -246,13 +250,14 @@ def adcirc_time_window(
 
         if title is not None:
             _ = fig.suptitle(f"{title}", fontsize=20)
-        save_figure(
-            "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
 
diff --git a/src/mud/examples/comparison.py b/src/mud/examples/comparison.py
index a0be18a..36553ef 100644
--- a/src/mud/examples/comparison.py
+++ b/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,11 +45,6 @@ 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
     -------
@@ -59,8 +52,9 @@ def comparison_plot(
         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,9 +205,10 @@ 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:
@@ -221,9 +216,13 @@ def run_comparison_example(
             _ = 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])
 
diff --git a/src/mud/examples/linear.py b/src/mud/examples/linear.py
index bce9cc9..0f59573 100644
--- a/src/mud/examples/linear.py
+++ b/src/mud/examples/linear.py
@@ -4,12 +4,11 @@
 Functions for examples for linear problems.
 """
 import logging
-from typing import List
+from typing import List, Optional
 
 import matplotlib.pyplot as plt  # type: ignore
 import numpy as np  # type: ignore
 import scipy as sp  # type: ignore
-from matplotlib import cm  # type: ignore
 
 from mud.base import IterativeLinearProblem, LinearGaussianProblem, LinearWMEProblem
 from mud.plot import mud_plot_params, save_figure
@@ -68,7 +67,6 @@ def random_linear_wme_problem(
     operator_list = []
     data_list = []
     for n, s in zip(num_observations, std_dev):
-
         if dist == "normal":
             M = np.random.randn(num_qoi, dim_input)
         else:
@@ -88,7 +86,7 @@ def random_linear_wme_problem(
         ref_data = M @ ref_input  # noqa: E221
         noise = np.diag(s) @ np.random.randn(n, 1)
         if ref_data.shape[0] == 1:
-            ref_data = float(ref_data)
+            ref_data = ref_data.ravel()
         data = ref_data + noise
 
         operator_list.append(M)
@@ -100,9 +98,9 @@ def random_linear_wme_problem(
 def random_linear_problem(
     dim_input: int = 10,
     dim_output: int = 10,
-    mean_i: np.typing.ArrayLike = None,
-    cov_i: np.typing.ArrayLike = None,
-    seed: int = None,
+    mean_i: Optional[np.typing.ArrayLike] = None,
+    cov_i: Optional[np.typing.ArrayLike] = None,
+    seed: Optional[int] = None,
 ):
     """Construct a random linear Gaussian Problem"""
 
@@ -264,9 +262,257 @@ def rotation_map_trials(
     ax.plot(avg_errs, color, lw=5, label=label)
 
 
+def call_consistent(lin_prob: LinearGaussianProblem, **kwargs):
+    fig, ax = plt.subplots(1, 1, figsize=(6, 6))
+    lin_prob.plot_fun_contours(
+        ax=ax, terms="reg_m", levels=50, cmap="viridis", alpha=1.0
+    )
+    lin_prob.plot_sol(
+        ax=ax,
+        point="initial",
+        label="Initial Mean",
+        pt_opts={
+            "color": "k",
+            "s": 100,
+            "marker": "o",
+            "label": "MUD",
+            "zorder": 10,
+        },
+    )
+    _ = ax.axis([0, 1, 0, 1])
+    if kwargs.get("save_path"):
+        save_figure("consistent_contour.png", **kwargs)
+
+
+def call_mismatch(lin_prob: LinearGaussianProblem, **kwargs):
+    fig, ax = plt.subplots(1, 1, figsize=(6, 6))
+    lin_prob.plot_fun_contours(
+        ax=ax, terms="data", levels=50, cmap="viridis", alpha=1.0
+    )
+    lin_prob.plot_contours(
+        ax=ax,
+        annotate=True,
+        note_loc=[0.1, 0.9],
+        label="Solution Contour",
+        plot_opts={"color": "r"},
+        annotate_opts={"fontsize": 20, "backgroundcolor": "w"},
+    )
+    ax.axis("equal")
+    _ = ax.set_xlim([0, 1])
+    _ = ax.set_ylim([0, 1])
+    if kwargs.get("save_path"):
+        save_figure("data_mismatch_contour.png", **kwargs)
+
+
+def call_tikhonov(lin_prob: LinearGaussianProblem, **kwargs):
+    fig, ax = plt.subplots(1, 1, figsize=(6, 6))
+    lin_prob.plot_fun_contours(ax=ax, terms="reg", levels=50, cmap="viridis", alpha=1.0)
+    lin_prob.plot_sol(
+        ax=ax,
+        point="initial",
+        label="Initial Mean",
+        pt_opts={
+            "color": "k",
+            "s": 100,
+            "marker": "o",
+            "label": "MUD",
+            "zorder": 10,
+        },
+    )
+    _ = ax.axis([0, 1, 0, 1])
+    if kwargs.get("save_path"):
+        save_figure("tikhonov_contour.png", **kwargs)
+
+
+def call_map(lin_prob: LinearGaussianProblem, **kwargs):
+    fig, ax = plt.subplots(1, 1, figsize=(6, 6))
+    lin_prob.plot_fun_contours(ax=ax, terms="bayes", levels=50, cmap="viridis")
+    lin_prob.plot_fun_contours(
+        ax=ax,
+        terms="data",
+        levels=25,
+        cmap="viridis",
+        alpha=0.5,
+        vmin=0,
+        vmax=4,
+    )
+    lin_prob.plot_sol(
+        ax=ax,
+        point="initial",
+        pt_opts={
+            "color": "k",
+            "s": 100,
+            "marker": "o",
+            "label": "MUD",
+            "zorder": 20,
+        },
+    )
+    lin_prob.plot_sol(
+        ax=ax,
+        point="ls",
+        label="Least Squares",
+        note_loc=[0.49, 0.55],
+        pt_opts={"color": "xkcd:blue", "s": 100, "marker": "d", "zorder": 10},
+        annotate_opts={"fontsize": 14, "backgroundcolor": "w"},
+    )
+    lin_prob.plot_sol(
+        ax=ax,
+        point="map",
+        label="MAP",
+        pt_opts={
+            "color": "tab:orange",
+            "s": 100,
+            "linewidths": 3,
+            "marker": "x",
+            "zorder": 10,
+        },
+        ln_opts=None,
+        annotate_opts={"fontsize": 14, "backgroundcolor": "w"},
+    )
+    lin_prob.plot_contours(
+        ax=ax,
+        annotate=False,
+        note_loc=[0.1, 0.9],
+        label="Solution Contour",
+        plot_opts={"color": "r"},
+        annotate_opts={"fontsize": 14, "backgroundcolor": "w"},
+    )
+    _ = ax.axis([0, 1, 0, 1])
+    if kwargs.get("save_path"):
+        save_figure("classical_solution.png", **kwargs)
+
+
+def call_mud(lin_prob: LinearGaussianProblem, **kwargs):
+    fig, ax = plt.subplots(1, 1, figsize=(6, 6))
+    lin_prob.plot_fun_contours(ax=ax, terms="dc", levels=50, cmap="viridis")
+    lin_prob.plot_fun_contours(
+        ax=ax,
+        terms="data",
+        levels=25,
+        cmap="viridis",
+        alpha=0.5,
+        vmin=0,
+        vmax=4,
+    )
+    lin_prob.plot_sol(
+        ax=ax,
+        point="initial",
+        pt_opts={
+            "color": "k",
+            "s": 100,
+            "marker": "o",
+            "label": "MUD",
+            "zorder": 20,
+        },
+    )
+    lin_prob.plot_sol(
+        ax=ax,
+        point="ls",
+        label="Least Squares",
+        note_loc=[0.49, 0.55],
+        pt_opts={"color": "k", "s": 100, "marker": "d", "zorder": 10},
+        annotate_opts={"fontsize": 14, "backgroundcolor": "w"},
+    )
+    lin_prob.plot_sol(
+        point="mud",
+        ax=ax,
+        label="MUD",
+        pt_opts={"color": "k", "s": 100, "marker": "*", "zorder": 10},
+        ln_opts={"color": "k", "marker": "*", "lw": 1, "zorder": 10},
+        annotate_opts={"fontsize": 14, "backgroundcolor": "w"},
+    )
+    lin_prob.plot_contours(
+        ax=ax,
+        annotate=False,
+        note_loc=[0.1, 0.9],
+        label="Solution Contour",
+        plot_opts={"color": "r"},
+        annotate_opts={"fontsize": 14, "backgroundcolor": "w"},
+    )
+    ax.axis("equal")
+    _ = ax.axis([0, 1, 0, 1])
+    if kwargs.get("save_path"):
+        save_figure("consistent_solution.png", **kwargs)
+
+
+def call_comparison(lin_prob: LinearGaussianProblem, **kwargs):
+    fig, ax = plt.subplots(1, 1, figsize=(8, 8))
+    lin_prob.plot_fun_contours(ax=ax, terms="bayes", levels=50, cmap="viridis")
+    lin_prob.plot_fun_contours(
+        ax=ax,
+        terms="data",
+        levels=25,
+        cmap="viridis",
+        alpha=0.5,
+        vmin=0,
+        vmax=4,
+    )
+    lin_prob.plot_sol(
+        ax=ax,
+        point="initial",
+        pt_opts={
+            "color": "k",
+            "s": 100,
+            "marker": "o",
+            "label": "MUD",
+            "zorder": 10,
+        },
+    )
+    lin_prob.plot_sol(
+        ax=ax,
+        point="ls",
+        label="Least Squares",
+        note_loc=[0.49, 0.55],
+        pt_opts={"color": "k", "s": 100, "marker": "d", "zorder": 10},
+    )
+    lin_prob.plot_sol(
+        ax=ax,
+        point="map",
+        label="MAP",
+        pt_opts={
+            "color": "tab:orange",
+            "s": 100,
+            "linewidth": 3,
+            "marker": "x",
+            "zorder": 10,
+        },
+        ln_opts=None,
+        annotate_opts={"fontsize": 14, "backgroundcolor": "w"},
+    )
+    lin_prob.plot_sol(
+        point="mud",
+        ax=ax,
+        label="MUD",
+        pt_opts={
+            "color": "k",
+            "s": 100,
+            "linewidth": 3,
+            "marker": "*",
+            "zorder": 10,
+        },
+        ln_opts={"color": "k", "marker": "*", "lw": 1, "zorder": 10},
+        annotate_opts={"fontsize": 14, "backgroundcolor": "w"},
+    )
+    lin_prob.plot_contours(
+        ax=ax,
+        annotate=False,
+        note_loc=[0.1, 0.9],
+        label="Solution Contour",
+        plot_opts={"color": "r"},
+        annotate_opts={"fontsize": 14, "backgroundcolor": "w"},
+    )
+    pt = [0.7, 0.3]
+    ax.scatter([pt[0]], [pt[1]], color="k", s=100, marker="s", zorder=11)
+    nc = (pt[0] - 0.02, pt[1] + 0.02)
+    ax.annotate("Truth", nc, fontsize=14, backgroundcolor="w")
+    _ = ax.axis([0, 1, 0, 1])
+    if kwargs.get("save_path"):
+        save_figure("map_compare_contour.png", **kwargs)
+
+
 def run_contours(
-    plot_fig: List[str] = None,
-    save_path: str = None,
+    plot_fig: Optional[List[str]] = None,
+    save_path: Optional[str] = None,
     dpi: int = 500,
     close_fig: bool = False,
     **kwargs,
@@ -323,251 +569,28 @@ def run_contours(
     _ = (lin_prob.solve("mud"), lin_prob.solve("map"), lin_prob.solve("ls"))
 
     if "data_mismatch" in plot_fig or "all" in plot_fig:
-        fig, ax = plt.subplots(1, 1, figsize=(6, 6))
-        lin_prob.plot_fun_contours(
-            ax=ax, terms="data", levels=50, cmap=cm.viridis, alpha=1.0
-        )
-        lin_prob.plot_contours(
-            ax=ax,
-            annotate=True,
-            note_loc=[0.1, 0.9],
-            label="Solution Contour",
-            plot_opts={"color": "r"},
-            annotate_opts={"fontsize": 20, "backgroundcolor": "w"},
-        )
-        ax.axis("equal")
-        _ = ax.set_xlim([0, 1])
-        _ = ax.set_ylim([0, 1])
-        save_figure(
-            "data_mismatch_contour.png",
-            save_path=save_path,
-            dpi=dpi,
-            close_fig=close_fig,
+        call_mismatch(
+            lin_prob=lin_prob, save_path=save_path, dpi=dpi, close_fig=close_fig
         )
     if "tikhonov" in plot_fig or "all" in plot_fig:
-        fig, ax = plt.subplots(1, 1, figsize=(6, 6))
-        lin_prob.plot_fun_contours(
-            ax=ax, terms="reg", levels=50, cmap=cm.viridis, alpha=1.0
-        )
-        lin_prob.plot_sol(
-            ax=ax,
-            point="initial",
-            label="Initial Mean",
-            pt_opts={
-                "color": "k",
-                "s": 100,
-                "marker": "o",
-                "label": "MUD",
-                "zorder": 10,
-            },
-        )
-        _ = ax.axis([0, 1, 0, 1])
-        save_figure(
-            "tikhonov_contour.png", save_path=save_path, dpi=dpi, close_fig=close_fig
+        call_tikhonov(
+            lin_prob=lin_prob, save_path=save_path, dpi=dpi, close_fig=close_fig
         )
+
     if "consistent" in plot_fig or "all" in plot_fig:
-        fig, ax = plt.subplots(1, 1, figsize=(6, 6))
-        lin_prob.plot_fun_contours(
-            ax=ax, terms="reg_m", levels=50, cmap=cm.viridis, alpha=1.0
-        )
-        lin_prob.plot_sol(
-            ax=ax,
-            point="initial",
-            label="Initial Mean",
-            pt_opts={
-                "color": "k",
-                "s": 100,
-                "marker": "o",
-                "label": "MUD",
-                "zorder": 10,
-            },
-        )
-        _ = ax.axis([0, 1, 0, 1])
-        save_figure(
-            "consistent_contour.png", save_path=save_path, dpi=dpi, close_fig=close_fig
+        call_consistent(
+            lin_prob=lin_prob, save_path=save_path, dpi=dpi, close_fig=close_fig
         )
+
     if "map" in plot_fig or "all" in plot_fig:
-        fig, ax = plt.subplots(1, 1, figsize=(6, 6))
-        lin_prob.plot_fun_contours(ax=ax, terms="bayes", levels=50, cmap=cm.viridis)
-        lin_prob.plot_fun_contours(
-            ax=ax,
-            terms="data",
-            levels=25,
-            cmap=cm.viridis,
-            alpha=0.5,
-            vmin=0,
-            vmax=4,
-        )
-        lin_prob.plot_sol(
-            ax=ax,
-            point="initial",
-            pt_opts={
-                "color": "k",
-                "s": 100,
-                "marker": "o",
-                "label": "MUD",
-                "zorder": 20,
-            },
-        )
-        lin_prob.plot_sol(
-            ax=ax,
-            point="ls",
-            label="Least Squares",
-            note_loc=[0.49, 0.55],
-            pt_opts={"color": "xkcd:blue", "s": 100, "marker": "d", "zorder": 10},
-            annotate_opts={"fontsize": 14, "backgroundcolor": "w"},
-        )
-        lin_prob.plot_sol(
-            ax=ax,
-            point="map",
-            label="MAP",
-            pt_opts={
-                "color": "tab:orange",
-                "s": 100,
-                "linewidths": 3,
-                "marker": "x",
-                "zorder": 10,
-            },
-            ln_opts=None,
-            annotate_opts={"fontsize": 14, "backgroundcolor": "w"},
-        )
-        lin_prob.plot_contours(
-            ax=ax,
-            annotate=False,
-            note_loc=[0.1, 0.9],
-            label="Solution Contour",
-            plot_opts={"color": "r"},
-            annotate_opts={"fontsize": 14, "backgroundcolor": "w"},
-        )
-        _ = ax.axis([0, 1, 0, 1])
-        save_figure(
-            "classical_solution.png", save_path=save_path, dpi=dpi, close_fig=close_fig
-        )
+        call_map(lin_prob=lin_prob, save_path=save_path, dpi=dpi, close_fig=close_fig)
+
     if "mud" in plot_fig or "all" in plot_fig:
-        fig, ax = plt.subplots(1, 1, figsize=(6, 6))
-        lin_prob.plot_fun_contours(ax=ax, terms="dc", levels=50, cmap=cm.viridis)
-        lin_prob.plot_fun_contours(
-            ax=ax,
-            terms="data",
-            levels=25,
-            cmap=cm.viridis,
-            alpha=0.5,
-            vmin=0,
-            vmax=4,
-        )
-        lin_prob.plot_sol(
-            ax=ax,
-            point="initial",
-            pt_opts={
-                "color": "k",
-                "s": 100,
-                "marker": "o",
-                "label": "MUD",
-                "zorder": 20,
-            },
-        )
-        lin_prob.plot_sol(
-            ax=ax,
-            point="ls",
-            label="Least Squares",
-            note_loc=[0.49, 0.55],
-            pt_opts={"color": "k", "s": 100, "marker": "d", "zorder": 10},
-            annotate_opts={"fontsize": 14, "backgroundcolor": "w"},
-        )
-        lin_prob.plot_sol(
-            point="mud",
-            ax=ax,
-            label="MUD",
-            pt_opts={"color": "k", "s": 100, "marker": "*", "zorder": 10},
-            ln_opts={"color": "k", "marker": "*", "lw": 1, "zorder": 10},
-            annotate_opts={"fontsize": 14, "backgroundcolor": "w"},
-        )
-        lin_prob.plot_contours(
-            ax=ax,
-            annotate=False,
-            note_loc=[0.1, 0.9],
-            label="Solution Contour",
-            plot_opts={"color": "r"},
-            annotate_opts={"fontsize": 14, "backgroundcolor": "w"},
-        )
-        ax.axis("equal")
-        _ = ax.axis([0, 1, 0, 1])
-        save_figure(
-            "consistent_solution.png", save_path=save_path, dpi=dpi, close_fig=close_fig
-        )
+        call_mud(lin_prob=lin_prob, save_path=save_path, dpi=dpi, close_fig=close_fig)
+
     if "comparison" in plot_fig or "all" in plot_fig:
-        fig, ax = plt.subplots(1, 1, figsize=(8, 8))
-        lin_prob.plot_fun_contours(ax=ax, terms="bayes", levels=50, cmap=cm.viridis)
-        lin_prob.plot_fun_contours(
-            ax=ax,
-            terms="data",
-            levels=25,
-            cmap=cm.viridis,
-            alpha=0.5,
-            vmin=0,
-            vmax=4,
-        )
-        lin_prob.plot_sol(
-            ax=ax,
-            point="initial",
-            pt_opts={
-                "color": "k",
-                "s": 100,
-                "marker": "o",
-                "label": "MUD",
-                "zorder": 10,
-            },
-        )
-        lin_prob.plot_sol(
-            ax=ax,
-            point="ls",
-            label="Least Squares",
-            note_loc=[0.49, 0.55],
-            pt_opts={"color": "k", "s": 100, "marker": "d", "zorder": 10},
-        )
-        lin_prob.plot_sol(
-            ax=ax,
-            point="map",
-            label="MAP",
-            pt_opts={
-                "color": "tab:orange",
-                "s": 100,
-                "linewidth": 3,
-                "marker": "x",
-                "zorder": 10,
-            },
-            ln_opts=None,
-            annotate_opts={"fontsize": 14, "backgroundcolor": "w"},
-        )
-        lin_prob.plot_sol(
-            point="mud",
-            ax=ax,
-            label="MUD",
-            pt_opts={
-                "color": "k",
-                "s": 100,
-                "linewidth": 3,
-                "marker": "*",
-                "zorder": 10,
-            },
-            ln_opts={"color": "k", "marker": "*", "lw": 1, "zorder": 10},
-            annotate_opts={"fontsize": 14, "backgroundcolor": "w"},
-        )
-        lin_prob.plot_contours(
-            ax=ax,
-            annotate=False,
-            note_loc=[0.1, 0.9],
-            label="Solution Contour",
-            plot_opts={"color": "r"},
-            annotate_opts={"fontsize": 14, "backgroundcolor": "w"},
-        )
-        pt = [0.7, 0.3]
-        ax.scatter([pt[0]], [pt[1]], color="k", s=100, marker="s", zorder=11)
-        nc = (pt[0] - 0.02, pt[1] + 0.02)
-        ax.annotate("Truth", nc, fontsize=14, backgroundcolor="w")
-        _ = ax.axis([0, 1, 0, 1])
-        save_figure(
-            "map_compare_contour.png", save_path=save_path, dpi=dpi, close_fig=close_fig
+        call_comparison(
+            lin_prob=lin_prob, save_path=save_path, dpi=dpi, close_fig=close_fig
         )
 
     return lin_prob
@@ -578,8 +601,8 @@ def run_wme_covariance(
     dim_output: int = 5,
     sigma: float = 1e-1,
     Ns: List[int] = [10, 100, 1000, 10000],
-    seed: int = None,
-    save_path: str = None,
+    seed: Optional[int] = None,
+    save_path: Optional[str] = None,
     dpi: int = 500,
     close_fig: bool = False,
 ):
@@ -676,12 +699,13 @@ def run_wme_covariance(
     _ = ax.set_xlabel("Index")
     _ = ax.set_ylabel("Eigenvalue")
     _ = ax.legend(loc="lower left")
-    save_figure(
-        "lin-meas-cov-sd-convergence.png",
-        save_path=save_path,
-        dpi=dpi,
-        close_fig=close_fig,
-    )
+    if save_path is not None:
+        save_figure(
+            "lin-meas-cov-sd-convergence.png",
+            save_path=save_path,
+            dpi=dpi,
+            close_fig=close_fig,
+        )
 
     return linear_wme_prob, ax
 
@@ -690,7 +714,7 @@ def run_high_dim_linear(
     dim_input: int = 100,
     dim_output: int = 100,
     seed: int = 21,
-    save_path: str = None,
+    save_path: Optional[str] = None,
     dpi: int = 500,
     close_fig: bool = True,
 ):
@@ -800,13 +824,13 @@ def err(xs):
     _ = ax.set_ylabel("Relative Error")
     _ = ax.set_xlabel("Dimension of Output Space")
     _ = ax.legend(["MUD", "MAP", "Least Squares"])
-
-    save_figure(
-        "lin-dim-cov-convergence.png",
-        save_path=save_path,
-        dpi=dpi,
-        close_fig=close_fig,
-    )
+    if save_path is not None:
+        save_figure(
+            "lin-dim-cov-convergence.png",
+            save_path=save_path,
+            dpi=dpi,
+            close_fig=close_fig,
+        )
 
     fig, ax = plt.subplots(1, 1, figsize=(10, 10))
     for idx, alpha in enumerate(alpha_list):
@@ -840,12 +864,12 @@ def err(xs):
     _ = ax.set_ylabel("Relative Error")
     _ = ax.set_xlabel("Rank(A)")
     _ = ax.legend(["MUD", "MAP", "Least Squares"])
-
-    save_figure(
-        "lin-rank-cov-convergence.png",
-        save_path=save_path,
-        dpi=dpi,
-        close_fig=close_fig,
-    )
+    if save_path is not None:
+        save_figure(
+            "lin-rank-cov-convergence.png",
+            save_path=save_path,
+            dpi=dpi,
+            close_fig=close_fig,
+        )
 
     return dim_errs, rank_errs
diff --git a/src/mud/examples/poisson.py b/src/mud/examples/poisson.py
index a1029b3..19a0b4b 100644
--- a/src/mud/examples/poisson.py
+++ b/src/mud/examples/poisson.py
@@ -6,7 +6,7 @@
 """
 import logging
 import random
-from typing import List, Union
+from typing import List, Optional, Union
 
 import matplotlib.pyplot as plt  # type: ignore
 import numpy as np
@@ -126,13 +126,99 @@ def spline_objective_function_2d(lam, aff=10000):
     return np.linalg.norm(g - vals)
 
 
+def plot_qoi_ex(mud_prob, num_components: int, **kwargs):
+    fig = plt.figure(figsize=(10, 5))
+    for i in range(num_components):
+        ax = fig.add_subplot(1, 2, i + 1)
+        mud_prob.plot_params_2d(ax=ax, y=i, contours=True, colorbar=True)
+        if i == 1:
+            ax.set_ylabel("")
+    if kwargs.get("save_path"):
+        save_figure("learned_qoi", **kwargs)
+    return ax
+
+
+def plot_response_ex(
+    poisson_prob,
+    closest,
+    raw_data,
+    sensors_mask,
+    group_idxs: List[List[int]] = [[0, 5], [5, 50], [50, -1]],
+    markers: List[str] = ["*", "+", "."],
+    colors: List[str] = ["red", "white", "k"],
+    **kwargs,
+):
+    fig = plt.figure(figsize=(10, 5))
+    ax = fig.add_subplot(1, 2, 1)
+
+    # Plot response surface from solving Eq. 30.
+    # u stores the mesh and values as solved by fenics
+    mesh, vals = raw_data["u"]
+    tcf = ax.tricontourf(mesh[:, 0], mesh[:, 1], vals, levels=20, vmin=-0.5, vmax=0)
+    fig.colorbar(tcf)
+
+    # Plot points used for each ordering
+    for idx, oi in enumerate(group_idxs):
+        start_idx, end_idx = group_idxs[idx][0], group_idxs[idx][1]
+        mask = sensors_mask[start_idx:end_idx]
+        poisson_prob.sensor_scatter_plot(
+            ax=ax,
+            mask=mask,
+            color=colors[idx],
+            marker=markers[idx],
+        )
+
+    # Label and format figure
+    _ = plt.xlim(0, 1)
+    _ = plt.ylim(0, 1)
+    _ = plt.title("Response Surface")
+    _ = plt.xlabel("$x_1$")
+    _ = plt.ylabel("$x_2$")
+
+    ax = fig.add_subplot(1, 2, 2)
+    # Plot closest solution in sample set to the reference solution
+    plot_solution_spline(
+        closest,
+        ax=ax,
+        lw=5,
+        c="green",
+        ls="--",
+        label=r"$\hat{g}(\lambda^\dagger)$",
+        zorder=50,
+    )
+    # Plot first 100 lambda initial
+    for i, lam in enumerate(poisson_prob.lam[0:50]):
+        plot_solution_spline(
+            lam,
+            plot_true=False,
+            ax=ax,
+            lw=1,
+            c="purple",
+            alpha=0.1,
+        )
+
+    plot_vert_line(ax, 1 / 3.0, color="b", linestyle="--", alpha=0.6)
+    plot_vert_line(ax, 2 / 3.0, color="b", linestyle="--", alpha=0.6)
+
+    ax.set_title("Boundary Condition")
+    ax.set_ylabel("")
+    ax.legend(
+        ["$g(x_2)$", r"$\hat{g}(x_2,\lambda^\dagger)$", r"$\hat{g}(x_2,\lambda_i)$"]
+    )
+
+    fig.tight_layout()
+    if kwargs.get("save_path"):
+        save_figure("response_surface", **kwargs)
+    return ax
+
+
 # TODO: Document group_idxs, markers, colors
 def run_2d_poisson_sol(
     data_file: str,
     sigma: float = 0.05,
-    seed: int = None,
+    seed: Optional[int] = None,
     plot_fig: Union[List[str], str] = "all",
-    save_path: str = None,
+    save_path: Optional[str] = None,
     dpi: int = 500,
     close_fig: bool = False,
     order: str = "random",
@@ -193,98 +279,47 @@ def run_2d_poisson_sol(
         method="pca", num_components=num_components, sensors_mask=idx_o
     )
     _ = mud_prob.estimate()
-    plot_fig = list(plot_fig) if type(plot_fig) != list else plot_fig
-    axes = []
-    if "response" in plot_fig or "all" in plot_fig:
-        fig = plt.figure(figsize=(10, 5))
-        ax = fig.add_subplot(1, 2, 1)
-
-        # Plot response surface from solving Eq. 30.
-        # u stores the mesh and values as solved by fenics
-        mesh, vals = raw_data["u"]
-        tcf = ax.tricontourf(mesh[:, 0], mesh[:, 1], vals, levels=20, vmin=-0.5, vmax=0)
-        fig.colorbar(tcf)
-
-        # Plot points used for each ordering
-        for idx, oi in enumerate(group_idxs):
-            start_idx, end_idx = group_idxs[idx][0], group_idxs[idx][1]
-            mask = idx_o[start_idx:end_idx]
-            poisson_prob.sensor_scatter_plot(
-                ax=ax,
-                mask=mask,
-                color=colors[idx],
-                marker=markers[idx],
-            )
 
-        # Label and format figure
-        _ = plt.xlim(0, 1)
-        _ = plt.ylim(0, 1)
-        _ = plt.title("Response Surface")
-        _ = plt.xlabel("$x_1$")
-        _ = plt.ylabel("$x_2$")
+    plot_fig = [plot_fig] if not isinstance(plot_fig, list) else plot_fig
 
-        ax = fig.add_subplot(1, 2, 2)
-        # Plot closest solution in sample set to the reference solution
-        plot_solution_spline(
-            closest,
-            ax=ax,
-            lw=5,
-            c="green",
-            ls="--",
-            label=r"$\hat{g}(\lambda^\dagger)$",
-            zorder=50,
-        )
-        # Plot first 100 lambda initial
-        for i, lam in enumerate(poisson_prob.lam[0:50]):
-            plot_solution_spline(
-                lam,
-                plot_true=False,
-                ax=ax,
-                lw=1,
-                c="purple",
-                alpha=0.1,
-            )
-
-        plot_vert_line(ax, 1 / 3.0, color="b", linestyle="--", alpha=0.6)
-        plot_vert_line(ax, 2 / 3.0, color="b", linestyle="--", alpha=0.6)
-
-        ax.set_title("Boundary Condition")
-        ax.set_ylabel("")
-        ax.legend(
-            ["$g(x_2)$", r"$\hat{g}(x_2,\lambda^\dagger)$", r"$\hat{g}(x_2,\lambda_i)$"]
-        )
-
-        fig.tight_layout()
-
-        save_figure(
-            "response_surface",
-            save_path,
+    axes = []
+    if "response" in plot_fig or "all" in plot_fig:
+        ax = plot_response_ex(
+            poisson_prob,
+            closest=closest,
+            raw_data=raw_data,
+            group_idxs=group_idxs,
+            markers=markers,
+            sensors_mask=idx_o,
+            colors=colors,
+            save_path=save_path,
             close_fig=close_fig,
             dpi=dpi,
             bbox_inches="tight",
         )
         axes.append(ax)
     if "qoi" in plot_fig or "all" in plot_fig:
-        fig = plt.figure(figsize=(10, 5))
-        for i in range(num_components):
-            ax = fig.add_subplot(1, 2, i + 1)
-            mud_prob.plot_params_2d(ax=ax, y=i, contours=True, colorbar=True)
-            if i == 1:
-                ax.set_ylabel("")
-        save_figure(
-            "learned_qoi", save_path, close_fig=close_fig, dpi=dpi, bbox_inches="tight"
+        ax = plot_qoi_ex(
+            mud_prob,
+            num_components=num_components,
+            save_path=save_path,
+            close_fig=close_fig,
+            dpi=dpi,
+            bbox_inches="tight",
         )
         axes.append(ax)
     if "densities" in plot_fig or "all" in plot_fig:
         ax1 = mud_prob.plot_param_space(param_idx=0, **param1_kwargs)
-        save_figure(
-            "lam1", save_path, close_fig=close_fig, dpi=dpi, bbox_inches="tight"
-        )
+        if save_path is not None:
+            save_figure(
+                "lam1", save_path, close_fig=close_fig, dpi=dpi, bbox_inches="tight"
+            )
 
         ax2 = mud_prob.plot_param_space(param_idx=1, **param2_kwargs)
-        save_figure(
-            "lam2", save_path, close_fig=close_fig, dpi=dpi, bbox_inches="tight"
-        )
+        if save_path is not None:
+            save_figure(
+                "lam2", save_path, close_fig=close_fig, dpi=dpi, bbox_inches="tight"
+            )
         axes.append([ax1, ax2])
 
     return (poisson_prob, mud_prob, axes)
@@ -301,8 +336,8 @@ def run_2d_poisson_trials(
     annotate_location_1: List[float] = [-2.6, 1.2, 0.8],
     annotate_location_2: List[float] = [-3.5, 0.83, 0.53],
     sigma: float = 0.05,
-    seed: int = None,
-    save_path: str = None,
+    seed: Optional[int] = None,
+    save_path: Optional[str] = None,
     dpi: int = 500,
     close_fig: bool = False,
 ):
@@ -384,7 +419,7 @@ def run_2d_poisson_trials(
         mud_prob.plot_param_space(
             ax=ax1, x_range=x_range, param_idx=0, mud_opts=None, true_opts=None
         )
-        ax1.set_ylim(ylim1)
+        ax1.set_ylim(*ylim1)
         mud_prob.plot_param_space(ax=ax1, true_val=closest, in_opts=None, up_opts=None)
         ax1.set_xlabel(r"$\lambda_1$")
         # annotate_location_1 = [-2.8, 1.2, 0.8]
@@ -405,7 +440,7 @@ def run_2d_poisson_trials(
             mud_opts=None,
             true_opts=None,
         )
-        ax2.set_ylim(ylim2)
+        ax2.set_ylim(*ylim2)
         mud_prob.plot_param_space(
             ax=ax2, param_idx=1, true_val=closest, in_opts=None, up_opts=None
         )
@@ -419,14 +454,14 @@ def run_2d_poisson_trials(
             e_r = mud_prob.expected_ratio()
             ax2.text(x, y2, rf"$\mathbb{{E}}(r) = {e_r:0.4}$", fontsize=18)
         ax2.legend()
-
-        save_figure(
-            f"solution_n{N}",
-            save_path,
-            close_fig=close_fig,
-            dpi=dpi,
-            bbox_inches="tight",
-        )
+        if save_path is not None:
+            save_figure(
+                f"solution_n{N}",
+                save_path,
+                close_fig=close_fig,
+                dpi=dpi,
+                bbox_inches="tight",
+            )
 
         axes.append([ax1, ax2])
         probs.append(mud_prob)
diff --git a/src/mud/plot.py b/src/mud/plot.py
index eebfcfa..672fbff 100644
--- a/src/mud/plot.py
+++ b/src/mud/plot.py
@@ -54,7 +54,9 @@ def _check_latex():
         plt.rcParams.update(mud_plot_params)
 
 
-def save_figure(fname: str, save_path: str = None, close_fig: bool = True, **kwargs):
+def save_figure(
+    fname: str, save_path: str = "figures", close_fig: bool = True, **kwargs
+):
     """
     Save Figure Utility
 
@@ -64,9 +66,8 @@ def save_figure(fname: str, save_path: str = None, close_fig: bool = True, **kwa
     ----------
     fname: str
         Name of image, with extension.
-    save_path: str, optional
-        Directory to save figure to. Assumed to exist. If not specified then the
-        figure is saved to the current working directory.
+    save_path: str
+        Directory to save figure to. Assumed to exist. Default: figures/
     close_fig: bool, default=True
         Whether to close the figure after saving it.
     kwargs: dict, optional
@@ -76,9 +77,9 @@ def save_figure(fname: str, save_path: str = None, close_fig: bool = True, **kwa
     """
     global mud_plot_params
 
-    if save_path is not None:
-        fname = str(Path(save_path) / Path(fname))
-        plt.savefig(fname, **kwargs)
+    fname = str(Path(save_path) / Path(fname))
+    plt.savefig(fname, **kwargs)
+
     if close_fig:
         plt.close()
 
diff --git a/src/mud/util.py b/src/mud/util.py
index e915f8e..fad2040 100644
--- a/src/mud/util.py
+++ b/src/mud/util.py
@@ -1,4 +1,4 @@
-from typing import List, Tuple, Union
+from typing import List, Optional, Tuple, Union
 
 import numpy as np
 from numpy.typing import ArrayLike
@@ -183,7 +183,7 @@ def make_2d_unit_mesh(N: int = 50, window: int = 1):
     return (X, Y, XX)
 
 
-def add_noise(signal: ArrayLike, sd: float = 0.05, seed: int = None):
+def add_noise(signal: ArrayLike, sd: float = 0.05, seed: Optional[int] = None):
     """
     Add Noise
 
@@ -241,7 +241,9 @@ def rank_decomposition(A: np.typing.ArrayLike) -> List[np.ndarray]:
 
 
 def fit_domain(
-    x: np.ndarray = None, min_max_bounds: np.ndarray = None, pad_ratio: float = 0.1
+    x: Optional[np.ndarray] = None,
+    min_max_bounds: Optional[np.ndarray] = None,
+    pad_ratio: float = 0.1,
 ) -> np.ndarray:
     """
     Fit domain bounding box to array x
diff --git a/tests/test_cli.py b/tests/test_cli.py
index e2ec68a..de66088 100644
--- a/tests/test_cli.py
+++ b/tests/test_cli.py
@@ -12,27 +12,33 @@
 from mud.cli import cli
 
 
-def test_comparison_example():
+def test_comparison_example(test_dir):
     runner = CliRunner()
-    result = runner.invoke(cli, ["examples", "-ns", "comparison"])
+    result = runner.invoke(
+        cli, ["examples", "--save-path", str(test_dir), "comparison"]
+    )
     assert result.exit_code == 0
 
 
-def test_contours_example():
+def test_contours_example(test_dir):
     runner = CliRunner()
-    result = runner.invoke(cli, ["examples", "-ns", "contours"])
+    result = runner.invoke(cli, ["examples", "--save-path", str(test_dir), "contours"])
     assert result.exit_code == 0
 
 
-def test_high_dim_linear():
+def test_high_dim_linear(test_dir):
     runner = CliRunner()
-    result = runner.invoke(cli, ["examples", "-ns", "high-dim-linear"])
+    result = runner.invoke(
+        cli, ["examples", "--save-path", str(test_dir), "high-dim-linear"]
+    )
     assert result.exit_code == 0
 
 
-def test_wme_covariance():
+def test_wme_covariance(test_dir):
     runner = CliRunner()
-    result = runner.invoke(cli, ["examples", "-ns", "wme-covariance"])
+    result = runner.invoke(
+        cli, ["examples", "--save-path", str(test_dir), "wme-covariance"]
+    )
     assert result.exit_code == 0
 
 
@@ -63,34 +69,55 @@ def test_poisson_generate(test_dir):
         assert np.abs(data["true_vals"][1] - 0.00183782) < 0.001
 
 
-def test_poisson_solve():
+def test_poisson_solve(test_dir):
     runner = CliRunner()
     data = str(Path(__file__).parent / "data" / "poisson_data")
     result = runner.invoke(
-        cli, ["examples", "-ns", "--seed", "21", "poisson-solve", data]
+        cli,
+        [
+            "examples",
+            "--save-path",
+            str(test_dir),
+            "--seed",
+            "21",
+            "poisson-solve",
+            data,
+        ],
     )
     assert result.exit_code == 0
     assert str(result.stdout) == "[-2.76754243 -1.6656349 ]\n"
 
 
-def test_poisson_trials():
+def test_poisson_trials(test_dir):
     runner = CliRunner()
     data = str(Path(__file__).parent / "data" / "poisson_data")
     result = runner.invoke(
-        cli, ["examples", "-ns", "--seed", "21", "poisson-trials", data, "-n", "2"]
+        cli,
+        [
+            "examples",
+            "--save-path",
+            str(test_dir),
+            "--seed",
+            "21",
+            "poisson-trials",
+            data,
+            "-n",
+            "2",
+        ],
     )
     assert result.exit_code == 0
     assert "0.018693404000" in str(result.stdout)
 
 
-def test_adcirc_solve():
+def test_adcirc_solve(test_dir):
     runner = CliRunner()
     data = str(Path(__file__).parent / "data" / "adcirc_data")
     result = runner.invoke(
         cli,
         [
             "examples",
-            "-ns",
+            "--save-path",
+            str(test_dir),
             "--seed",
             "21",
             "adcirc-solve",