1-century-dcip-inversion.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Inversion of DC & IP data at the Century Deposit\n",
    "\n",
    "The Century Deposit is a Zinc-lead-silver deposit is located 250 km to the NNW of the Mt Isa region in Queensland Australia (Location: 18° 43' 15\"S, 138° 35' 54\"E). \n",
    "\n",
    "<img src=\"./images/Mutton-Figure1-1.png\" width=45%>\n",
    "\n",
    "\n",
    "**References**\n",
    "- [Mutton (2000)](https://doi.org/10.1190/1.1444878)\n",
    "\n",
    "- [geologic background](http://www.portergeo.com.au/database/mineinfo.asp?mineid=mn075)\n",
    "\n",
    "For background on DC Resistivity Surveys, see [GPG: DC Resistivity](https://gpg.geosci.xyz/content/DC_resistivity/index.html)\n",
    "\n",
    "In this notebook, we are going to invert the data and see if we can replicate the results presented in [Mutton, 2000](https://doi.org/10.1190/1.1444878). Specifically, we will work with data collected along the line 46800mE. \n",
    "\n",
    "\n",
    "## Motivation: replicate inversion results from 20 years ago!\n",
    "\n",
    "<img src=\"images/Mutton-Figure2-1.png\" width=70%>\n",
    "\n",
    "The image below shows the Induced Polarization inversion result generated using the UBC 2D DCIP inversion code that was published in [Mutton, 2000](https://doi.org/10.1190/1.1444878). \n",
    "\n",
    "**Can we replicate it with SimPEG?** \n",
    "\n",
    "(and update the colorbar 😉)\n",
    "\n",
    "<img src=\"images/Mutton-Figure5-1.png\" width=80%>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 0: Import Dependencies\n",
    "\n",
    "Here are links to the documentation for each of the packages we rely on for this tutorial. \n",
    "\n",
    "**Core Python / Jupyter**\n",
    "- [numpy](https://numpy.org/doc/1.18/): arrays\n",
    "- [matplotlib](https://matplotlib.org/3.2.1/contents.html): plotting\n",
    "- [ipywidgets](https://ipywidgets.readthedocs.io/en/latest/index.html): interactive widgets in jupyter\n",
    "\n",
    "**SimPEG ecosystem**\n",
    "- [discretize](http://discretize.simpeg.xyz/): meshes and differential operators \n",
    "- [pymatsolver](http://pymatsolver.rtfd.io/): wrappers on linear solvers\n",
    "- [SimPEG](http://docs.simpeg.xyz/): Simulation and Parameter Estimation in Geophysics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# core python \n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.colors import LogNorm\n",
    "import ipywidgets\n",
    "import os\n",
    "\n",
    "# tools in the simPEG Ecosystem \n",
    "import discretize  # for creating computational meshes\n",
    "\n",
    "# linear solvers\n",
    "try: \n",
    "    from pymatsolver import Pardiso as Solver  # this is a fast linear solver \n",
    "except ImportError:\n",
    "    from SimPEG import SolverLU as Solver  # this will be slower\n",
    "\n",
    "# SimPEG inversion machinery\n",
    "from SimPEG import (\n",
    "    Data, maps,\n",
    "    data_misfit, regularization, optimization, inverse_problem, \n",
    "    inversion, directives\n",
    ") \n",
    "\n",
    "# DC resistivity and IP modules\n",
    "from SimPEG.electromagnetics import resistivity as dc\n",
    "from SimPEG.electromagnetics import induced_polarization as ip"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# set the font size in the plots\n",
    "from matplotlib import rcParams\n",
    "rcParams[\"font.size\"] = 14"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 1: Load the data\n",
    "\n",
    "\n",
    "**What are our data?**\n",
    "\n",
    "The survey collected at the Century Deposit uses a dipole-dipole geometry with:\n",
    "- Current Electrodes (A, B): we refer to the source electrodes as \"A\" for the positive electrode and \"B\" for the negative electrode. \n",
    "- Potential Electrodes (M, N): we refer to the positive potential electrode as \"M\" and the negative as \"N\".\n",
    "\n",
    "A datum is the potential difference between the M and N electrodes. \n",
    "\n",
    "<img src=\"https://gpg.geosci.xyz/_images/add1.gif\">\n",
    "\n",
    "<img src=\"https://gpg.geosci.xyz/_images/add2.gif\">"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The data are stored in the `century` directory. Here, we take a look at what is in that directory. Each folder contains data associated with the named line. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "os.listdir('century')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will focus on the `46800E` line. In that directory, there are several files. The `.OBS` files are the observed data. The other files are previous inversion results from the UBC inversion. For the two `.OBS` files: `POT` referrs to \"Potentials\" - these are our DC data and `IP` referrs to \"Induced Polarization\". We will start by working with the DC resistivity data and so will load up that file first. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "line = \"46800E\"\n",
    "print(os.listdir(os.path.join('century',line)))\n",
    "\n",
    "dc_data_file = f\"./century/{line}/{line[:-1]}POT.OBS\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "These file formats are an older version of the UBC DC format and are very similar to the current [UBC DC IP Surface Format](https://giftoolscookbook.readthedocs.io/en/latest/content/fileFormats/dcip2dObsfile.html#surface-format)\n",
    "\n",
    "```\n",
    "Comment line: description of data\n",
    "   n_sources dipole(1) dipole(1) \n",
    "     A_location  B_location   n_receivers\n",
    "   M_location  N_location  measured_potential  standard_deviation\n",
    "   M_location  N_location  measured_potential  standard_deviation\n",
    "     A_location  B_location   n_receivers\n",
    "   M_location  N_location  measured_potential  standard_deviation\n",
    "   M_location  N_location  measured_potential  standard_deviation\n",
    "   ....\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "with open(dc_data_file) as f:\n",
    "    print(f.read())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we write a function to read in these data. Within SimPEG, there are lots of utilities for reading in data, particularly if they are in common formats such as the more recent UBC formats (e.g. see `dc.utils.readUBC_DC2Dpre`). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def read_dcip_data(filename, verbose=True):\n",
    "    \"\"\"\n",
    "    Read in a .OBS file from the Century data set into a python dictionary. \n",
    "    The format is the old UBC-GIF DCIP format.\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    filename : str\n",
    "        Path to the file to be parsed\n",
    "    \n",
    "    verbose: bool\n",
    "        Print some things? \n",
    "    \n",
    "    \n",
    "    Returns\n",
    "    -------\n",
    "    dict\n",
    "        A dictionary with the locations of\n",
    "        - a_locations: the positive source electrode locations (numpy array) \n",
    "        - b_locations: the negative source electrode locations (numpy array) \n",
    "        - m_locations: the receiver locations (list of numpy arrays)\n",
    "        - n_locations: the receiver locations (list of numpy arrays)\n",
    "        - observed_data: observed data (list of numpy arrays)\n",
    "        - standard_deviations: assigned standard deviations (list of numpy arrays)\n",
    "        - n_sources: number of sources (int)\n",
    "    \n",
    "    \"\"\"\n",
    "    \n",
    "    # read in the text file as a numpy array of strings (each row is an entry)\n",
    "    contents = np.genfromtxt(filename, delimiter=' \\n', dtype=np.str)\n",
    "    \n",
    "    # the second line has the number of sources, current, and data type (voltages if 1)\n",
    "    n_sources = int(contents[1].split()[0])\n",
    "    \n",
    "    if verbose is True: \n",
    "        print(f\"number of sources: {n_sources}\")\n",
    "    \n",
    "    # initialize storage for the electrode locations and data\n",
    "    a_locations = np.zeros(n_sources)\n",
    "    b_locations = np.zeros(n_sources)\n",
    "    m_locations = []\n",
    "    n_locations = []\n",
    "    observed_data = []\n",
    "    standard_deviations = []\n",
    "    \n",
    "    # index to track where we have read in content \n",
    "    content_index = 1 \n",
    "    \n",
    "    # loop over sources \n",
    "    for i in range(n_sources):\n",
    "        # start by reading in the source info \n",
    "        content_index = content_index + 1  # read the next line\n",
    "        a_location, b_location, nrx = contents[content_index].split()  # this is a string\n",
    "\n",
    "        # convert the strings to a float for locations and an int for the number of receivers\n",
    "        a_locations[i] = float(a_location)\n",
    "        b_locations[i] = float(b_location)\n",
    "        nrx = int(nrx)\n",
    "\n",
    "        if verbose is True: \n",
    "            print(f\"Source {i}: A-loc: {a_location}, B-loc: {b_location}, N receivers: {nrx}\")\n",
    "\n",
    "        # initialize space for receiver locations, observed data associated with this source\n",
    "        m_locations_i, n_locations_i = np.zeros(nrx), np.zeros(nrx)\n",
    "        observed_data_i, standard_deviations_i = np.zeros(nrx), np.zeros(nrx)\n",
    "\n",
    "        # read in the receiver info \n",
    "        for j in range(nrx):\n",
    "            content_index = content_index + 1  # read the next line\n",
    "            m_location, n_location, datum, std = contents[content_index].split()\n",
    "\n",
    "            # convert the locations and data to floats, and store them\n",
    "            m_locations_i[j] = float(m_location)\n",
    "            n_locations_i[j] = float(n_location)\n",
    "            observed_data_i[j] = float(datum)\n",
    "            standard_deviations_i[j] = float(std)\n",
    "\n",
    "        # append the receiver info to the lists\n",
    "        m_locations.append(m_locations_i)\n",
    "        n_locations.append(n_locations_i)\n",
    "        observed_data.append(observed_data_i)\n",
    "        standard_deviations.append(standard_deviations_i)\n",
    "    \n",
    "    return {\n",
    "        \"a_locations\": a_locations,\n",
    "        \"b_locations\": b_locations, \n",
    "        \"m_locations\": m_locations,\n",
    "        \"n_locations\": n_locations,\n",
    "        \"observed_data\": observed_data, \n",
    "        \"standard_deviations\": standard_deviations,\n",
    "        \"n_sources\": n_sources, \n",
    "    }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dc_data_dict = read_dcip_data(dc_data_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for key, value in dc_data_dict.items():\n",
    "    print(f\"{key:<20}: {type(value)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dc_data_dict[\"a_locations\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "At this point, we have not used anything in SimPEG, this is all pure python. Now, we will look at using SimPEG. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 2: Create a SimPEG survey\n",
    "\n",
    "In SimPEG, the `survey` object keeps track of the geometry and type of sources and receivers.  Similar to the way the data file was structured, each `source` takes a `receiver_list` which knows about the different receiver types that are \"listening\" to that source. \n",
    "\n",
    "For example, for the first source, we have dipole receivers at 2 different locations. A SimPEG receiver object can have multiple locations, so conceptually, we can think of a receiver object containing all receivers of one type. If we also had a pole-source, then we would add another receiver type to our receiver list. Since we will run 2D simulations and inversions, then the receivers need to be defined in 2D space `[x, z]`. The receivers are on the surface of the Earth, and in this example, we will not consider topography, so they are at `z=0`.\n",
    "\n",
    "```python\n",
    "\n",
    "m_locs = np.array([\n",
    "    [26700.0, 0.],\n",
    "    [26800.0, 0.]\n",
    "])\n",
    "\n",
    "n_locs = np.array([\n",
    "    [26800.0, 0.],\n",
    "    [26900.0, 0.]\n",
    "])\n",
    "\n",
    "rx = dc.receivers.Dipole(locations_m=m_locs, locations_n=n_locs) \n",
    "rx_list = [rx]\n",
    "\n",
    "```\n",
    "\n",
    "Note that in SimPEG, we use column-major / Fortran ordering. So `m_locs[:, 0]` corresponds to all of the x-locations and `m_locs[:, 1]` corresponds to all of the z-locations.  \n",
    "\n",
    "A source in SimPEG has a single location. For the first source: \n",
    "```python \n",
    "\n",
    "a_loc = np.r_[26000.0, 0.]\n",
    "b_loc = np.r_[26100.0, 0.]\n",
    "\n",
    "src = dc.sources.Dipole(location_a=a_loc, location_b=b_loc, receiver_list=rx_list)\n",
    "```\n",
    "\n",
    "Finally, the survey is instantiated with a source list which contains all of the sources and their associated receiver lists\n",
    "```python\n",
    "survey = dc.Survey(source_list=[src1, src2, ...])\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# initialize an empty list for each \n",
    "source_list = []\n",
    "\n",
    "for i in range(dc_data_dict[\"n_sources\"]):\n",
    "    \n",
    "    # receiver electrode locations in 2D \n",
    "    m_locs = np.vstack([\n",
    "        dc_data_dict[\"m_locations\"][i], \n",
    "        np.zeros_like(dc_data_dict[\"m_locations\"][i])\n",
    "    ]).T\n",
    "    n_locs = np.vstack([\n",
    "        dc_data_dict[\"n_locations\"][i],\n",
    "        np.zeros_like(dc_data_dict[\"n_locations\"][i])\n",
    "    ]).T\n",
    "    \n",
    "    # construct the receiver object \n",
    "    receivers = dc.receivers.Dipole(locations_m=m_locs, locations_n=n_locs)\n",
    "    \n",
    "    # construct the source \n",
    "    source = dc.sources.Dipole(\n",
    "        location_a=np.r_[dc_data_dict[\"a_locations\"][i], 0.],\n",
    "        location_b=np.r_[dc_data_dict[\"b_locations\"][i], 0.],\n",
    "        receiver_list=[receivers]\n",
    "    )\n",
    "    \n",
    "    # append the new source to the source list\n",
    "    source_list.append(source)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "survey = dc.Survey(source_list=source_list)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Observed Data \n",
    "\n",
    "The data object keeps track of our observed data, the uncertainties and the survey parameters corresponding to each datum. The order of the data and standard deviations should be identical to how we would unwrap the source and receiver lists. E.g.\n",
    "\n",
    "```\n",
    "dobs = [d(source1, receiver1), d(source1, receiver2), d(source2, receiver1), ...]\n",
    "```\n",
    "\n",
    "Since this is exactly the same order as how we read the data in, we can simply concatenate the data in our dictionary. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dc_data = Data(\n",
    "    survey=survey, \n",
    "    dobs=np.hstack(dc_data_dict[\"observed_data\"]),\n",
    "    standard_deviation=np.hstack(dc_data_dict[\"standard_deviations\"])\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Visualizing data in a Pseudosection\n",
    "\n",
    "Pseudosections give us a way to visualize a profile of DC resistivity or IP data. The horizontal location is the mid-point between the center of the source AB electrodes and the MN electrodes. The vertical location corresponds to the \"n-spacing\" which is the separation distance between the AB and MN electrodes divided by the length of the AB / MN dipoles. For this survey, the dipole length is 100m, so an n-spacing of 2 corresponds to a distance of 200m between the midpoints of the AB and MN electrodes. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# plot psuedosection\n",
    "fig, ax = plt.subplots(1, 1, figsize=(12, 4))\n",
    "dc.utils.plot_pseudosection(\n",
    "    dc_data, data_type=\"potential\",\n",
    "    plot_type=\"contourf\", data_location=True, ax=ax,\n",
    ")\n",
    "ax.set_aspect(1.5)  # some vertical exxageration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_building_pseudosection(source_ind=0):\n",
    "    \"\"\"\n",
    "    A plotting function to visualize how a pseudosection is built up. \n",
    "    \"\"\"\n",
    "    \n",
    "    fig, ax = plt.subplots(1, 1, figsize=(12, 4))\n",
    "    mid_x, mid_z = dc.utils.source_receiver_midpoints(dc_data.survey)\n",
    "    ax.plot(mid_x, mid_z, '.k', alpha=0.5)\n",
    "    \n",
    "    # plot the source location\n",
    "    a_loc = dc_data_dict[\"a_locations\"][source_ind]\n",
    "    b_loc = dc_data_dict[\"b_locations\"][source_ind]\n",
    "    src_mid = (a_loc+b_loc)/2\n",
    "    ax.plot(np.r_[a_loc, b_loc], np.r_[0, 0], 'C0.')\n",
    "    ax.plot(src_mid, np.r_[0], 'C0x', ms=6)\n",
    "\n",
    "    # plot the receiver locations\n",
    "    m_locs = dc_data_dict[\"m_locations\"][source_ind]\n",
    "    n_locs = dc_data_dict[\"n_locations\"][source_ind]\n",
    "    rx_mid = (m_locs+n_locs)/2\n",
    "    ax.plot(np.r_[m_locs, n_locs], np.zeros(2*len(m_locs)), 'C1.')\n",
    "    ax.plot(rx_mid, np.zeros_like(m_locs), 'C1x', ms=6)\n",
    "    \n",
    "    # plot where the pseudosection points should be \n",
    "    pseudo_x = (rx_mid + src_mid)/2.\n",
    "    pseudo_z = -np.abs(rx_mid-src_mid)/2.\n",
    "    \n",
    "    ax.plot(np.r_[src_mid, pseudo_x], np.r_[0, pseudo_z], '-k', alpha=0.3)\n",
    "    \n",
    "    # draw lines to the points we are plotting\n",
    "    for rx_x, px, pz in zip(rx_mid, pseudo_x, pseudo_z): \n",
    "        ax.plot(np.r_[px, rx_x], np.r_[pz, 0], '-k', alpha=0.3)\n",
    "    \n",
    "    # these are the data corresponding to the source-receiver pairs we are looking at\n",
    "    ax.plot(pseudo_x, pseudo_z, 'C2o')\n",
    "    ax.set_xlabel(\"Northing (m)\")\n",
    "    ax.set_ylabel(\"n-spacing\")\n",
    "    \n",
    "    ax.set_yticklabels([])\n",
    "    ax.set_xlim([25900, 29300])\n",
    "    ax.set_aspect(1.5) \n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ipywidgets.interact(\n",
    "    plot_building_pseudosection, \n",
    "    source_ind=ipywidgets.IntSlider(min=0, max=int(dc_data_dict[\"n_sources\"])-1, value=0)\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Apparent resistivity\n",
    "\n",
    "Our measured data are potential differences in units of volts. Both the geometry of the sources and receivers, as well as the geology, influence these values. Having both of these factors influence the values can make it challenging to make sense of the data. In order to translate them to something more interpretable, it is common to work with apparent resistivity values, which are derived by considering the potentials that would be measured if the earth were a half-space: \n",
    "\n",
    "$$\n",
    "\\rho_a = \\frac{2 \\pi \\Delta V}{IG}\n",
    "$$\n",
    "\n",
    "Where:\n",
    "- $\\rho_a$ is the apparent resistivity ($\\Omega$ m)\n",
    "- $\\Delta V$ is the measured potential difference (V)\n",
    "- $I$ is the source current (A)\n",
    "- $G$ is the geometric factor which depends on the location of all four electrodes. \n",
    "    $$\n",
    "    G = \\left(\\frac{1}{r_{AM}} - \\frac{1}{r_{BM}} - \\frac{1}{r_{AN}} + \\frac{1}{r_{BN}} \\right)\n",
    "    $$\n",
    "\n",
    "Pseudosections of the apparent resistivity are a valuable tool for QC'ing and processing the data. It is easier to spot outliers and see trends in the data than using plots of the potentials."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot psuedosection\n",
    "fig, ax = plt.subplots(1, 1, figsize=(12, 4))\n",
    "dc.utils.plot_pseudosection(\n",
    "    dc_data, data_type=\"apparent resistivity\", \n",
    "    plot_type=\"contourf\", data_location=True, ax=ax, cbar_opts={\"pad\":0.25}\n",
    ")\n",
    "ax.set_aspect(1.5)  # some vertical exxageration\n",
    "ax.set_title(f\"DC: {line} Pseudosection\")\n",
    "ax.set_xlabel(\"Northing (m)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Histogram of our data\n",
    "\n",
    "In addition to looking at a 2D pseudosection of the data, a histogram gives us an idea of the distribution of apparent resistivities. For more complex survey geometries (e.g. fully 3D), it is not always clear how to plot data in a pseudosection form, so histograms are a valueable tool for visualizing data and estimating a good background resistivity for the starting and reference models. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "apparent_resistivity = dc.utils.apparent_resistivity(dc_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, 1)\n",
    "out = ax.hist(np.log10(apparent_resistivity), bins=30)\n",
    "ax.set_xlabel(\"log$_{10}(\\\\rho_a)$\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "estimate a background model "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# rho0 = 10**np.mean(np.log10(apparent_resistivity))\n",
    "rho0 = np.median(apparent_resistivity)\n",
    "rho0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Step 3: set up the forward simulation machinery\n",
    "\n",
    "### The DC resistivity equations\n",
    "Now that we have a sense of our data, we will set up a forward simulation which is responsible for simulating predicted data given a model. \n",
    "\n",
    "The equations we are solving are the steady-state Maxwell's equations: \n",
    "\n",
    "$$\n",
    "\\nabla \\times \\vec{e} = 0 \\quad \\to \\quad \\vec{e} = -\\nabla \\phi \\\\\n",
    "\\nabla \\cdot \\vec{j} = - \\nabla \\cdot \\vec{j}_s = I \\delta(r)\n",
    "$$\n",
    "\n",
    "Since the curl of the electric field is zero, we can represent it as a the gradient of a scalar potential $\\phi$.\n",
    "\n",
    "Ohm's law relates the current density to the electric field through the resistivity, $\\rho$, in units of $\\Omega$m (or equivalently through the electrical conductivity $\\sigma = 1/\\rho$.  \n",
    "\n",
    "$$\n",
    "\\vec{j} = \\frac{1}{\\rho} \\vec{e}\n",
    "$$\n",
    "\n",
    "and assembling everything\n",
    "\n",
    "$$\n",
    "\\nabla \\cdot \\frac{1}{\\rho} \\nabla \\phi = I \\delta(r) \n",
    "$$\n",
    "\n",
    "### Discretizing\n",
    "\n",
    "In SimPEG, we use a Finite Volume approach to solving this partial differential equation (PDE). To do so, we define where each component of the equation lives on the mesh. Scalars can either be defined at cell centers or on nodes whereas vectors are defined on cell faces or edges. There are two discretizations implemented in SimPEG for the DC resistivity problem: Cell-Centered and Nodal. In both, we discretize physical properties (resistivity) at cell centers, and you can think of the resistivity filling that entire \"lego-block\". The difference between the two discretizations depends upon where we discretize the potential $\\phi$: either at the cell centers or the nodes. For this tutorial, we will use the **Nodal discretization** where\n",
    "- $\\rho$: discretized at cell centers\n",
    "- $\\phi$: discretized at nodes\n",
    "- $\\vec{e}$: discretized on cell edges\n",
    "\n",
    "<img src=\"images/finite-volume-02.png\" width=25%>\n",
    "\n",
    "We will not go into the details here, but we highlight the importance of the boundary conditions. The Nodal discretization uses Neuman boundary conditions, which means that there is no flux of current through the boundary. Thus when we design our mesh, it needs to extend sufficiently far as to satisfy this condition. \n",
    "\n",
    "$$\n",
    "\\frac{1}{\\rho}\\vec{e} = 0 | \\partial\\Omega\n",
    "$$\n",
    "\n",
    "\n",
    "\n",
    "For an overview of the derivation for the cell centered discretization, see the leading edge article [_Pixels and their Neighbors_](https://github.com/simpeg/tle-finitevolume/blob/master/Pixels_and_neighbors.pdf). \n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.1 Design a Mesh\n",
    "\n",
    "When designing a mesh, there are 3 main items to consider\n",
    "- minimum cell size: \n",
    "    - controls resolution \n",
    "    - typically ~4 per minimum electrode spacing\n",
    "    \n",
    "    \n",
    "- extent of the core domain:  \n",
    "    - in the horizontal dimension: typically use the extent of the survey (or slightly beyond)\n",
    "    - in the vertical direction, we want the core region to extend to the depth to which we are sensitive to variations in physical properties. As a rule of thumb, this is typically a fraction of the maximum source-receiver separation\n",
    "    \n",
    "    \n",
    "- extent of the full modelling domain: \n",
    "   - need to satisfy boundary conditions! \n",
    "   - typically add ~the maximum source-receiver separation in each dimension\n",
    "   \n",
    "\n",
    "Here, we will use a `TensorMesh`, where each dimension can be defined with just a cell width. \n",
    "A tensor in each dimension is enough to define the entire mesh.\n",
    "For larger domains, we recommend using an `TreeMesh`. There are [examples in the discretize documentation](http://discretize.simpeg.xyz/en/master/tutorials/mesh_generation/4_tree_mesh.html) for creating QuadTree (2D) and OcTree (3D) meshes. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "min_electrode_spacing = np.min(np.abs(survey.locations_a[:, 0] - survey.locations_b[:, 0]))\n",
    "n_cells_per_spacing = 4\n",
    "\n",
    "dx = min_electrode_spacing / n_cells_per_spacing\n",
    "dz = min_electrode_spacing / n_cells_per_spacing\n",
    "\n",
    "print(\n",
    "    f\"We will use a base mesh with {dx} m x {dz} m cells\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "core_domain_x = np.r_[\n",
    "    survey.electrode_locations[:, 0].min(),\n",
    "    survey.electrode_locations[:, 0].max()\n",
    "]\n",
    "\n",
    "core_domain_x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# find the maximum spacing between source, receiver midpoints\n",
    "mid_ab = (survey.locations_a + survey.locations_b)/2\n",
    "mid_mn = (survey.locations_m + survey.locations_n)/2\n",
    "separation_ab_mn = np.abs(mid_ab - mid_mn)\n",
    "max_separation = separation_ab_mn.max()\n",
    "print(f\"The maximum separation is {max_separation}m\")\n",
    "\n",
    "core_domain_z = np.r_[-max_separation/3., 0.]\n",
    "core_domain_z"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_core_extra = 4 \n",
    "n_core_x = np.ceil(np.diff(core_domain_x)/dx) + n_core_extra*2  # on each side\n",
    "n_core_z = np.ceil(np.diff(core_domain_z)/dz) + n_core_extra  # just below\n",
    "\n",
    "padding_factor = 1.3\n",
    "n_pad_x = 10\n",
    "n_pad_z = 10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hx = [(dx, n_pad_x, -padding_factor), (dx, n_core_x), (dx, n_pad_x, padding_factor)]\n",
    "hz = [(dz, n_pad_z, -padding_factor), (dz, n_core_z)]\n",
    "\n",
    "mesh = discretize.TensorMesh([hx, hz])\n",
    "\n",
    "# origin of the mesh\n",
    "mesh.x0 = np.r_[\n",
    "    -mesh.hx[:n_pad_x + n_core_extra].sum() + core_domain_x.min(),\n",
    "    -mesh.hy.sum()\n",
    "]\n",
    "\n",
    "mesh"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, 1, figsize=(12, 4))\n",
    "mesh.plotGrid(ax=ax)\n",
    "ax.set_aspect(1.5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.2: Construct the forward simulation\n",
    "\n",
    "A `simulation` object contains everything we need to compute potentials given a resistivity model. It is instantiated with:\n",
    "\n",
    "- `mesh`: which discretizes the world and provides us with all of the differential operators\n",
    "\n",
    "\n",
    "- `rhoMap`: in SimPEG, the `maps` module provides flexibility for defining what our inversion model is. Each `Map` is a mathematical transform (that has a derivative - which we need in the inversion!). For example, the `ExpMap` takes the exponential of a model vector; we would define our model as log-resistivity. This is useful in electrical and electromagnetic inversions where we know that electrical resistivity should be strictly positive, and tends to vary logarithmically. \n",
    "<img src=\"https://gpg.geosci.xyz/_images/resistivity_table.png\" width=70%>\n",
    "  Similarly, mappings can be used to define active and inactive cells (for example if we have topography and don't want to invert for air-cells), or to define parametric models. \n",
    "  \n",
    "\n",
    "- `solver`: A linear solver (e.g. LU) provided in the [pymatsolver repository](https://pymatsolver.readthedocs.io/en/latest/). Here we use Pardiso. \n",
    "\n",
    "\n",
    "- `storeJ`: for smaller problems, where we aren't memory limited, the inversion is faster if we store the sensitivity matrix rather than "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# forward simulation \n",
    "# Use Exponential Map: m = log(rho)\n",
    "mapping = maps.ExpMap(mesh)\n",
    "\n",
    "# Generate 2.5D DC problem\n",
    "simulation_dc = dc.Simulation2DNodal(\n",
    "    mesh, rhoMap=mapping, solver=Solver, survey=survey, storeJ=True\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To check our mesh and forward simulation, we run a simulation over a half-space. The apparent resistivities we compute from the simulated data should be equal to the true half-space resistivity. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time \n",
    "\n",
    "# run the forward simulation over the half-space & plot apparent resistivities\n",
    "m0 = np.ones(mesh.nC) * np.log(rho0)\n",
    "d0 = simulation_dc.make_synthetic_data(m0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot psuedosection\n",
    "fig, ax = plt.subplots(1, 1, figsize=(12, 4))\n",
    "\n",
    "clim = np.r_[apparent_resistivity.min(), apparent_resistivity.max()]\n",
    "dc.utils.plot_pseudosection(\n",
    "    d0, data_type=\"apparent resistivity\",  clim=clim,\n",
    "    plot_type=\"contourf\", data_location=True, ax=ax, \n",
    "    cbar_opts={\"pad\":0.25}\n",
    ")\n",
    "ax.set_aspect(1.5)  # some vertical exxageration\n",
    "ax.set_title(f\"{line} Pseudosection\")\n",
    "ax.set_xlabel(\"Northing (m)\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This next plot shows the percent error in the apparent resistivity. We want the error to be less than the errors in our data: 5% is common for assigning a percent uncertainty, so if the errors are below that, then we can be satisfied with our setup. If not, there are several things to check. In general (for nearly every geophysical survey):\n",
    "- boundary conditions - does the mesh extend far enough?\n",
    "- discretization near source and receivers - are we using fine-enough cells\n",
    "- specifically for 2.5 DC: are we using enough filters (we will revisit this in the next notebook) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "apparent_resistivity_m0 = dc.utils.apparent_resistivity(d0)\n",
    "percent_error = (apparent_resistivity_m0 - rho0)/rho0*100\n",
    "\n",
    "fig, ax = plt.subplots(1, 1)\n",
    "ax.plot(separation_ab_mn[:, 0], percent_error, '.', alpha=0.4)\n",
    "ax.set_xlabel(\"AB - MN separation\")\n",
    "ax.set_ylabel(\"percent difference (%)\")\n",
    "ax.set_title(\"Error in apparent resistivities\")\n",
    "ax.set_ylim(7*np.r_[-1, 1])\n",
    "ax.grid(alpha=0.3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 4. DC resistivity inversion\n",
    "\n",
    "We forumlate the inverse problem as an optimization problem consisting of a data misfit and a regularization \n",
    "\n",
    "$$\n",
    "\\min_{\\mathbf{m}} \\Phi(\\mathbf{m}) = \\Phi_d(\\mathbf{m}) + \\beta\\Phi_m(\\mathbf{m}) \\\\ s.t. ~ \\Phi_d \\leq \\Phi_d^* \\quad \\mathbf{m}_i^{\\rm L} \\leq \\mathbf{m}_ \\leq \\mathbf{m}_i^{\\rm U}\n",
    "$$\n",
    "\n",
    "where:\n",
    "- $\\mathbf{m}$ is our inversion model - a vector containing the set of parameters that we invert for\n",
    "- $\\Phi_d$ is the data misfit\n",
    "- $\\Phi_m$ is the regularization\n",
    "- $\\beta$ is a trade-off parameter that weights the relative importance of the data misfit and regularization terms\n",
    "- $\\Phi_d^*$ is our target misfit, which is typically set to $N/2$ where $N$ is the number of data (Parker, 1994) (or also see [Oldenburg & Li (2005)](https://www.researchgate.net/profile/Douglas_Oldenburg/publication/238708196_5_Inversion_for_Applied_Geophysics_A_Tutorial/links/004635282572529927000000.pdf))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4.1 Data Misfit\n",
    "\n",
    "The data misfit, $\\Phi_d(\\mathbf{m})$, is often taken to be a weighted $\\ell_2$-norm, where the weights capture the noise model (eg. we want to assign higher weights and do a good job fitting data that we are confident are less noisy, and assign less weight / influence to data that are noisy). The $\\ell_2$ norm is the correct norm to choose when noise is Gaussian (or approximately Gaussian, or if you have no additional information and assume it is Gaussian). An $\\ell_2$ data misfit is captured mathematically by  \n",
    "\n",
    "$$\n",
    "\\Phi_d(\\mathbf{m}) = \\frac{1}{2}\\|\\mathbf{W_d} (\\mathcal{F}(\\mathbf{m}) - \\mathbf{d}^{\\text{obs}})\\|^2\n",
    "$$\n",
    "\n",
    "where \n",
    "\n",
    "- $\\mathbf{W_d} = \\text{diag}(1/\\boldsymbol{\\epsilon})$ is a diagonal matrix with diagonal entries $1/\\epsilon_i$, where $\\epsilon_i$ is an estimated standard deviation of the $i$th datum.\n",
    "- $\\mathcal{F}$ is the forward modelling operator that simulates predicted data given a model\n",
    "- $\\mathbf{d}^{\\text{obs}}$ is the observed data\n",
    "\n",
    "To instantiate the data misfit term, we provide a `data` object (which also has the associated uncertainties) along with a `simulation` which has the machinery to compute predicted data. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# set up the data misfit\n",
    "dmisfit = data_misfit.L2DataMisfit(data=dc_data, simulation=simulation_dc)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4.2 Regularization\n",
    "The inverse problem is an [ill posed problem](https://en.wikipedia.org/wiki/Well-posed_problem). There are multiple (actually infinitely many!) models that can fit the data.  If we start by thinking about a linear problem $\\mathbf{G}\\mathbf{m} = \\mathbf{d}$, the matrix $\\mathbf{G}$ has more columns than rows, so it is not directly invertible (eg. see Matt Hall's [Linear Inversion Tutorial](http://library.seg.org/doi/full/10.1190/tle35121085.1)). Here, we are dealing with a non-linear system of equations, but the principle is the same. \n",
    "\n",
    "Thus, to choose from the infinitely many solutions and arrive at a sensible one, we employ a regularization: $\\Phi_m$. Tikhonov regularization, which again uses $\\ell_2$-norms, is a standard choice (It has a few nice features: it is convex and easy to differentiate). It takes the form: \n",
    "$$\n",
    "\\Phi_m(\\mathbf{m}) = \\frac{1}{2}\\big(\\alpha_s\\|\\mathbf{W_s} (\\mathbf{m} - \\mathbf{m}_{\\text{ref}})\\|^2 + \\alpha_x\\|\\mathbf{W_x} (\\mathbf{m})\\|^2 + \\alpha_z\\|\\mathbf{W_z} (\\mathbf{m})\\|^2 \\big)\n",
    "$$\n",
    "\n",
    "\n",
    "The first term is often referred to as the \"smallness\" as it measures the \"size\" of the model (in the $\\ell_2$ sense). The matrix $\\mathbf{W_s}$ is generally taken to be a diagonal matrix that may contain information about the length scales of the model or be used to weight the relative importance of various parameters in the model. The scalar $\\alpha_s$ weights the relative importance of this term in the regularization. \n",
    "\n",
    "Notice that we include a reference model ($\\mathbf{m}_{\\text{ref}}$. Often this is defined as a constant value, but if more information is known about the background, that can be used to construct a reference model. Note that saying \"I am not going to use a reference model\" means that you are actually using $\\mathbf{m}_{\\text{ref}} = 0$, this is important to realize... in the inversion we demonstrate here, our `model` will be $\\mathbf{m} = \\text{log}(\\rho)$. If we set $\\mathbf{m}_{\\text{ref}} = 0$, then we are favoring models close to 1 $\\Omega$m - which is quite conductive! \n",
    "\n",
    "The second term is often referred to as the \"smoothness\". The matrix $\\mathbf{W_{\\{x, z\\}}}$ approximate the derivative of the model with respect to depth, and is hence a measure of how \"smooth\" the model is. The term $\\alpha_{\\{x,z\\}}$ weights its relative importance in the regularization. \n",
    "\n",
    "The values $\\alpha_x, \\alpha_z$ have dimensions of $(length)$^2 as compared to $\\alpha_s$. So if we choose our length scale to be 25m, the size of our finest cells and $\\alpha_x = 1$, then we set $\\alpha_s = 1/25^2$. In notebook [3-inversions-dc.ipynb](3-inversions-dc.ipynb). We will explore these choices further.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Related to inversion\n",
    "reg = regularization.Tikhonov(\n",
    "    mesh, \n",
    "    alpha_s=1./mesh.hx.min()**2,  # we will talk about these choices in notebook 3 on inversion\n",
    "    alpha_x=1, \n",
    "    alpha_y=1,  # since this is a 2D problem, the second dimension is \"y\" \n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we assemble the data misfit and regularization into a statement of the inverse problem. We select an optimization method: in this case, a second-order Inexact Gauss Newton (see [Nocedal and Wright, Ch 7](https://link.springer.com/book/10.1007/978-0-387-40065-5)). The `inverse_problem` then has all of the machinery necessary to evaluate the objective function and compute an update.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "opt = optimization.InexactGaussNewton(maxIter=20, maxIterCG=20)\n",
    "inv_prob = inverse_problem.BaseInvProblem(dmisfit, reg, opt)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4.3 Assemble the inversion\n",
    "\n",
    "In SimPEG, we use `directives` to orchestrate updates during the inversion. The three we use below are\n",
    "\n",
    "- `betaest`: sets an initial value for $\\beta$ by estimating the largest eigenvalue of  $\\Phi_d$ and of $\\Phi_m$ and then taking their ratio. The value is then scaled by the `beta0_ratio` value, so for example if you wanted the regularization to be ~10 times more important than the data misfit, then we would set `beta0_ratio=10`\n",
    "\n",
    "\n",
    "- `beta_schedule`: this reduces the value of beta during the course of the inversion. Particularly for non-linear problems, it can be advantageous to start off by seeking a smooth model which fits the regularization and gradually decreasing the influence of the regularization to increase the influence of the data. \n",
    "\n",
    "\n",
    "- `target`: this directive checks at the end of each iteration if we have reached the target misfit. If we have, the inversion terminates, it not, it continues. \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# directives\n",
    "beta_est = directives.BetaEstimate_ByEig(beta0_ratio=1e0)\n",
    "beta_schedule = directives.BetaSchedule(coolingFactor=4, coolingRate=2)\n",
    "target = directives.TargetMisfit()\n",
    "\n",
    "inv = inversion.BaseInversion(\n",
    "    inv_prob, directiveList=[beta_est, beta_schedule, target]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4.4 Run the inversion!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run inversion\n",
    "mopt = inv.run(m0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# plot psuedosection\n",
    "fig, ax = plt.subplots(2, 1, figsize=(12, 8))\n",
    "\n",
    "clim = np.r_[apparent_resistivity.min(), apparent_resistivity.max()]\n",
    "\n",
    "for a, data, title in zip(\n",
    "    ax,\n",
    "    [dc_data.dobs, inv_prob.dpred],\n",
    "    [\"true\", \"recovered\"]\n",
    "): \n",
    "    dc.utils.plot_pseudosection(\n",
    "        dc_data, dobs=data, data_type=\"apparent resistivity\",\n",
    "        plot_type=\"contourf\", data_location=True, ax=a, \n",
    "        cbar_opts={\"pad\":0.25}, clim=clim\n",
    "    )\n",
    "    a.set_title(title)\n",
    "\n",
    "\n",
    "for a in ax: \n",
    "    a.set_aspect(1.5)  # some vertical exxageration\n",
    "    a.set_xlabel(\"Northing (m)\")\n",
    "    a.set_ylabel(\"Pseudo-depth\")\n",
    "\n",
    "plt.tight_layout()\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot misfit\n",
    "fig, ax = plt.subplots(1, 1, figsize=(12, 4))\n",
    "\n",
    "normalized_misfit = (inv_prob.dpred - dc_data.dobs) / dc_data.standard_deviation\n",
    "\n",
    "dc.utils.plot_pseudosection(\n",
    "    dc_data, dobs=normalized_misfit, data_type=\"misfit\",\n",
    "    plot_type=\"contourf\", data_location=True, ax=ax, \n",
    "    cbar_opts={\"pad\":0.25}\n",
    ")\n",
    "\n",
    "ax.set_title(\"DC misfit\")\n",
    "ax.set_aspect(1.5)  # some vertical exxageration\n",
    "ax.set_xlabel(\"Northing (m)\")\n",
    "ax.set_ylabel(\"n-spacing\")\n",
    "\n",
    "cb_axes = plt.gcf().get_axes()[-1]\n",
    "cb_axes.set_xlabel('normalized misfit')\n",
    "\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# run inversion, plot tikhonov curves\n",
    "\n",
    "def load_leapfrog_geologic_section(filename=\"./century/geologic_section.csv\"):\n",
    "    \"\"\"\n",
    "    Load the geologic cross section. \n",
    "    \"\"\"\n",
    "    fid = open(filename, 'r')\n",
    "    lines = fid.readlines()\n",
    "    data = []\n",
    "    data_tmp = []\n",
    "    for line in lines[2:]:\n",
    "        line_data = (line.split(',')[:3])\n",
    "        if 'End' in line:\n",
    "            data.append(np.vstack(data_tmp)[:,[0, 2]])\n",
    "            data_tmp = []\n",
    "        else:\n",
    "            data_tmp.append(np.array(line_data, dtype=float))\n",
    "    return data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "geologic_section = load_leapfrog_geologic_section()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1,1, figsize=(10, 5))\n",
    "rho = (mapping*mopt)\n",
    "out = mesh.plotImage(\n",
    "    rho, pcolorOpts={'norm':LogNorm(), 'cmap':'Spectral'}, ax=ax,\n",
    "    clim=(30, 1000)\n",
    ")\n",
    "for data in geologic_section:\n",
    "    ax.plot(data[:,0], data[:,1], 'k--', alpha=0.5)\n",
    "ax.set_xlim(core_domain_x)\n",
    "ax.set_ylim((-400, 0))\n",
    "cb = plt.colorbar(out[0], fraction=0.05, orientation='horizontal', ax=ax, pad=0.2)\n",
    "cb.set_label(\"Resistivity ($\\Omega$m)\")\n",
    "ax.set_xlabel('Easting (m)')\n",
    "ax.set_ylabel('Elevation (m)')\n",
    "ax.set_aspect(1.5)  # some vertical exxageration\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Compare to UBC results obtained 20 years ago(!)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mesh_ubc = discretize.TensorMesh.readUBC(\"century/46800E/468MESH.DAT\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def read_ubc_model(filename, mesh_ubc=mesh_ubc): \n",
    "    \"\"\"\n",
    "    A function to read a UBC conductivity or chargeability model. \n",
    "    \"\"\"\n",
    "    values = np.genfromtxt(\n",
    "        filename, delimiter=' \\n',\n",
    "        dtype=str, comments='!', skip_header=1\n",
    "    )\n",
    "    tmp = np.hstack([np.array(value.split(), dtype=float) for value in values])\n",
    "    model_ubc = discretize.utils.mkvc(tmp.reshape(mesh_ubc.vnC, order='F')[:,::-1])\n",
    "    return model_ubc\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sigma_ubc = read_ubc_model(\"century/46800E/DCMODA.CON\")\n",
    "rho_ubc = 1./sigma_ubc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(2, 1, figsize=(10, 8))\n",
    "rho = (mapping*mopt)\n",
    "clim = (30, 1000)\n",
    "\n",
    "out = mesh.plotImage(\n",
    "    rho, pcolorOpts={'norm':LogNorm(), 'cmap':'Spectral'}, ax=ax[0],\n",
    "    clim=clim \n",
    ")\n",
    "cb = plt.colorbar(out[0], fraction=0.05, orientation='horizontal', ax=ax[0], pad=0.2)\n",
    "cb.set_label(\"Resistivity ($\\Omega$m)\")\n",
    "ax[0].set_title(\"SimPEG\")\n",
    "\n",
    "out = mesh_ubc.plotImage(\n",
    "    rho_ubc, pcolorOpts={'norm':LogNorm(), 'cmap':'Spectral'}, ax=ax[1],\n",
    "    clim=clim \n",
    ")\n",
    "cb = plt.colorbar(out[0], fraction=0.05, orientation='horizontal', ax=ax[1], pad=0.2)\n",
    "cb.set_label(\"Resistivity ($\\Omega$m)\")\n",
    "ax[1].set_title(\"UBC\")\n",
    "\n",
    "for a in ax: \n",
    "    for data in geologic_section:\n",
    "        a.plot(data[:,0], data[:,1], 'k--', alpha=0.5)\n",
    "    a.set_xlim(core_domain_x)\n",
    "    a.set_ylim((-400, 0))\n",
    "    a.set_aspect(1.5)  # some vertical exxageration\n",
    "    a.set_xlabel('Easting (m)')\n",
    "    a.set_ylabel('Elevation (m)')\n",
    "\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "----\n",
    "# Induced Polarization \n",
    "\n",
    "Induced Polarization (IP) surveys are sensitive to chargeability - the physical property which describes a material's tendency to become polarized and act like a capacitor. \n",
    "\n",
    "<img src=\"images/chargeability.png\" width=30%>\n",
    "\n",
    "The physical explanation for causes of chargeability are complex and are an active area of research. There are, however, two conceptual models that are insightful for characterizing the chareable behaviour of rocks: membrane polarization and electrode polarization. \n",
    "\n",
    "<img src=\"images/conceptual-model-ip.png\" width = 70%>\n",
    "\n",
    "Physically, the impact of chargeabile material causes the measured voltage to increase in the on-time, and after the current switch-off, the voltage decays. In the on-time, currents are primarily due to the potential difference made by a generator, and distorted by conductivity constrasts in the earth (DC-effects). In addition, polarization charges start to build, and reach steady state in the late on-time. After the current is switched off, all DC currents are immediately gone, but polarization charges remain and will start to discharge and generate a \n",
    "secondary voltage decay (IP). \n",
    "\n",
    "<img src=\"https://em.geosci.xyz/_images/DCIP.png\" width=60%>\n",
    "\n",
    "For more background information on IP, including case histories, see the [EM GeoSci pages on IP](https://em.geosci.xyz/content/geophysical_surveys/ip/index.html).\n",
    "\n",
    "\n",
    "## IP simulation \n",
    "\n",
    "To simulate IP, we use a linearized model where we assume that the IP signals are due to a perturbation in electrical conductvity (inverse of resistivity) \n",
    "\n",
    "$$\n",
    "\\rho_\\eta = \\rho\\frac{1}{1-\\eta} \\quad \\eta \\in [0, 1)\n",
    "$$\n",
    "\n",
    "with an IP datum written as \n",
    "\n",
    "$$\n",
    "d_i^{IP} = \\sum_{j=1}^M J_{ij}\\eta_j \\quad i = 1, ..., N\n",
    "$$\n",
    "\n",
    "where the entries of J are the sensitivities for the DC resistivity problem \n",
    "\n",
    "$$\n",
    "J_{i,j} = \\frac{\\partial\\log\\phi^i}{\\partial\\log\\rho_j}\n",
    "$$\n",
    "\n",
    "or in matrix form \n",
    "\n",
    "$$\n",
    "\\mathbf{d^{IP}} = \\mathbf{J}\\mathbf{\\eta}\n",
    "$$\n",
    "\n",
    "For further discussion, see the [derivation here](https://em.geosci.xyz/content/geophysical_surveys/ip/interpretation.html#linearization)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 5: Load the IP data\n",
    "\n",
    "There are several ways to define IP data, for discussion, please see the [reference page on IP data in em.geosci.xyz](https://em.geosci.xyz/content/geophysical_surveys/ip/data.html). For this survey, the data that we are provided are apparent chargeabilities (mV/V). This is analagous to apparent resistivity that we discussed for DC resistivity. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The file format and organization is identical to what we implemented for the DC data, so we can re-use that function to load our data into a dictionary. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ip_data_file = f\"./century/{line}/{line[:-1]}IP.OBS\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ip_data_dict = read_dcip_data(ip_data_file, verbose=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The organization of an IP survey is very similar to the DC case. Here, we are calling sources and receivers from the `ip` module, and also specifying that the receivers are measring apparent chargeability data. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# initialize an empty list for each \n",
    "source_list_ip = []\n",
    "\n",
    "for i in range(ip_data_dict[\"n_sources\"]):\n",
    "    \n",
    "    # receiver electrode locations in 2D \n",
    "    m_locs = np.vstack([\n",
    "        ip_data_dict[\"m_locations\"][i], \n",
    "        np.zeros_like(ip_data_dict[\"m_locations\"][i])\n",
    "    ]).T\n",
    "    n_locs = np.vstack([\n",
    "        ip_data_dict[\"n_locations\"][i],\n",
    "        np.zeros_like(ip_data_dict[\"n_locations\"][i])\n",
    "    ]).T\n",
    "    \n",
    "    # construct the receiver object \n",
    "    receivers = ip.receivers.Dipole(\n",
    "        locations_m=m_locs, locations_n=n_locs, data_type=\"apparent_chargeability\"\n",
    "    )\n",
    "    \n",
    "    # construct the source \n",
    "    source = ip.sources.Dipole(\n",
    "        location_a=np.r_[ip_data_dict[\"a_locations\"][i], 0.],\n",
    "        location_b=np.r_[ip_data_dict[\"b_locations\"][i], 0.],\n",
    "        receiver_list=[receivers]\n",
    "    )\n",
    "    \n",
    "    # append the new source to the source list\n",
    "    source_list_ip.append(source)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "survey_ip = ip.Survey(source_list_ip)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ip_data = Data(\n",
    "    survey=survey, \n",
    "    dobs=np.hstack(ip_data_dict[\"observed_data\"]), \n",
    "    standard_deviation=np.hstack(ip_data_dict[\"standard_deviations\"])\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot psuedosection\n",
    "fig, ax = plt.subplots(1, 1, figsize=(12, 4))\n",
    "ip.utils.plot_pseudosection(\n",
    "    ip_data, data_type=\"apparent chargeability\",\n",
    "    plot_type=\"contourf\", data_location=True, ax=ax, \n",
    ")\n",
    "ax.set_title(f\"IP: {line} Pseudosection\")\n",
    "ax.set_aspect(1.5)  # some vertical exxageration"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 6: Define the IP simulation\n",
    "\n",
    "The general setup of an IP simulation is very similar to what we saw earlier for DC. Now, instead of inverting for resistivity (`rho`), we invert for chargeability (`eta`). The only additional parameter that is needed is the electrical resistivity (`rho`) from the DC inversion. This is used to construct the forward operator in the linear model of IP as described in the [IP Simulation](#IP-simulation) section above. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "simulation_ip = ip.simulation_2d.Simulation2DNodal(\n",
    "    mesh=mesh, solver=Solver, rho=np.exp(mopt), \n",
    "    etaMap=maps.IdentityMap(mesh), survey=survey_ip,\n",
    "    storeJ=True\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 7: IP inversion\n",
    "\n",
    "The setup for the IP inversion is very similar to what we did earlier for DC. The only change you will notice is that for the optimization, we are now using `ProjectedGNCG`: this is a Projected Gauss Newton with Conjugate Gradient (an iterative solver) optimization routine. It is similar to Inexact Gauss Newton but allows bound contstraints to be imposed. Here, we know that apparent chargeability is positive, so we use a lower bound of zero. \n",
    "\n",
    "We do not use an `ExpMap` as we did for resistivity in the DC inversion because chargeability tends to vary on a linear scale. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dmisfit_ip = data_misfit.L2DataMisfit(data=ip_data, simulation=simulation_ip)\n",
    "\n",
    "# Related to inversion\n",
    "reg_ip = regularization.Tikhonov(\n",
    "    mesh,\n",
    "    alpha_s=1/mesh.hx.min()**2.,\n",
    "    alpha_x=1.,\n",
    "    alpha_y=1.,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "opt_ip = optimization.ProjectedGNCG(\n",
    "    maxIter=20, maxIterCG=20, upper=np.inf, lower=0.\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "inv_prob_ip = inverse_problem.BaseInvProblem(dmisfit_ip, reg_ip, opt_ip)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Set a small starting model** - this is often assumed to be zero (no chargeable material). Here, since we will be using bound constraints in the optimization, we use a small value to reduce challenges of \"getting stuck\" on the lower bound of zero. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "m0_ip =  1e-3 * np.ones(mesh.nC)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "beta_est_ip = directives.BetaEstimate_ByEig(beta0_ratio=1e0)\n",
    "beta_schedule_ip = directives.BetaSchedule(coolingFactor=2, coolingRate=1)\n",
    "target_ip = directives.TargetMisfit()\n",
    "\n",
    "inv_ip = inversion.BaseInversion(\n",
    "    inv_prob_ip, directiveList=[beta_est_ip, beta_schedule_ip, target_ip]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run inversion\n",
    "charge_opt = inv_ip.run(m0_ip)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot misfit\n",
    "fig, ax = plt.subplots(1, 1, figsize=(12, 4))\n",
    "\n",
    "normalized_misfit = (inv_prob_ip.dpred - ip_data.dobs) / ip_data.standard_deviation\n",
    "\n",
    "dc.utils.plot_pseudosection(\n",
    "    ip_data, dobs=normalized_misfit, data_type=\"misfit\",\n",
    "    plot_type=\"contourf\", data_location=True, ax=ax, \n",
    "    cbar_opts={\"pad\":0.25}\n",
    ")\n",
    "\n",
    "ax.set_title(\"IP misfit\")\n",
    "ax.set_aspect(1.5)  # some vertical exxageration\n",
    "ax.set_xlabel(\"Northing (m)\")\n",
    "\n",
    "cb_axes = plt.gcf().get_axes()[-1]\n",
    "cb_axes.set_xlabel('normalized misfit')\n",
    "\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1,1, figsize=(10, 5))\n",
    "out = mesh.plotImage(\n",
    "    charge_opt, pcolorOpts={'cmap':'Spectral_r'}, ax=ax,\n",
    "    clim=(0, 30)\n",
    ")\n",
    "for data in geologic_section:\n",
    "    ax.plot(data[:,0], data[:,1], 'k--', alpha=0.5)\n",
    "ax.set_ylim((-400, 0))\n",
    "ax.set_xlim(core_domain_x)\n",
    "cb = plt.colorbar(out[0], label=\"chargeability\", fraction=0.05, orientation='horizontal', ax=ax, pad=0.2)\n",
    "ax.set_xlabel('Easting (m)')\n",
    "ax.set_ylabel('Elevation (m)')\n",
    "ax.set_aspect(1.5)\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "charge_ubc = read_ubc_model(\"century/46800E/IPMODA.CHG\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(2, 1, figsize=(10, 8))\n",
    "rho = (mapping*mopt)\n",
    "clim = (0, 30)\n",
    "\n",
    "out = mesh.plotImage(\n",
    "    charge_opt, pcolorOpts={'cmap':'Spectral_r'}, ax=ax[0],\n",
    "    clim=clim\n",
    ")\n",
    "cb = plt.colorbar(out[0], label=\"Chargeability\", fraction=0.05, orientation='horizontal', ax=ax[0], pad=0.2)\n",
    "ax[0].set_title(\"SimPEG\")\n",
    "\n",
    "out = mesh_ubc.plotImage(\n",
    "    charge_ubc, pcolorOpts={'cmap':'Spectral_r'}, ax=ax[1],\n",
    "    clim=clim\n",
    ")\n",
    "cb = plt.colorbar(out[0], label=\"Chargeability\", fraction=0.05, orientation='horizontal', ax=ax[1], pad=0.2)\n",
    "ax[1].set_title(\"UBC\")\n",
    "\n",
    "for a in ax: \n",
    "    for data in geologic_section:\n",
    "        a.plot(data[:,0], data[:,1], 'k--', alpha=0.5)\n",
    "        \n",
    "    a.set_xlim(core_domain_x)\n",
    "    a.set_ylim((-400, 0))\n",
    "    \n",
    "    a.set_aspect(1.5)\n",
    "    a.set_xlabel('Easting (m)')\n",
    "    a.set_ylabel('Elevation (m)')\n",
    "    \n",
    "plt.tight_layout()\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Homework ✏️\n",
    "\n",
    "There are 6 lines of data in total for the century deposit. Can you invert them all? \n",
    "\n",
    "Send us images of your results on Slack! We will share images of the UBC results in 1 week. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}