.. DO NOT EDIT.
.. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY.
.. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE:
.. "examples/generated/scripts_synth_data/three_survey_lines_smt_bayesbay.py"
.. LINE NUMBERS ARE GIVEN BELOW.

.. only:: html

    .. note::
        :class: sphx-glr-download-link-note

        :ref:`Go to the end <sphx_glr_download_examples_generated_scripts_synth_data_three_survey_lines_smt_bayesbay.py>`
        to download the full example code.

.. rst-class:: sphx-glr-example-title

.. _sphx_glr_examples_generated_scripts_synth_data_three_survey_lines_smt_bayesbay.py:


Invert three surveys line for a thin plate using the surrogate model
====================================================================

.. raw:: html

   <!-- Please leave the cell below as it is -->

.. GENERATED FROM PYTHON SOURCE LINES 13-18

|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/airborne_em/airborne_em_three_lines_transmitters.ipynb


.. GENERATED FROM PYTHON SOURCE LINES 21-40

.. raw:: html

   <!-- Again, please don't touch the markdown cell above. We'll generate badge 
        automatically from the above cell. -->

.. raw:: html

   <!-- This cell describes things related to environment setup, so please add more text 
        if something special (not listed below) is needed to run this notebook -->

..

   If you are running this notebook locally, make sure you’ve followed
   `steps
   here <https://github.com/inlab-geo/cofi-examples#run-the-examples-with-cofi-locally>`__
   to set up the environment. (This
   `environment.yml <https://github.com/inlab-geo/cofi-examples/blob/main/envs/environment.yml>`__
   file specifies a list of packages required to run the notebooks)


.. GENERATED FROM PYTHON SOURCE LINES 43-49

This notebook assumes that you have created a surrogate model by
executing the following two notebooks: - `Latin Hypercube
Sampling <./three_survey_lines_latin_hypercube_sampling.ipynb>`__ -
`Surrogate model
creation <./three_survey_lines_surrogate_model_creation.ipynb>`__


.. GENERATED FROM PYTHON SOURCE LINES 49-62

.. code-block:: Python


    # -------------------------------------------------------- #
    #                                                          #
    #     Uncomment below to set up environment on "colab"     #
    #                                                          #
    # -------------------------------------------------------- #

    # !pip install -U cofi
    # !pip install git+https://github.com/JuergHauser/PyP223.git
    # !pip install smt
    # !git clone https://github.com/inlab-geo/cofi-examples.git
    # %cd cofi-examples/examples/vtem_max








.. GENERATED FROM PYTHON SOURCE LINES 64-70

.. code-block:: Python


    # If this notebook is run locally PyP223 and smt need to be installed separately by uncommenting the following lines, 
    # that is by removing the # and the white space between it and the exclamation mark.
    # !pip install git+https://github.com/JuergHauser/PyP223.git
    # !pip install smt








.. GENERATED FROM PYTHON SOURCE LINES 72-95

.. code-block:: Python


    import arviz
    import bayesbay
    from vtem_max_forward_lib import (
        problem_setup, 
        system_spec, 
        survey_setup, 
        ForwardWrapper, 
        plot_predicted_profile, 
        plot_transient, 
        plot_plate_faces, 
        plot_plate_faces_single
    )
    import functools
    import matplotlib.pyplot as plt
    from matplotlib.lines import Line2D
    import pickle
    import numpy as np

    import cofi

    np.random.seed(42)








.. GENERATED FROM PYTHON SOURCE LINES 100-114

Background
----------

This example inverts three survey line of VTEM max data using the
vertical component for a thin plate target. It thus becomes possible to
invert for the easting,northing, depth of the plate reference point, the
plate dip and plate azimuth. Solving the forward problem, that is
calculating the objective function, usess the surrogate model that has
been created by the `Kriging
approach <./three_survey_lines_surrogate_model_creation.ipynb>`__
applied to the `latin hypercube
samples <three_survey_lines_latin_hypercube_sampling.ipynb>`__ of the
objective function.


.. GENERATED FROM PYTHON SOURCE LINES 117-120

Problem definition
------------------


.. GENERATED FROM PYTHON SOURCE LINES 120-134

.. code-block:: Python


    tx_min = 115
    tx_max = 281
    tx_interval = 15
    ty_min = 25
    ty_max = 176
    ty_interval = 75
    tx_points = np.arange(tx_min, tx_max, tx_interval)
    ty_points = np.arange(ty_min, ty_max, ty_interval)
    n_transmitters = len(tx_points) * len(ty_points)
    tx, ty = np.meshgrid(tx_points, ty_points)
    tx = tx.flatten()
    ty = ty.flatten()








.. GENERATED FROM PYTHON SOURCE LINES 136-143

