From 94084e7c8cc24ea59f8fddd8e0693ae02d684757 Mon Sep 17 00:00:00 2001 From: "Michael Pilosov, PhD" <40366263+mathematicalmichael@users.noreply.github.com> Date: Sun, 21 Jan 2024 22:13:25 -0700 Subject: [PATCH] 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 | 4 +- .github/workflows/main.yml | 2 +- src/mud/base.py | 32 +- src/mud/examples/adcirc.py | 86 ++--- src/mud/examples/comparison.py | 49 ++- src/mud/examples/linear.py | 559 +++++++++++++++++---------------- src/mud/examples/poisson.py | 217 +++++++------ src/mud/plot.py | 15 +- src/mud/util.py | 8 +- 9 files changed, 521 insertions(+), 451 deletions(-) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index c5de5f3..3d551a7 100644 --- a/.github/workflows/build.yml +++ b/.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/ diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index 410d7dc..5313a1d 100644 --- a/.github/workflows/main.yml +++ b/.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 diff --git a/src/mud/base.py b/src/mud/base.py index 21c2bf0..bc9b719 100644 --- a/src/mud/base.py +++ b/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): diff --git a/src/mud/examples/adcirc.py b/src/mud/examples/adcirc.py index 27345d8..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,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 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 3d86325..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 @@ -99,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""" @@ -263,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, @@ -322,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 @@ -577,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, ): @@ -675,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 @@ -689,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, ): @@ -799,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): @@ -839,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 b1ac3d0..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", @@ -198,95 +284,42 @@ def run_2d_poisson_sol( 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$") - - 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, + 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) @@ -303,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, ): @@ -386,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] @@ -407,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 ) @@ -421,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