.. DO NOT EDIT. .. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY. .. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE: .. "examples/generated/scripts_synth_data/linear_regression_neighpy.py" .. LINE NUMBERS ARE GIVEN BELOW. .. only:: html .. note:: :class: sphx-glr-download-link-note :ref:`Go to the end ` to download the full example code. .. rst-class:: sphx-glr-example-title .. _sphx_glr_examples_generated_scripts_synth_data_linear_regression_neighpy.py: Polynomial Linear Regression with the Neighbourhood Algorithm ============================================================= .. GENERATED FROM PYTHON SOURCE LINES 9-14 |Open In Colab| .. |Open In Colab| image:: https://img.shields.io/badge/open%20in-Colab-b5e2fa?logo=googlecolab&style=flat-square&color=ffd670 :target: https://colab.research.google.com/github/inlab-geo/cofi-examples/blob/main/examples/linear_regression/linear_regression.ipynb .. GENERATED FROM PYTHON SOURCE LINES 17-24 If you are running this notebook locally, make sure you’ve followed `steps here `__ to set up the environment. (This `environment.yml `__ file specifies a list of packages required to run the notebooks) .. GENERATED FROM PYTHON SOURCE LINES 27-29 -------------- .. GENERATED FROM PYTHON SOURCE LINES 32-128 1. Introduction ---------------- 1.1 Neighbourhood Algorithm ~~~~~~~~~~~~~~~~~~~~~~~~~~~~ The Neighbourhood Algorithm (NA) is described in two papers (`Sambridge, 1999a <>`__ and `Sambridge, 1999b <>`__) and became popular for geophysical inverse problems. We give a breif summary of the NA here for intuition. It can be divided into two main phases: 1. Direct Search Phase 2. Appraisal Phase As decribed in the papers, these phases use geometrical structures called Voronoi cells to define neighbourhoods of parameter space within which the value of the objective function is constant. These neighbourhoods are then repeatedly sampled and refined to find the optimum position in the parameter space. Direct Search Phase ^^^^^^^^^^^^^^^^^^^ This is a derivative-free optimisation aglorithm that results in an ensemble of points/samples/neighbourhoods distributed according to the objective function being solved. More dense regions of points correspond to the more optimum areas of the objective function. In brief, the phase algorithm progresses as follows: 1. Create an initial ensemble (of size :math:`n_i`) of points by uniformly sampling the parameter space 2. Interate :math:`n` times: 1. Rank the ensemble according to the objective function 2. Select the best :math:`n_r` points in the ensemble to be resampled 3. Resample the neighbourhoods/Voronoi cells of best :math:`n_r` points to obtain :math:`n_s` new samples 4. Add the new samples to the ensemble In the original papers the resampling is performed using a Gibbs sampler, although in theory any sampling algorithm can be used within each neighbourhood - the trick is in computing the neighbourhoods. The result of this phase is a large ensemble of size :math:`N = n_i + n \times n_s` and their associated objective values, effectivly a discretisation of the objective function. Appraisal Phase ^^^^^^^^^^^^^^^ This phase is used to reffine an ensemble of points, and produce new samples that are distributed according to the posterior distribution function. Effectively, this is just a Bayesian sampling of a step-wise constant posterior. In theory any sampling approach can be used to create the ensemble to be refined, and any sampling approach can be used to do the refining. In the original papers the initial ensemble is created using the Direct Search Phase, and the refining is done with a Gibbs sampler. A key advantage of this phase is that it avoids computing the forward function for all the new samples. So if you have already partially sampled the parameter space but the forward function is too slow that your initial method will take too long to get a sufficient number of samples, you can use the appraisal phase of the NA. 1.2 CoFI and the NA ~~~~~~~~~~~~~~~~~~~~ The implementation of the NA that ``cofi`` wraps is called ```neighpy`` <>`__. This implementation implements both phases of the NA as describled in the original papers. Because of the multi-phase nature of the NA, ``cofi`` gives you 3 options: 1. Only run the Direct Search Phase 2. Only run the Appraisal Phase, using samples obtained from any method 3. Run both phases, using the Direct Seach Phase samples as the initial ensemble for the Appraisal Phase We will look at all of these options in this tutorial. 1.3 Problem ~~~~~~~~~~~~ This tutorial focusses on linear regression - that is, fitting curves to datasets. We will work with polynomial curves, .. math:: y(x) = \sum_{n=0}^N m_n x^n\,. Here, :math:`N` is the ‘order’ of the polynomial: if N=1 we have a straight line, if N=2 it will be a quadratic, and so on. The :math:`m_n` are the ‘model coefficients’. We have a set of noisy data values, Y, measured at known locations, X, and wish to find the best fit degree 3 polynomial. The function we are going to fit is: :math:`y=-6-5x+2x^2+x^3` .. GENERATED FROM PYTHON SOURCE LINES 131-134 1. Introduction ---------------- .. GENERATED FROM PYTHON SOURCE LINES 134-143 .. code-block:: Python # display theory on travel time tomography from IPython.display import display, Markdown with open("../../theory/neighbourhood_algorithm.md", "r") as f: content = f.read() display(Markdown(content)) .. rst-class:: sphx-glr-script-out .. code-block:: none .. GENERATED FROM PYTHON SOURCE LINES 148-182 1.2 CoFI and the NA ~~~~~~~~~~~~~~~~~~~~ The implementation of the NA that ``cofi`` wraps is called ```neighpy`` <>`__. This implementation implements both phases of the NA as describled in the original papers. Because of the multi-phase nature of the NA, ``cofi`` gives you 3 options: 1. Only run the Direct Search Phase 2. Only run the Appraisal Phase, using samples obtained from any method 3. Run both phases, using the Direct Seach Phase samples as the initial ensemble for the Appraisal Phase We will look at all of these options in this tutorial. 1.3 Problem ~~~~~~~~~~~~ This tutorial focusses on linear regression - that is, fitting curves to datasets. We will work with polynomial curves, .. math:: y(x) = \sum_{n=0}^N m_n x^n\,. Here, :math:`N` is the ‘order’ of the polynomial: if N=1 we have a straight line, if N=2 it will be a quadratic, and so on. The :math:`m_n` are the ‘model coefficients’. We have a set of noisy data values, Y, measured at known locations, X, and wish to find the best fit degree 3 polynomial. The function we are going to fit is: :math:`y=-6-5x+2x^2+x^3` .. GENERATED FROM PYTHON SOURCE LINES 185-187 -------------- .. GENERATED FROM PYTHON SOURCE LINES 190-193 2. Import Modules and Plot Functions ------------------------------------ .. GENERATED FROM PYTHON SOURCE LINES 193-202 .. code-block:: Python # -------------------------------------------------------- # # # # Uncomment below to set up environment on "colab" # # # # -------------------------------------------------------- # # !pip install -U cofi .. GENERATED FROM PYTHON SOURCE LINES 204-215 .. code-block:: Python import numpy as np import matplotlib.pyplot as plt from scipy.spatial import Voronoi, voronoi_plot_2d from functools import partial from numbers import Number from cofi import BaseProblem, InversionOptions, Inversion np.random.seed(42) .. GENERATED FROM PYTHON SOURCE LINES 217-406 .. code-block:: Python ############## PLOTTING ############################################################### def _add_extra(plot_fn, extra): # plot_fn = ax.plot, ax.scatter .... # extra = [x, y, kwargs] try: kwargs = extra[2] except IndexError: kwargs = {} plot_fn(extra[0], extra[1], **kwargs) def plot_data_space(basis_func, _m_true, x, y_observed, sigma, samples=None, samples_label="", extras=None): # extras = [[model, kwargs], ...] _x_plot = np.linspace(-3.5, 2.5) _G_plot = basis_func(_x_plot) fig, ax = plt.subplots(figsize=(12, 5)) if samples is not None: ax.plot( _x_plot, np.apply_along_axis(lambda m: _G_plot @ m, 1, samples).T, color="k", alpha=0.1, label=samples_label, ) if extras is not None: for extra in extras: _y_plot = _G_plot @ extra[0] _add_extra(ax.plot, [_x_plot, _y_plot, extra[1]]) # [x, y, kwargs] _y_plot = _G_plot @ _m_true ax.plot(_x_plot, _y_plot, color="r", label="True model") ax.errorbar( x, y_observed, yerr=sigma, fmt=".", color="lightcoral", label="observed data" ) ax.set_xlabel("X") ax.set_ylabel("Y") ax.set_ylim(-11, 10) handles, labels = [x[-5:] for x in ax.get_legend_handles_labels()] by_label = dict(zip(labels, handles)) ax.legend(by_label.values(), by_label.keys()) def plot_direct_search_evolution(objectives, n_initial, n_total): fig, ax = plt.subplots(figsize=(12, 5)) ax.plot(objectives, marker=".", linestyle="", markersize=2, color="black") ax.set_yscale("log") ax.set_xlabel("Sample number") ax.set_ylabel("Objective value") ax.set_title("Evolution of the objective function during the direct search") ax.plot( [0, n_initial], [1e1, 1e1], linestyle="--", color="gray", marker="|", markersize=10, ) ax.plot( [n_initial, n_total], [1e1, 1e1], linestyle="--", color="gray", marker="|", markersize=10, ) ax.text(n_initial // 2, 1e1, "Initial random search", ha="center", va="bottom") ax.text( n_initial + (n_total - n_initial) // 2, 1e1, "Optimisation search", ha="center", va="bottom", ) def plot_direct_search_marginals(_m_true, bounds, ds_voronoi, extras=[]): # extras = [[model, kwargs], ...] fig, axs = plt.subplots( 3, 3, figsize=(7, 7), tight_layout=True, gridspec_kw={"hspace": 0, "wspace": 0} ) for i in range(1, 4): for j in range(3): if j < i: vor = Voronoi(ds_voronoi[:, [i, j]]) voronoi_plot_2d( vor, ax=axs[i - 1, j], show_vertices=False, show_points=False, line_width=0.1, ) axs[i - 1, j].scatter( _m_true[i], _m_true[j], c="r", marker="x", s=100, label="True model", zorder=10, ) for extra in extras: m, kwargs = extra _add_extra(axs[i - 1, j].scatter, [m[i], m[j], kwargs]) # [x, y, kwargs] axs[i - 1, j].set_xlim(bounds[i]) axs[i - 1, j].set_ylim(bounds[j]) if j == 0: axs[i - 1, j].set_ylabel(f"$m_{i}$") else: axs[i - 1, j].set_yticks([]) if i == 3: axs[i - 1, j].set_xlabel(f"$m_{j}$") else: axs[i - 1, j].set_xticks([]) else: axs[i - 1, j].set_visible(False) handles, labels = axs[1, 0].get_legend_handles_labels() by_label = dict(zip(labels, handles)) fig.legend( by_label.values(), by_label.keys(), loc="lower left", bbox_to_anchor=(0.4, 0.65) ) fig.suptitle("Direct Search Phase Results") def plot_appraisal_marginals(_m_true, bounds, appraised, extras=[]): appraised_mean = appraised.mean(axis=0) fig, axs = plt.subplots( 4, 4, figsize=(7, 7), tight_layout=True, gridspec_kw={"hspace": 0, "wspace": 0} ) for i in range(4): for j in range(4): if j == i: axs[i, j].hist(appraised[:, i], bins=20, histtype="step", color="k", label="Appraisal samples") axs[i, j].axvline(appraised_mean[i], color="b", linestyle="--", label="Appraisal mean") axs[i, j].axvline(_m_true[i], color="r", linestyle="--", label="True model") for extra in extras: m, _kwargs = extra kwargs = {"linestyle": "--"} if "color" in _kwargs.keys(): kwargs = {**kwargs, "color":_kwargs["color"]} axs[i, j].axvline(m[i], **kwargs) axs[i, j].set_xlim(bounds[i]) axs[i, j].set_yticks([]) elif j < i: axs[i, j].scatter( appraised[:, i], appraised[:, j], c="gray", s=0.1, label="Appraisal samples", ) axs[i, j].scatter( appraised_mean[i], appraised_mean[j], c="b", marker="x", s=100, label="Appraisal mean", zorder=10, ) axs[i, j].scatter( _m_true[i], _m_true[j], c="r", marker="x", s=100, label="True model", zorder=10, ) for extra in extras: m, kwargs = extra _add_extra(axs[i, j].scatter, [m[i], m[j], kwargs]) # [x, y, kwargs] axs[i, j].set_xlim(bounds[i]) axs[i, j].set_ylim(bounds[j]) else: axs[i, j].set_visible(False) if j == 0: if i != 0: axs[i, j].set_ylabel(f"$m_{i}$") else: axs[i, j].set_yticks([]) if i == 3: axs[i, j].set_xlabel(f"$m_{j}$") else: axs[i, j].set_xticks([]) handles, labels = axs[1, 0].get_legend_handles_labels() by_label = dict(zip(labels, handles)) fig.legend( by_label.values(), by_label.keys(), loc="lower left", bbox_to_anchor=(0.535, 0.51) ) fig.suptitle("Appraisal Phase Results") .. GENERATED FROM PYTHON SOURCE LINES 411-413 -------------- .. GENERATED FROM PYTHON SOURCE LINES 416-454 3. Define the problem ---------------------- Here we compute :math:`y(x)` for multiple :math:`x`-values simultaneously, so write the forward operator in the following form: .. math:: \left(\begin{array}{c}y_1\\y_2\\\vdots\\y_N\end{array}\right) = \left(\begin{array}{ccc}1&x_1&x_1^2&x_1^3\\1&x_2&x_2^2&x_2^3\\\vdots&\vdots&\vdots\\1&x_N&x_N^2&x_N^3\end{array}\right)\left(\begin{array}{c}m_0\\m_1\\m_2\end{array}\right) This clearly has the required general form, :math:`\mathbf{y=Gm}`, and so the best-fitting model can be identified using the least-squares algorithm. In the following code block, we’ll define the forward function and generate some random data points as our dataset. .. math:: \begin{align} \text{forward}(\textbf{m}) &= \textbf{G}\textbf{m}\\ &= \text{basis\_func}(\textbf{x})\cdot\textbf{m} \end{align} where: - :math:`\text{forward}` is the forward function that takes in a model and produces synthetic data, - :math:`\textbf{m}` is the model vector, - :math:`\textbf{G}` is the basis matrix (i.e. design matrix) of this linear regression problem and looks like the following: .. math:: \left(\begin{array}{ccc}1&x_1&x_1^2&x_1^3\\1&x_2&x_2^2&x_2^3\\\vdots&\vdots&\vdots\\1&x_N&x_N^2&x_N^3\end{array}\right) - :math:`\text{basis\_func}` is the basis function that converts :math:`\textbf{x}` into :math:`\textbf{G}` Recall that the function we are going to fit is: :math:`y=-6-5x+2x^2+x^3` .. GENERATED FROM PYTHON SOURCE LINES 454-473 .. code-block:: Python # generate data with random Gaussian noise def basis_func(x): return np.array([x**i for i in range(4)]).T # x -> G _m_true = np.array([-6, -5, 2, 1]) # m sample_size = 20 # N x = np.random.choice(np.linspace(-3.5, 2.5), size=sample_size) # x def forward_func(m): return basis_func(x) @ m # m -> y_synthetic sigma = 1 Cd = np.eye(sample_size) * sigma**2 # Cd y_observed = forward_func(_m_true) + np.random.normal(0, sigma, sample_size) # d .. GENERATED FROM PYTHON SOURCE LINES 475-479 .. code-block:: Python # quickly redefining some plot functions plot_data_space = partial(plot_data_space, basis_func=basis_func, _m_true=_m_true, x=x, y_observed=y_observed, sigma=sigma) .. GENERATED FROM PYTHON SOURCE LINES 481-484 .. code-block:: Python plot_data_space() .. image-sg:: /examples/generated/scripts_synth_data/images/sphx_glr_linear_regression_neighpy_001.png :alt: linear regression neighpy :srcset: /examples/generated/scripts_synth_data/images/sphx_glr_linear_regression_neighpy_001.png :class: sphx-glr-single-img .. GENERATED FROM PYTHON SOURCE LINES 489-497 We will use a simple weighted :math:`\ell_2`-norm misfit as our objective function .. math:: F(\mathbf{m}) = (\mathbf{d} - \mathbf{Gm})^T\mathbf{C_d}^{-1}(\mathbf{d} - \mathbf{Gm}) where :math:`\mathbf{d}` is the observed data vector and :math:`\mathbf{C_d}` is the data covariance matrix. .. GENERATED FROM PYTHON SOURCE LINES 497-504 .. code-block:: Python Cd_inv = np.linalg.inv(Cd) def objective(m: np.ndarray) -> Number: y_predicted = forward_func(m) return (y_observed - y_predicted).T @ Cd_inv @ (y_observed - y_predicted) .. GENERATED FROM PYTHON SOURCE LINES 506-515 .. code-block:: Python # define the problem in cofi inv_problem = BaseProblem() inv_problem.name = "Polynomial Regression" inv_problem.set_data(y_observed) inv_problem.set_objective(objective) inv_problem.summary() .. rst-class:: sphx-glr-script-out .. code-block:: none ===================================================================== Summary for inversion problem: Polynomial Regression ===================================================================== Model shape: Unknown --------------------------------------------------------------------- List of functions/properties set by you: ['objective', 'data'] --------------------------------------------------------------------- List of functions/properties created based on what you have provided: -- none -- --------------------------------------------------------------------- List of functions/properties that can be further set for the problem: ( not all of these may be relevant to your inversion workflow ) ['log_posterior', 'log_posterior_with_blobs', 'log_likelihood', 'log_prior', 'gradient', 'hessian', 'hessian_times_vector', 'residual', 'jacobian', 'jacobian_times_vector', 'data_misfit', 'regularization', 'regularization_matrix', 'forward', 'data_covariance', 'data_covariance_inv', 'initial_model', 'model_shape', 'blobs_dtype', 'bounds', 'constraints'] .. GENERATED FROM PYTHON SOURCE LINES 520-522 -------------- .. GENERATED FROM PYTHON SOURCE LINES 525-528 4. Straight Up NA ------------------ .. GENERATED FROM PYTHON SOURCE LINES 531-535 Here we’ll just run both phases of the NA in one go. For this we use the ``cofi`` tool called ``Neighpy``. .. GENERATED FROM PYTHON SOURCE LINES 535-556 .. code-block:: Python bounds=[[-10, 10]] * 4 # some loose prior bounds on all the polynomial coefficients n_samples_per_iteration=100 # number of samples to draw per iteration n_cells_to_resample=20 # number of cells to resample - the larger this is the more the algorithm will explore n_initial_samples=5000 # number of initial samples to draw n_iterations=100 # number of iterations (direct search) to run n_resample=1000 # number of resamples to draw from appraisal n_walkers=5 # number of walkers for the appraisal inv_options = InversionOptions() inv_options.set_tool("neighpy") inv_options.set_params( bounds=bounds, n_samples_per_iteration=n_samples_per_iteration, n_cells_to_resample=n_cells_to_resample, n_initial_samples=n_initial_samples, n_iterations=n_iterations, n_resample=n_resample, n_walkers=n_walkers, ) .. GENERATED FROM PYTHON SOURCE LINES 558-563 .. code-block:: Python # quickly redefining some plot functions plot_direct_search_marginals = partial(plot_direct_search_marginals, _m_true=_m_true, bounds=bounds) plot_appraisal_marginals = partial(plot_appraisal_marginals, _m_true=_m_true, bounds=bounds) .. GENERATED FROM PYTHON SOURCE LINES 565-570 .. code-block:: Python inv = Inversion(inv_problem, inv_options) inv_result = inv.run() inv_result.summary() .. rst-class:: sphx-glr-script-out .. code-block:: none NAI - Initial Random Search NAI - Optimisation Loop: 0%| | 0/100 [00:00 u: return -np.inf # model lies outside bounds -> return log(0) return 0.0 # model lies within bounds -> return log(1) nwalkers = 8 ndim = 4 nsteps = 10000 walkers_start = np.array([0.,0.,0.,0.]) + np.random.randn(nwalkers, ndim) inv_problem.set_log_prior(log_prior) inv_problem.set_log_likelihood(log_likelihood) inv_problem.set_model_shape(ndim) inv_options = InversionOptions() inv_options.set_tool("emcee") inv_options.set_params(nwalkers=nwalkers, nsteps=nsteps, initial_state=walkers_start) inv = Inversion(inv_problem, inv_options) inv_result = inv.run() emcee_samples = inv_result.sampler.get_chain(flat=True) emcee_logppd = inv_result.sampler.get_log_prob(flat=True) .. GENERATED FROM PYTHON SOURCE LINES 800-819 .. code-block:: Python # Appraise the emcee ensemble inv_options = InversionOptions() inv_options.set_tool("neighpyII") inv_options.set_params( bounds=bounds, initial_ensemble=emcee_samples, log_ppd=emcee_logppd, n_resample=n_resample, n_walkers=n_walkers, ) inv = Inversion(inv_problem, inv_options) inv_result = inv.run() inv_result.summary() appraised = inv_result.new_samples appraised_mean = inv_result.new_samples.mean(axis=0) .. rst-class:: sphx-glr-script-out .. code-block:: none ============================ Summary for inversion result ============================ SUCCESS ---------------------------- new_samples: [[-5.49030657 -5.13867999 1.84577882 1.10040005] [-6.32820487 -9.80679864 -1.54231433 -8.96284046] [-8.75845898 -9.14943555 -4.92184315 2.02114239] ... [-4.35883717 -9.47734368 -0.28714092 -2.0451282 ] [-0.81304734 -7.06419206 0.75777164 0.52980791] [-7.01403911 -7.81563622 -0.5894208 8.33705105]] .. GENERATED FROM PYTHON SOURCE LINES 821-824 .. code-block:: Python plot_appraisal_marginals(appraised=appraised, extras=[]) .. image-sg:: /examples/generated/scripts_synth_data/images/sphx_glr_linear_regression_neighpy_009.png :alt: Appraisal Phase Results :srcset: /examples/generated/scripts_synth_data/images/sphx_glr_linear_regression_neighpy_009.png :class: sphx-glr-single-img .. GENERATED FROM PYTHON SOURCE LINES 826-835 .. code-block:: Python plot_data_space( samples=appraised, samples_label="Appraisal samples", extras=[ [appraised_mean, {"label": "Appraisal mean", "color": "b"}], ], ) .. image-sg:: /examples/generated/scripts_synth_data/images/sphx_glr_linear_regression_neighpy_010.png :alt: linear regression neighpy :srcset: /examples/generated/scripts_synth_data/images/sphx_glr_linear_regression_neighpy_010.png :class: sphx-glr-single-img .. GENERATED FROM PYTHON SOURCE LINES 840-842 -------------- .. GENERATED FROM PYTHON SOURCE LINES 845-848 Watermark --------- .. GENERATED FROM PYTHON SOURCE LINES 848-854 .. code-block:: Python watermark_list = ["cofi", "numpy", "scipy", "matplotlib", "emcee", "arviz"] for pkg in watermark_list: pkg_var = __import__(pkg) print(pkg, getattr(pkg_var, "__version__")) .. rst-class:: sphx-glr-script-out .. code-block:: none cofi 0.2.11+71.gb28b5b0 numpy 2.2.6 scipy 1.17.1 matplotlib 3.10.9 emcee 3.1.6 arviz 1.1.0 .. GENERATED FROM PYTHON SOURCE LINES 860-860 sphinx_gallery_thumbnail_number = -1 .. rst-class:: sphx-glr-timing **Total running time of the script:** (1 minutes 9.318 seconds) .. _sphx_glr_download_examples_generated_scripts_synth_data_linear_regression_neighpy.py: .. only:: html .. container:: sphx-glr-footer sphx-glr-footer-example .. container:: sphx-glr-download sphx-glr-download-jupyter :download:`Download Jupyter notebook: linear_regression_neighpy.ipynb ` .. container:: sphx-glr-download sphx-glr-download-python :download:`Download Python source code: linear_regression_neighpy.py ` .. container:: sphx-glr-download sphx-glr-download-zip :download:`Download zipped: linear_regression_neighpy.zip ` .. only:: html .. rst-class:: sphx-glr-signature `Gallery generated by Sphinx-Gallery `_