.. code-block:: Python


    fiducial_id = np.arange(len(tx))
    line_id = np.zeros(len(tx), dtype=int)
    line_id[ty == ty_points[0]] = 0
    line_id[ty == ty_points[1]] = 1
    line_id[ty == ty_points[2]] = 2








.. GENERATED FROM PYTHON SOURCE LINES 145-162

.. code-block:: Python


    survey_setup = {
        "tx": tx,                                                   # transmitter easting/x-position
        "ty": ty,                                                   # transmitter northing/y-position
        "tz": np.array([50] * n_transmitters),                     # transmitter height/z-position
        "tazi": np.deg2rad(np.array([90] * n_transmitters)),    # transmitter azimuth
        "tincl": np.deg2rad(np.array([6] * n_transmitters)),    # transmitter inclination
        "rx": tx,                                                   # receiver easting/x-position
        "ry": np.array([100] * n_transmitters),                    # receiver northing/y-position
        "rz": np.array([50] * n_transmitters),                     # receiver height/z-position
        "trdx": np.array([0] * n_transmitters),                    # transmitter receiver separation inline
        "trdy": np.array([0] * n_transmitters),                    # transmitter receiver separation crossline
        "trdz": np.array([0] * n_transmitters),                    # transmitter receiver separation vertical
        "fiducial_id": fiducial_id,                                 # unique id for each transmitter
        "line_id": line_id                  # id for each line
    }








.. GENERATED FROM PYTHON SOURCE LINES 164-180

.. code-block:: Python


    true_model = {
        "res": np.array([300, 1000]), 
        "thk": np.array([20]), 
        "peast": np.array([175]), 
        "pnorth": np.array([100]), 
        "ptop": np.array([30]), 
        "pres": np.array([0.1]), 
        "plngth1": np.array([100]), 
        "plngth2": np.array([100]), 
        "pwdth1": np.array([0.1]), 
        "pwdth2": np.array([90]), 
        "pdzm": np.array([75]),
        "pdip": np.array([60])
    }








.. GENERATED FROM PYTHON SOURCE LINES 182-186

.. code-block:: Python


    forward = ForwardWrapper(true_model, problem_setup, system_spec, survey_setup,
                             ["pdip","pdzm", "peast", "ptop", "pwdth2"], data_returned=["vertical"])





.. rst-class:: sphx-glr-script-out

 .. code-block:: none

    ['pdip', 'pdzm', 'peast', 'ptop', 'pwdth2']




.. GENERATED FROM PYTHON SOURCE LINES 188-192

.. code-block:: Python


    # check the order of parameters in a model vector
    forward.params_to_invert





.. rst-class:: sphx-glr-script-out

 .. code-block:: none


    ['pdip', 'pdzm', 'peast', 'ptop', 'pwdth2']



.. GENERATED FROM PYTHON SOURCE LINES 194-197

.. code-block:: Python


    true_param_value = np.array([60.,65., 175., 30., 90.])








.. GENERATED FROM PYTHON SOURCE LINES 202-205

Ensemble method using the surrogate model
-----------------------------------------


.. GENERATED FROM PYTHON SOURCE LINES 205-210

.. code-block:: Python


    filename = "kriging_surrogate_model.pkl"
    with open(filename, "rb") as f:
       sm = pickle.load(f)








.. GENERATED FROM PYTHON SOURCE LINES 215-217

**Initialise a model for inversion**


.. GENERATED FROM PYTHON SOURCE LINES 217-222

.. code-block:: Python


    init_param_value = np.array([45, 90, 160, 35, 80])
    m_min = np.array([15, 35, 155, 30, 65])
    m_max = np.array([75, 145, 185, 40, 115])








.. GENERATED FROM PYTHON SOURCE LINES 227-229

**Define uniform priors for each model parameter**


.. GENERATED FROM PYTHON SOURCE LINES 229-252

.. code-block:: Python


    def initialize_param(param, position=None, value=1):
        return np.array([value]) + 0.5 * np.random.randn()

    parameters = []
    for iparam, (vmin, vmax) in enumerate(zip(m_min, m_max)):
        parameter = bayesbay.prior.UniformPrior(
            name=f"m{iparam}",
            vmin=m_min[iparam],
            vmax=m_max[iparam],
            perturb_std=(vmax - vmin) / 20,
        )
        custom_init = functools.partial(initialize_param, value=init_param_value[iparam])
        parameter.set_custom_initialize(custom_init)
        parameters.append(parameter)

    param_space = bayesbay.parameterization.ParameterSpace(
        name="param_space",
        n_dimensions=1,
        parameters=parameters,
    )
    parameterization = bayesbay.parameterization.Parameterization(param_space)








.. GENERATED FROM PYTHON SOURCE LINES 257-259

**Define the log likelihood**


.. GENERATED FROM PYTHON SOURCE LINES 259-275

