diff --git a/.gitignore b/.gitignore
index e4bc802..9de7d08 100644
--- a/.gitignore
+++ b/.gitignore
@@ -2,12 +2,15 @@
 __pycache__/
 *.nc
 *.pdf
+*.zip
 
 # Sphinx
 docs/_build/
 docs/api/resmda*
 docs/savefig/
 docs/gallery/*
+download/
+docs/sg_execution_times.rst
 
 # Pytest and coverage related
 htmlcov
diff --git a/docs/conf.py b/docs/conf.py
index eb39077..e74345b 100644
--- a/docs/conf.py
+++ b/docs/conf.py
@@ -1,7 +1,6 @@
 import time
 import warnings
 from resmda import __version__
-from sphinx_gallery.sorting import FileNameSortKey
 
 # ==== 1. Extensions  ====
 
@@ -39,17 +38,17 @@
 sphinx_gallery_conf = {
     'examples_dirs': ['../examples', ],
     'gallery_dirs': ['gallery', ],
-    'capture_repr': ('_repr_html_', '__repr__'),
+    'capture_repr': ('_repr_html_', ),
     # Patter to search for example files
     "filename_pattern": r"\.py",
     # Sort gallery example by file name instead of number of lines (default)
-    "within_subsection_order": FileNameSortKey,
+    "within_subsection_order": "FileNameSortKey",
     # Remove the settings (e.g., sphinx_gallery_thumbnail_number)
     'remove_config_comments': True,
     # Show memory
     'show_memory': True,
     # Custom first notebook cell
-    'first_notebook_cell': '%matplotlib notebook',
+    'first_notebook_cell': '%matplotlib widget',
 }
 
 # https://github.com/sphinx-gallery/sphinx-gallery/pull/521/files
diff --git a/examples/basicESMDA.py b/examples/basicESMDA.py
index f7e525a..5562435 100644
--- a/examples/basicESMDA.py
+++ b/examples/basicESMDA.py
@@ -19,8 +19,8 @@
 .. math::
     x_{j,i+1} &= x_{j,i} + (C^e_{xy})_i \left((C^e_{yy})_i +
                  \alpha_iC^e_{dd}\right)^{-1} \left(d + \sqrt{\alpha_i}
-                 \varepsilon_j - g(x_{j,i})\right) \\
-    y_{j,i+1} &= g(x_{j,i+1})
+                 \varepsilon_j - g(x_{j,i})\right) \ , \\
+    y_{j,i+1} &= g(x_{j,i+1}) \ .
 
 The model used for this example is
 
@@ -53,7 +53,7 @@ def forward(x, beta):
 px = np.linspace(-5, 5, 301)
 for i, b in enumerate([0.0, 0.2]):
     axs[i].set_title(
-            f"{['Linear model', 'Nonlinear model'][i]}: $\\beta$ = {b}")
+            f"{['Linear model', 'Nonlinear model'][i]}: β = {b}")
     axs[i].plot(px, forward(px, b))
     axs[i].set_xlabel('x')
     axs[i].set_ylabel('y')
@@ -109,7 +109,6 @@ def plot_result(mpost, dpost, dobs, title, ylim):
     ax2.set_xlabel('y')
     ax2.set_xlim([-5, 8])
     ax2.legend()
-    fig.show()
 
 
 ###############################################################################
diff --git a/examples/basicreservoir.py b/examples/basicreservoir.py
index 767052a..d786718 100644
--- a/examples/basicreservoir.py
+++ b/examples/basicreservoir.py
@@ -21,7 +21,7 @@
 # ----------------
 
 # Grid extension
-nx = 25
+nx = 30
 ny = 25
 nc = nx*ny
 
@@ -32,7 +32,6 @@
 
 # ESMDA parameters
 ne = 100                  # Number of ensembles
-#  dt = np.logspace(-5, -3, 10)
 dt = np.zeros(10)+0.0001  # Time steps (could be irregular, e.g., increasing!)
 time = np.r_[0, np.cumsum(dt)]
 nt = time.size
@@ -61,25 +60,23 @@
 perm_prior = RP(ne, random=rng)
 
 
-# TODO: change scale in imshow to represent meters
-
 # QC covariance, reference model, and first two random models
-fig, axs = plt.subplots(2, 2, constrained_layout=True)
-axs[0, 0].set_title('Model')
-im = axs[0, 0].imshow(perm_true.T, vmin=perm_min, vmax=perm_max)
-axs[0, 1].set_title('Lower Covariance Matrix')
-im2 = axs[0, 1].imshow(RP.cov, cmap='plasma')
-axs[1, 0].set_title('Random Model 1')
-axs[1, 0].imshow(perm_prior[0, ...].T, vmin=perm_min, vmax=perm_max)
-axs[1, 1].set_title('Random Model 2')
-axs[1, 1].imshow(perm_prior[1, ...].T, vmin=perm_min, vmax=perm_max)
-fig.colorbar(im, ax=axs[1, :], orientation='horizontal',
-             label='Log of Permeability (mD)')
-for ax in axs[1, :]:
-    ax.set_xlabel('x-direction')
-for ax in axs[:, 0]:
-    ax.set_ylabel('y-direction')
-fig.show()
+pinp1 = {"origin": "lower", "vmin": perm_min, "vmax": perm_max}
+fig, axs = plt.subplots(2, 2, figsize=(6, 6), constrained_layout=True)
+axs[0, 0].set_title("Model")
+im = axs[0, 0].imshow(perm_true.T, **pinp1)
+axs[0, 1].set_title("Lower Covariance Matrix")
+im2 = axs[0, 1].imshow(RP.cov, cmap="plasma")
+axs[1, 0].set_title("Random Model 1")
+axs[1, 0].imshow(perm_prior[0, ...].T, **pinp1)
+axs[1, 1].set_title("Random Model 2")
+axs[1, 1].imshow(perm_prior[1, ...].T, **pinp1)
+fig.colorbar(im, ax=axs[1, :], orientation="horizontal",
+             label="Log of Permeability (mD)")
+for ax in axs[1, :].ravel():
+    ax.set_xlabel("x-direction")
+for ax in axs[:, 0].ravel():
+    ax.set_ylabel("y-direction")
 
 
 ###############################################################################
@@ -101,15 +98,14 @@ def sim(x):
 data_obs = rng.normal(data_true, dstd)
 
 # QC data and priors
-fig, ax = plt.subplots(1, 1)
-ax.set_title('Observed and prior data')
-ax.plot(time, data_prior.T, color='.6', alpha=0.5)
-ax.plot(time, data_true, 'ko', label='True data')
-ax.plot(time, data_obs, 'C3o', label='Obs. data')
+fig, ax = plt.subplots(1, 1, constrained_layout=True)
+ax.set_title("Observed and prior data")
+ax.plot(time*24*60*60, data_prior.T, color=".6", alpha=0.5)
+ax.plot(time*24*60*60, data_true, "ko", label="True data")
+ax.plot(time*24*60*60, data_obs, "C3o", label="Obs. data")
 ax.legend()
-ax.set_xlabel('Time (???)')
-ax.set_ylabel('Pressure (???)')
-fig.show()
+ax.set_xlabel("Time (s)")
+ax.set_ylabel("Pressure (bar)")
 
 
 ###############################################################################
@@ -143,14 +139,14 @@ def restrict_permeability(x):
 # and posterior ensembles to see how the models have been updated.
 
 # Plot posterior
-fig, ax = plt.subplots(1, 3, figsize=(12, 5))
-ax[0].set_title('Prior Mean')
-ax[0].imshow(perm_prior.mean(axis=0).T)
-ax[1].set_title('Post Mean')
-ax[1].imshow(perm_post.mean(axis=0).T)
-ax[2].set_title('"Truth"')
-ax[2].imshow(perm_true.T)
-fig.show()
+fig, ax = plt.subplots(1, 2, figsize=(8, 5), constrained_layout=True)
+pinp2 = {"origin": "lower", "vmin": 2.5, "vmax": 3.5}
+ax[0].set_title("Prior Mean")
+im = ax[0].imshow(perm_prior.mean(axis=0).T, **pinp2)
+ax[1].set_title("Post Mean")
+ax[1].imshow(perm_post.mean(axis=0).T, **pinp2)
+fig.colorbar(im, ax=ax, label="Log of Permeability (mD)",
+             orientation="horizontal")
 
 
 ###############################################################################
@@ -162,15 +158,14 @@ def restrict_permeability(x):
 
 
 # Compare posterior to prior and observed data
-fig, ax = plt.subplots(1, 1)
-ax.set_title('Prior and posterior data')
-ax.plot(time, data_prior.T, color='.6', alpha=0.5)
-ax.plot(time, data_post.T, color='C0', alpha=0.5)
-ax.plot(time, data_true, 'ko')
-ax.plot(time, data_obs, 'C3o')
-ax.set_xlabel('Time (???)')
-ax.set_ylabel('Pressure (???)')
-fig.show()
+fig, ax = plt.subplots(1, 1, constrained_layout=True)
+ax.set_title("Prior and posterior data")
+ax.plot(time*24*60*60, data_prior.T, color=".6", alpha=0.5)
+ax.plot(time*24*60*60, data_post.T, color="C0", alpha=0.5)
+ax.plot(time*24*60*60, data_true, "ko")
+ax.plot(time*24*60*60, data_obs, "C3o")
+ax.set_xlabel("Time (s)")
+ax.set_ylabel("Pressure (bar)")
 
 
 ###############################################################################
diff --git a/examples/fluvialreservoir.py b/examples/fluvialreservoir.py
new file mode 100644
index 0000000..14fb3ef
--- /dev/null
+++ b/examples/fluvialreservoir.py
@@ -0,0 +1,393 @@
+r"""
+2D Fluvial Reservoir ESMDA example
+==================================
+
+In contrast to the basic reservoir example
+:ref:`sphx_glr_gallery_basicreservoir.py`, where a single facies was used, this
+example uses fluvial models containing different facies. It also compares the
+use of ES-MDA with and without localization, as explained in the example
+:ref:`sphx_glr_gallery_localization.py`.
+
+The fluvial models were generated with ``FLUVSIM`` through ``geomodpy``, for
+more information see towards the end of the example where the code is shown to
+reproduce the facies.
+
+.. note::
+
+    To retrieve the data, you need to have ``pooch`` installed:
+
+    .. code-block:: bash
+
+        pip install pooch
+
+    or
+
+    .. code-block:: bash
+
+        conda install -c conda-forge pooch
+
+"""
+import os
+import json
+
+import pooch
+import numpy as np
+import matplotlib.pyplot as plt
+
+import resmda
+
+# For reproducibility, we instantiate a random number generator with a fixed
+# seed. For production, remove the seed!
+rng = np.random.default_rng(1513)
+
+# Adjust this path to a folder of your choice.
+data_path = os.path.join("..", "download", "")
+
+# sphinx_gallery_thumbnail_number = 3
+
+###############################################################################
+# Load and plot the facies
+# ------------------------
+
+fname = "facies.npy"
+pooch.retrieve(
+    "https://raw.github.com/tuda-geo/data/2024-06-18/resmda/"+fname,
+    "4bfe56c836bf17ca63453c37e5da91cb97bbef8cc6c08d605f70bd64fe7488b2",
+    fname=fname,
+    path=data_path,
+)
+facies = np.load(data_path + fname)
+ne, nx, ny = facies.shape
+
+# Define mean permeability per facies
+perm_means = [0.1, 5.0, 3.0]
+
+# Plot the facies
+fig, axs = plt.subplots(
+    2, 5, figsize=(8, 3), sharex=True, sharey=True, constrained_layout=True)
+axs = axs.ravel()
+fig.suptitle(f"Facies {[f'{i} = {p}' for i, p in enumerate(perm_means)]}")
+for i in range(ne):
+    im = axs[i].imshow(
+        facies[i, ...], cmap=plt.get_cmap("Accent", 3),
+        clim=[-0.5, 2.5], origin="lower"
+    )
+fig.colorbar(im, ax=axs, ticks=[0, 1, 2], label="Facies code")
+
+
+###############################################################################
+# Assign random permeabilities to the facies
+# ------------------------------------------
+
+perm_min = 0.05
+perm_max = 10.0
+
+# Instantiate a random permeability instance
+RP = resmda.RandomPermeability(
+    nx, ny, perm_mean=None, perm_min=perm_min, perm_max=perm_max
+)
+
+# Fill the different facies with their permeabilities
+permeabilities = np.empty(facies.shape)
+for code, mean in enumerate(perm_means):
+    mask = facies == code
+    permeabilities[mask] = RP(ne, perm_mean=mean)[mask]
+
+fig, axs = plt.subplots(
+    2, 5, figsize=(8, 3), sharex=True, sharey=True, constrained_layout=True)
+axs = axs.ravel()
+fig.suptitle("Permeabilities")
+for i in range(ne):
+    im = axs[i].imshow(permeabilities[i, ...], origin="lower")
+fig.colorbar(im, ax=axs, label="Log Permeability (mD)")
+
+
+###############################################################################
+# Model parameters
+# ----------------
+
+# We take the first model as "true/reference", and the other for ES-MDA.
+perm_true = permeabilities[0, ...][None, ...]
+perm_prior = permeabilities[1:, ...]
+
+# Time steps
+dt = np.zeros(10)+0.0001
+time = np.r_[0, np.cumsum(dt)]
+nt = time.size
+
+# Assumed standard deviation of our data
+dstd = 0.5
+
+# Measurement points
+ox = (5, 15, 24)
+oy = (5, 10, 24)
+
+# Number of data points
+nd = time.size * len(ox)
+
+# Wells
+wells = np.array([
+    [ox[0], oy[0], 180], [5, 12, 120],
+    [ox[1], oy[1], 180], [20, 5, 120],
+    [ox[2], oy[2], 180], [24, 17, 120]
+])
+
+
+###############################################################################
+# Run the prior models and the reference case
+# -------------------------------------------
+
+# Instantiate reservoir simulator
+RS = resmda.Simulator(nx, ny, wells=wells)
+
+
+def sim(x):
+    """Custom fct to use exp(x), and specific dt & location."""
+    return RS(np.exp(x), dt=dt, data=(ox, oy)).reshape((x.shape[0], -1))
+
+
+# Simulate data for the prior and true fields
+data_prior = sim(perm_prior)
+data_true = sim(perm_true)
+data_obs = rng.normal(data_true, dstd)
+data_obs[0, :3] = data_true[0, :3]
+
+
+###############################################################################
+# Localization Matrix
+# -------------------
+
+# Vector of nd length with the well x and y position for each nd data point
+nd_positions = np.tile(np.array([ox, oy]), time.size).T
+
+# Create matrix
+loc_mat = resmda.build_localization_matrix(RP.cov, nd_positions, (nx, ny))
+
+
+###############################################################################
+# ES-MDA
+# ------
+
+
+def restrict_permeability(x):
+    """Restrict possible permeabilities."""
+    np.clip(x, perm_min, perm_max, out=x)
+
+
+inp = {
+    "model_prior": perm_prior,
+    "forward": sim,
+    "data_obs": data_obs,
+    "sigma": dstd,
+    "alphas": 4,
+    "data_prior": data_prior,
+    "callback_post": restrict_permeability,
+    "random": rng,
+}
+
+
+###############################################################################
+# Without localization
+# ''''''''''''''''''''
+
+nl_perm_post, nl_data_post = resmda.esmda(**inp)
+
+
+###############################################################################
+# With localization
+# '''''''''''''''''
+
+wl_perm_post, wl_data_post = resmda.esmda(**inp, localization_matrix=loc_mat)
+
+
+###############################################################################
+# Compare permeabilities
+# ----------------------
+
+# Plot posterior
+fig, axs = plt.subplots(
+    1, 3, figsize=(8, 4), sharex=True, sharey=True, constrained_layout=True)
+
+par = {"vmin": perm_min, "vmax": perm_max, "origin": "lower"}
+
+axs[0].set_title("Prior Mean")
+im = axs[0].imshow(perm_prior.mean(axis=0).T, **par)
+
+
+axs[1].set_title("Post Mean; No localization")
+axs[1].imshow(nl_perm_post.mean(axis=0).T, **par)
+
+axs[2].set_title("Post Mean: Localization")
+axs[2].imshow(wl_perm_post.mean(axis=0).T, **par)
+axs[2].contour(loc_mat.sum(axis=2).T, levels=[2.0, ], colors="w")
+
+fig.colorbar(im, ax=axs, label="Log Permeabilities (mD)",
+             orientation="horizontal")
+
+for ax in axs:
+    for well in wells:
+        ax.plot(well[0], well[1], ["C3v", "C1^"][int(well[2] == 120)])
+for ax in axs:
+    ax.set_xlabel('x-direction')
+axs[0].set_ylabel('y-direction')
+
+
+###############################################################################
+# Compare data
+# ------------
+
+# QC data and priors
+fig, axs = plt.subplots(
+    2, 3, figsize=(8, 5), sharex=True, sharey=True, constrained_layout=True)
+fig.suptitle("Prior and posterior data")
+for i, ax in enumerate(axs[0, :]):
+    ax.set_title(f"Well ({ox[i]}; {oy[i]})")
+    ax.plot(time*24*60*60, data_prior[..., i::3].T, color=".6", alpha=0.5)
+    ax.plot(time*24*60*60, nl_data_post[..., i::3].T, color="C0", alpha=0.5)
+    ax.plot(time*24*60*60, data_obs[0, i::3], "C3o")
+for i, ax in enumerate(axs[1, :]):
+    ax.plot(time*24*60*60, data_prior[..., i::3].T, color=".6", alpha=0.5)
+    ax.plot(time*24*60*60, wl_data_post[..., i::3].T, color="C0", alpha=0.5)
+    ax.plot(time*24*60*60, data_obs[0, i::3], "C3o")
+    ax.set_xlabel("Time (s)")
+for i, ax in enumerate(axs[:, 0]):
+    ax.set_ylabel("Pressure (bar)")
+for i, txt in enumerate(["No l", "L"]):
+    axs[i, 2].yaxis.set_label_position("right")
+    axs[i, 2].set_ylabel(f"{txt}ocalization")
+
+
+###############################################################################
+# Reproduce the facies
+# --------------------
+#
+# .. note::
+#
+#     The following cell is about how to reproduce the facies data loaded in
+#     ``facies.npy``. This was created using ``geomodpy``.
+#
+#     ``geomodpy`` (Guillaume Rongier, 2023) is not open-source yet. The
+#     functionality of geomodpy that we use here is a python wrapper for the
+#     Fortran library ``FLUVSIM`` published in:
+#
+#         **Deutsch, C. V., and T. T. Tran**, 2002, FLUVSIM: a program for
+#         object-based stochastic modeling of fluvial depositional systems:
+#         Computers & Geosciences, 28(4), 525--535.
+#
+#         DOI: `10.1016/S0098-3004(01)00075-9
+#         <https://doi.org/10.1016/S0098-3004(01)00075-9>`_.
+#
+#
+# .. code-block:: python
+#
+#     # ==== Required imports ====
+#     import json
+#     import numpy as np
+#
+#     # FLUVSIM Version used: 2.900
+#     from geomodpy.wrapper.fluvsim import FLUVSIM
+#
+#     # For reproducibility, we instantiate a random number generator with a
+#     # fixed seed. For production, remove the seed!
+#     rng = np.random.default_rng(1848)
+#
+#
+#     # ==== Define the geological parameters ====
+#
+#     # Here we define the geological parameters by means of their normal
+#     # distribution parameters
+#
+#     # Each tuple stands for (mean, std); lists contain several of them.
+#     geol_distributions = {
+#         "channel_orientation": (60, 20),
+#         "channel_amplitude": [(100, 1), (250, 1), (400, 1)],
+#         "channel_wavelength": [(1000, 5), (2000, 5), (3000, 5)],
+#         "channel_thickness": [(4, 0.1), (8, 0.1), (11, 0.1)],
+#         "channel_thickness_undulation": (1, 0.02),
+#         "channel_thickness_undulation_wavelength": [
+#             (250, 1), (400, 1), (450, 1)
+#         ],
+#         "channel_width_thickness_ratio": [(40, 0.5), (50, 0.5), (60, 0.5)],
+#         "channel_width_undulation": (1, 0.02),
+#         "channel_width_undulation_wavelength": (250, 1),
+#         "channel_prop": (0.4, 0.005),
+#     }
+#
+#
+#     def generate_geol_params(geol_dists):
+#         """Generate geological parameters from normal distributions.
+#
+#         Expects for each parameter a tuple of two values, or a list
+#         containing tuples of two values each.
+#         """
+#         geol_params = {}
+#         for param, dist in geol_dists.items():
+#             if isinstance(dist, list):
+#                 geol_params[param] = tuple(
+#                     rng.normal(mean, std) for mean, std in dist
+#                 )
+#             else:
+#                 geol_params[param] = rng.normal(*dist)
+#         return geol_params
+#
+#
+#     # ==== Create the facies ====
+#
+#     # Number of sets and realizations
+#     nsets = 2
+#     nreal = 5
+#
+#     # Model dimension
+#     nx, ny, nz = 64, 64, 1
+#
+#     # Pre-allocate containers to store all realizations and their
+#     # corresponding parameters
+#     all_params = {}
+#     facies = np.zeros((nsets * nreal, nz, nx, ny), dtype="i4")
+#
+#     for i in range(nsets):  # We create two sets
+#         print(f"Generating realization {i+1} of {nsets}")
+#
+#         params = generate_geol_params(geol_distributions)
+#         all_params[f"set-{i}"] = params
+#
+#         fluvsim = FLUVSIM(
+#             shape=(nx, ny, nz),
+#             spacing=(50, 50, 1),
+#             origin=(25, 25, 0.5),
+#             n_realizations=nreal,
+#             **params
+#         )
+#
+#         realizations = fluvsim.run().data_vars["facies code"].values
+#         facies[i*nreal:(i+1)*nreal, ...] = realizations.astype("i4")
+#
+#
+#     # ==== Save the output ====
+#
+#     # Save the input parameters to FLUVSIM as a json.
+#     with open("facies.json", "w") as f:
+#         json.dump(all_params, f, indent=2)
+#     # Save the facies values as a compressed npy-file.
+#     np.save("facies", facies.squeeze(), allow_pickle=False)
+
+
+###############################################################################
+# Input parameters to ``FLUVSIM``
+# -------------------------------
+# For reproducibility purposes we print here the used input values to FLUVSIM.
+# These are, just as the data themselves, online at
+# https://github.com/tuda-geo/data/resmda.
+
+fname = "facies.json"
+pooch.retrieve(
+    "https://raw.github.com/tuda-geo/data/2024-06-18/resmda/"+fname,
+    "db2cb8a620775c68374c24a4fa811f6350381c7fc98a823b9571136d307540b4",
+    fname=fname,
+    path=data_path,
+)
+with open(data_path + fname, "r") as f:
+    print(json.dumps(json.load(f), indent=2))
+
+###############################################################################
+resmda.Report()
diff --git a/examples/localization.py b/examples/localization.py
index 8d42bf7..fa6ed6f 100644
--- a/examples/localization.py
+++ b/examples/localization.py
@@ -22,7 +22,7 @@
 # ----------------
 
 # Grid extension
-nx = 30
+nx = 35
 ny = 30
 nc = nx*ny
 
@@ -39,7 +39,7 @@
 # Assumed sandard deviation of our data
 dstd = 0.5
 
-# Observation location indices (should be well locations)
+# Measurement points
 ox = (5, 15, 24)
 oy = (5, 10, 24)
 
@@ -48,9 +48,9 @@
 
 # Wells
 wells = np.array([
-    [5, 5, 180], [5, 12, 120],
-    [15, 10, 180], [20, 5, 120],
-    [24, 24, 180], [24, 17, 120]
+    [ox[0], oy[0], 180], [5, 12, 120],
+    [ox[1], oy[1], 180], [20, 5, 120],
+    [ox[2], oy[2], 180], [24, 17, 120]
 ])
 
 
@@ -99,11 +99,14 @@ def sim(x):
 loc_mat = resmda.build_localization_matrix(RP.cov, nd_positions, (nx, ny))
 
 # QC localization matrix
-fig, ax = plt.subplots(1, 1, sharex=True, sharey=True, constrained_layout=True)
-ax.imshow(loc_mat.sum(axis=2).T)
-ax.contour(loc_mat.sum(axis=2).T, levels=[2.0, ], colors='w')
+fig, ax = plt.subplots(1, 1, constrained_layout=True)
+im = ax.imshow(loc_mat.sum(axis=2).T/time.size, origin='lower', cmap="plasma")
+ax.contour(loc_mat.sum(axis=2).T/time.size, levels=[0.2, ], colors='w')
+ax.set_xlabel('x-direction')
+ax.set_ylabel('y-direction')
 for well in wells:
     ax.plot(well[0], well[1], ['C3v', 'C1^'][int(well[2] == 120)])
+fig.colorbar(im, ax=ax, label="Localization Weight (-)")
 
 
 ###############################################################################
@@ -148,22 +151,30 @@ def restrict_permeability(x):
 
 # Plot posterior
 fig, axs = plt.subplots(
-    2, 2, figsize=(6, 6), sharex=True, sharey=True, constrained_layout=True)
-axs[0, 0].set_title('Prior Mean')
-axs[0, 0].imshow(perm_prior.mean(axis=0).T)
-axs[0, 1].set_title('"Truth"')
-axs[0, 1].imshow(perm_true.T)
+    1, 3, figsize=(8, 4), sharex=True, sharey=True, constrained_layout=True)
 
+par = {"vmin": perm_min, "vmax": perm_max, "origin": "lower"}
 
-axs[1, 0].set_title('Post Mean without localization')
-axs[1, 0].imshow(nl_perm_post.mean(axis=0).T)
-axs[1, 1].set_title('Post Mean with localization')
-axs[1, 1].imshow(wl_perm_post.mean(axis=0).T)
-axs[1, 1].contour(loc_mat.sum(axis=2).T, levels=[2.0, ], colors='w')
+axs[0].set_title("Prior Mean")
+im = axs[0].imshow(perm_prior.mean(axis=0).T, **par)
 
-for ax in axs.ravel():
+
+axs[1].set_title("Post Mean; No localization")
+axs[1].imshow(nl_perm_post.mean(axis=0).T, **par)
+
+axs[2].set_title("Post Mean: Localization")
+axs[2].imshow(wl_perm_post.mean(axis=0).T, **par)
+axs[2].contour(loc_mat.sum(axis=2).T, levels=[2.0, ], colors="w")
+
+fig.colorbar(im, ax=axs, label="Log Permeabilities (mD)",
+             orientation="horizontal")
+
+for ax in axs:
     for well in wells:
-        ax.plot(well[0], well[1], ['C3v', 'C1^'][int(well[2] == 120)])
+        ax.plot(well[0], well[1], ["C3v", "C1^"][int(well[2] == 120)])
+for ax in axs:
+    ax.set_xlabel('x-direction')
+axs[0].set_ylabel('y-direction')
 
 
 ###############################################################################
@@ -175,16 +186,20 @@ def restrict_permeability(x):
     2, 3, figsize=(8, 5), sharex=True, sharey=True, constrained_layout=True)
 fig.suptitle('Prior and posterior data')
 for i, ax in enumerate(axs[0, :]):
-    ax.plot(time, data_prior[..., i::3].T, color='.6', alpha=0.5)
-    ax.plot(time, nl_data_post[..., i::3].T, color='C0', alpha=0.5)
-    ax.plot(time, data_obs[0, i::3], 'C3o')
-    ax.set_ylabel('Pressure')
+    ax.set_title(f"Well ({ox[i]}; {oy[i]})")
+    ax.plot(time*24*60*60, data_prior[..., i::3].T, color='.6', alpha=0.5)
+    ax.plot(time*24*60*60, nl_data_post[..., i::3].T, color='C0', alpha=0.5)
+    ax.plot(time*24*60*60, data_obs[0, i::3], 'C3o')
 for i, ax in enumerate(axs[1, :]):
-    ax.plot(time, data_prior[..., i::3].T, color='.6', alpha=0.5)
-    ax.plot(time, wl_data_post[..., i::3].T, color='C0', alpha=0.5)
-    ax.plot(time, data_obs[0, i::3], 'C3o')
-    ax.set_xlabel('Time')
-    ax.set_ylabel('Pressure')
+    ax.plot(time*24*60*60, data_prior[..., i::3].T, color='.6', alpha=0.5)
+    ax.plot(time*24*60*60, wl_data_post[..., i::3].T, color='C0', alpha=0.5)
+    ax.plot(time*24*60*60, data_obs[0, i::3], 'C3o')
+    ax.set_xlabel("Time (s)")
+for i, ax in enumerate(axs[:, 0]):
+    ax.set_ylabel("Pressure (bar)")
+for i, txt in enumerate(["No l", "L"]):
+    axs[i, 2].yaxis.set_label_position("right")
+    axs[i, 2].set_ylabel(f"{txt}ocalization")
 
 
 ###############################################################################
diff --git a/requirements-dev.txt b/requirements-dev.txt
index 230467f..0b5e3aa 100644
--- a/requirements-dev.txt
+++ b/requirements-dev.txt
@@ -8,14 +8,15 @@ matplotlib
 setuptools_scm
 
 # FOR DOCUMENTATION
-sphinx
+sphinx>=7.3
 numpydoc
 sphinx_design
 sphinx_numfig
-sphinx_gallery
+sphinx_gallery>=0.16
 pydata_sphinx_theme
 sphinx_automodapi
 ipykernel
+pooch
 
 # FOR TESTING
 flake8