.. code-block:: Python


    def my_objective(model):
        val = sm.predict_values(model)[0][0]
        if val < 1e-3:
            return 1e-3
        else:
            return val
        
    def my_log_likelihood(state, *args, **kwargs):
        model = np.array(
            [state["param_space"][f"m{i}"] for i in range(init_param_value.size)]
        )
        return -0.5 * my_objective(model.T)

    log_likelihood = bayesbay.likelihood.LogLikelihood(log_like_func=my_log_likelihood)








.. GENERATED FROM PYTHON SOURCE LINES 280-282

**Initialize the Markov chains**


.. GENERATED FROM PYTHON SOURCE LINES 282-291

.. code-block:: Python


    n_chains = 12
    walkers_start = []
    for i in range(n_chains):
        walkers_start.append(
            parameterization.initialize()
        )  # A bayesbay.State is appended to walkers_start for each chain









.. GENERATED FROM PYTHON SOURCE LINES 296-298

**Define CoFI options**


.. GENERATED FROM PYTHON SOURCE LINES 298-312

.. code-block:: Python


    inv_options = cofi.InversionOptions()
    inv_options.set_tool("bayesbay")
    inv_options.set_params(
        walkers_starting_states=walkers_start,
        perturbation_funcs=parameterization.perturbation_funcs,
        log_like_ratio_func=log_likelihood,
        n_chains=n_chains,
        n_iterations=5_000,
        burnin_iterations=500,
        verbose=False,
        save_every=25,
    )








.. GENERATED FROM PYTHON SOURCE LINES 317-319

**Define CoFI inverse problem and run**


.. GENERATED FROM PYTHON SOURCE LINES 319-323

.. code-block:: Python


    inv = cofi.Inversion(cofi.BaseProblem(), inv_options)
    inv_result = inv.run()








.. GENERATED FROM PYTHON SOURCE LINES 328-331

Plotting
--------


.. GENERATED FROM PYTHON SOURCE LINES 331-416

.. code-block:: Python


    import arviz_base
    from scipy.stats import gaussian_kde
    arviz.style.use("arviz-variat")
    var_names = [
        "Dip (°)",
        "Dip Azimuth (°)",
        "Easting (m)",
        "Depth (m)",
        "Width (m)",
    ]
    clean_names = ["Dip_deg", "Dip_Azimuth_deg", "Easting_m", "Depth_m", "Width_m"]

    results = inv_result.models
    # Keep raw 1D samples for plate visualization in next cell
    posterior_samples = {
        clean_names[i]: np.concatenate(results[f"param_space.m{i}"])
        for i in range(init_param_value.size)
    }

    true_values = {
        var_names[i]: true_param_value[i] for i in range(init_param_value.size)
    }

    # Set to True for KDE contours (old style), False for scatter (arviz 1.0 default)
    USE_KDE_CONTOURS = True

    # Convert to DataTree with fake chain for plot_pair
    posterior_for_arviz = {
        k: v[np.newaxis, :] for k, v in posterior_samples.items()
    }
    arviz_base.rcParams["plot.max_subplots"] = 80
    az_idata = arviz_base.from_dict({"posterior": posterior_for_arviz})

    if USE_KDE_CONTOURS:
        pm = arviz.plot_pair(
            az_idata,
            marginal=True,
            triangle="lower",
            visuals={"scatter": False},
        )
        # Add KDE contours to off-diagonal panels
        n = len(clean_names)
        for i in range(n):
            for j in range(n):
                if i <= j:
                    continue
                try:
                    ax = pm.iget_target(i, j)
                except (ValueError, IndexError):
                    continue
                x = posterior_samples[clean_names[j]]
                y = posterior_samples[clean_names[i]]
                kde = gaussian_kde(np.vstack([x, y]))
                xmin, xmax = x.min(), x.max()
                ymin, ymax = y.min(), y.max()
                xx, yy = np.mgrid[xmin:xmax:100j, ymin:ymax:100j]
                zz = kde(np.vstack([xx.ravel(), yy.ravel()])).reshape(xx.shape)
                ax.contourf(xx, yy, zz, levels=10, cmap="Blues")
                ax.contour(xx, yy, zz, levels=10, colors="grey", linewidths=0.5, alpha=0.5)
    else:
        pm = arviz.plot_pair(
            az_idata,
            marginal=True,
            triangle="lower",
        )

    # Add reference values
    ref_vals = list(true_values.values())
    n = len(ref_vals)
    for i in range(n):
        for j in range(n):
            try:
                ax = pm.iget_target(i, j)
            except (ValueError, IndexError):
                continue
            if i == j:
                ax.axvline(ref_vals[i], color="green", linestyle="--", lw=1, alpha=0.5)
            elif i > j:
                ax.plot(
                    ref_vals[j], ref_vals[i], "o",
                    color="yellow", markeredgecolor="k", ms=10, zorder=5,
                )





.. image-sg:: /examples/generated/scripts_synth_data/images/sphx_glr_three_survey_lines_smt_bayesbay_001.png
   :alt: three survey lines smt bayesbay
   :srcset: /examples/generated/scripts_synth_data/images/sphx_glr_three_survey_lines_smt_bayesbay_001.png
   :class: sphx-glr-single-img





.. GENERATED FROM PYTHON SOURCE LINES 418-494

.. code-block:: Python


    _, axes = plt.subplots(2, 2)
    axes[1, 1].axis("off")
    plot_plate_faces(
        "plate_true",
        forward,
        true_param_value,
        axes[0, 0],
        axes[0, 1],
        axes[1, 0],
        color="purple",
        label="True model",
    )
    plot_plate_faces(
        "plate_init",
        forward,
        init_param_value,
        axes[0, 0],
        axes[0, 1],
        axes[1, 0],
        color="green",
        label="Starting model",
    )

    plt.tight_layout()

    idraw = np.random.randint(0, len(posterior_samples[clean_names[0]]))
    sample = np.array([posterior_samples[name][idraw] for name in clean_names])

    plot_plate_faces(
        "plate_inverted",
        forward,
        sample,
        axes[0, 0],
        axes[0, 1],
        axes[1, 0],
        color="red",
        label="Posterior sample",
        linestyle="dotted",
    )

    point = Line2D(
        [0],
        [0],
        label="Fiducial",
        marker="o",
        markersize=5,
        markeredgecolor="orange",
        markerfacecolor="orange",
        linestyle="",
    )

    handles, labels = axes[1, 0].get_legend_handles_labels()
    handles.extend([point])

    axes[1, 0].legend(handles=handles, bbox_to_anchor=(1.04, 0), loc="lower left")

    # plot 10 randomly selected samples of the posterior distribution
    idraws = np.random.choice(
        np.arange(0, len(posterior_samples[clean_names[0]])), 10, replace=False
    )
    for idraw in idraws:
        sample = np.array([posterior_samples[name][idraw] for name in clean_names])
        plot_plate_faces(
            "plate_inverted",
            forward,
            sample,
            axes[0, 0],
            axes[0, 1],
            axes[1, 0],
            color="red",
            label="Posterior sample",
            linestyle="dotted",
        )





.. image-sg:: /examples/generated/scripts_synth_data/images/sphx_glr_three_survey_lines_smt_bayesbay_002.png
   :alt: three survey lines smt bayesbay
   :srcset: /examples/generated/scripts_synth_data/images/sphx_glr_three_survey_lines_smt_bayesbay_002.png
   :class: sphx-glr-single-img





.. GENERATED FROM PYTHON SOURCE LINES 499-512

--------------

Watermark
---------

.. raw:: html

   <!-- Feel free to add more modules in the watermark_list below, if more packages are used -->

.. raw:: html

   <!-- Otherwise please leave the below code cell unchanged -->


.. GENERATED FROM PYTHON SOURCE LINES 512-518

.. code-block:: Python


    watermark_list = ["bayesbay", "cofi", "numpy", "scipy", "matplotlib","smt"]
    for pkg in watermark_list:
        pkg_var = __import__(pkg)
        print(pkg, getattr(pkg_var, "__version__"))





.. rst-class:: sphx-glr-script-out

 .. code-block:: none

    bayesbay 0.3.10
    cofi 0.2.11+71.gb28b5b0
    numpy 2.2.6
    scipy 1.17.1
    matplotlib 3.10.9
    smt 2.13.0




.. GENERATED FROM PYTHON SOURCE LINES 519-519

sphinx_gallery_thumbnail_number = -1


.. rst-class:: sphx-glr-timing

   **Total running time of the script:** (0 minutes 9.964 seconds)


.. _sphx_glr_download_examples_generated_scripts_synth_data_three_survey_lines_smt_bayesbay.py:

.. only:: html

  .. container:: sphx-glr-footer sphx-glr-footer-example

    .. container:: sphx-glr-download sphx-glr-download-jupyter

      :download:`Download Jupyter notebook: three_survey_lines_smt_bayesbay.ipynb <three_survey_lines_smt_bayesbay.ipynb>`

    .. container:: sphx-glr-download sphx-glr-download-python

      :download:`Download Python source code: three_survey_lines_smt_bayesbay.py <three_survey_lines_smt_bayesbay.py>`

    .. container:: sphx-glr-download sphx-glr-download-zip

      :download:`Download zipped: three_survey_lines_smt_bayesbay.zip <three_survey_lines_smt_bayesbay.zip>`


.. only:: html

 .. rst-class:: sphx-glr-signature

    `Gallery generated by Sphinx-Gallery <https://sphinx-gallery.github.io>`_