gravity-inversion.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Building a non-linear gravity inversion from scratch (almost)\n",
    "\n",
    "In this notebook, we'll build a non-linear gravity inversion to estimate the relief of a sedimentary basin. We'll implement smoothness regularization and see its effects on the solution. We'll also see how we can break the inversion by adding random noise, abusing regularization, and breaking the underlying assumptions.\n",
    "\n",
    "This notebook is part of the Geophysics Library lesson: [GeophysicsLibrary/non-linear-gravity-inversion](https://github.com/GeophysicsLibrary/non-linear-gravity-inversion) \n",
    "\n",
    "See the lesson material for extra code and setup instructions."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Imports\n",
    "\n",
    "We'll use the basic scientific Python stack for this tutorial plus a custom module with the forward modelling function (based on the code from the [Harmonica](https://github.com/fatiando/harmonica) library)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "# Our custom code (cheatcodes.py) with forward modelling and some utilities.\n",
    "import cheatcodes as cc "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is a little trick to make the resolution of the matplotlib figures better for larger screens."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.rc(\"figure\", dpi=120)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Assumptions\n",
    "\n",
    "Here are some assumptions we'll work with:\n",
    "\n",
    "1. The basin is much larger in the y-dimension so we'll assume it's infinite (reducing the problem to 2D)\n",
    "1. The gravity disturbance is entirely due to the sedimentary basin\n",
    "1. The top of the basin is a flat surface at $z=0$\n",
    "1. The data are measured at a constant height of $z=1\\ m$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Making synthetic data\n",
    "\n",
    "First, we'll explore the forward modelling function and create some synthetic data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "depths, basin_boundaries = cc.synthetic_model()\n",
    "\n",
    "print(basin_boundaries)\n",
    "print(depths)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plot the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cc.plot_prisms(depths, basin_boundaries)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Forward model some gravity data at a set of locations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = np.linspace(-5e3, 105e3, 60)\n",
    "density = -300  # kg/m³\n",
    "data = cc.forward_model(depths, basin_boundaries, density, x)\n",
    "\n",
    "plt.figure(figsize=(9, 3))\n",
    "plt.plot(x / 1000, data, \".k\")\n",
    "plt.xlabel(\"x [km]\")\n",
    "plt.ylabel(\"gravity disturbance [mGal]\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Calculating the Jacobian matrix\n",
    "\n",
    "The first step to most inverse problems is being able to calculate the Jacobian matrix. We'll do this for our problem by a first-order finite differences approximation. If you can get analytical derivatives, that's usually a lot better."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_jacobian(parameters, basin_boundaries, density, x):\n",
    "    \"\"\"\n",
    "    Calculate the Jacobian matrix by finite differences.\n",
    "    \"\"\"\n",
    "    jacobian = np.empty((x.size, parameters.size))\n",
    "    step = np.zeros_like(parameters)\n",
    "    delta = 10\n",
    "    for j in range(jacobian.shape[1]):\n",
    "        step[j] += delta\n",
    "        jacobian[:, j] = (\n",
    "            (\n",
    "                cc.forward_model(parameters + step, basin_boundaries, density, x)\n",
    "                - cc.forward_model(parameters, basin_boundaries, density, x)\n",
    "            ) \n",
    "            / delta\n",
    "        )\n",
    "        step[j] = 0\n",
    "    return jacobian"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Calculate and plot an example so we can see what this matrix looks like. We'll use a parameter vector with constant depths at first."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "parameters = np.zeros(30) + 5000\n",
    "\n",
    "jacobian = make_jacobian(parameters, basin_boundaries, density, x)\n",
    "\n",
    "plt.figure()\n",
    "plt.imshow(jacobian)\n",
    "plt.colorbar(label=\"mGal/m\")\n",
    "plt.xlabel(\"columns\")\n",
    "plt.ylabel(\"rows\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Solve the inverse problem \n",
    "\n",
    "Now that we have a way of forward modelling and calculating the Jacobian matrix, we can implement the Gauss-Newton method for solving the non-linear inverse problem. The function below takes the input data, model configuration, and an initial estimate and outputs the estimated parameters and a list with the goal function value per iteration. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def basin2d_inversion(x, data, basin_boundaries, density, initial, max_iterations=10):\n",
    "    \"\"\"\n",
    "    Solve the inverse problem using the Gauss-Newton method.\n",
    "    \"\"\"\n",
    "    parameters = initial.astype(np.float64).copy() \n",
    "    predicted = cc.forward_model(parameters, basin_boundaries, density, x)\n",
    "    residuals = data - predicted\n",
    "    goal_function = [np.linalg.norm(residuals)**2]\n",
    "    for i in range(max_iterations):       \n",
    "        jacobian = make_jacobian(parameters, basin_boundaries, density, x)\n",
    "        hessian = jacobian.T @ jacobian\n",
    "        gradient = jacobian.T @ residuals\n",
    "        deltap = np.linalg.solve(hessian, gradient)\n",
    "        new_parameters = parameters + deltap\n",
    "        predicted = cc.forward_model(new_parameters, basin_boundaries, density, x)\n",
    "        residuals = data - predicted\n",
    "        current_goal = np.linalg.norm(residuals)**2\n",
    "        if current_goal > goal_function[-1]:\n",
    "            break\n",
    "        parameters = new_parameters\n",
    "        goal_function.append(current_goal)\n",
    "    return parameters, goal_function"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can use this function to invert our synthetic data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "estimated, goal_function = basin2d_inversion(\n",
    "    x, data, basin_boundaries, density, initial=np.full(30, 1000),\n",
    ")\n",
    "predicted = cc.forward_model(estimated, basin_boundaries, density, x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plot the observed vs predicted data so we can inspect the fit."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(9, 3))\n",
    "plt.plot(x / 1e3, data, \".k\", label=\"observed\")\n",
    "plt.plot(x / 1e3, predicted, \"-r\", label='predicted')\n",
    "plt.legend()\n",
    "plt.xlabel(\"x [km]\")\n",
    "plt.ylabel(\"gravity disturbance [mGal]\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Look at the convergence of the method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure()\n",
    "plt.plot(goal_function)\n",
    "plt.yscale(\"log\")\n",
    "plt.xlabel(\"iteration\")\n",
    "plt.ylabel(\"goal function (mGal²)\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And finally see if our estimate is close to the true model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ax = cc.plot_prisms(depths, basin_boundaries)\n",
    "cc.plot_prisms(estimated, basin_boundaries, edgecolor=\"blue\", ax=ax)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Perfect! It seems that our inversion works well under these conditions (this initial estimate and no noise in the data). **Now let's break it!**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## **Your turn**\n",
    "\n",
    "**Add pseudo-random noise to the data using `np.random.normal` function and investigate the effect this has on the inversion results.** A typical gravity survey has accuracy in between 0.5-1 mGal. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Stability test\n",
    "\n",
    "We can go one step further and run several of these inversion in a loop for different random noise realizations. This will give us an idea of how stable the inversion is."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "estimates = []\n",
    "for i in range(5):\n",
    "    noise = np.random.normal(loc=0, scale=1, size=data.size)\n",
    "    noisy_data = data + noise\n",
    "    estimated, goal_function = basin2d_inversion(\n",
    "        x, noisy_data, basin_boundaries, density, initial=np.full(30, 1000),\n",
    "    )\n",
    "    estimates.append(estimated)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ax = cc.plot_prisms(depths, basin_boundaries)\n",
    "for estimated in estimates:\n",
    "    cc.plot_prisms(estimated, basin_boundaries, edgecolor=\"#0000ff66\", ax=ax)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## **Question time!**\n",
    "\n",
    "**Why does the inversion become more unstable for the deeper portions of the model?**\n",
    "\n",
    "Hint: It's related to the physics of the forward modelling and the Jacobian matrix."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Regularization to the rescue\n",
    "\n",
    "To deal with the instability issues we encountered, we will apply **first-order Tikhonov regularization** (aka \"smoothness\"). \n",
    "\n",
    "First thing we need to do is create the finite difference matrix $\\bar{\\bar{R}}$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def finite_difference_matrix(nparams):\n",
    "    \"\"\"\n",
    "    Create the finite difference matrix for regularization.\n",
    "    \"\"\"\n",
    "    fdmatrix = np.zeros((nparams - 1, nparams))\n",
    "    for i in range(fdmatrix.shape[0]):\n",
    "        fdmatrix[i, i] = -1\n",
    "        fdmatrix[i, i + 1] = 1\n",
    "    return fdmatrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "finite_difference_matrix(10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can use this to make a new inversion function with smoothness."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def basin2d_smooth_inversion(x, data, basin_boundaries, density, initial, smoothness, max_iterations=10):\n",
    "    \"\"\"\n",
    "    Solve the regularized inverse problem using the Gauss-Newton method.\n",
    "    \"\"\"\n",
    "    parameters = initial.astype(np.float64).copy() \n",
    "    predicted = cc.forward_model(parameters, basin_boundaries, density, x)\n",
    "    residuals = data - predicted\n",
    "    goal_function = [np.linalg.norm(residuals)**2]\n",
    "    fdmatrix = finite_difference_matrix(parameters.size)\n",
    "    for i in range(max_iterations):       \n",
    "        jacobian = make_jacobian(parameters, basin_boundaries, density, x)\n",
    "        hessian = jacobian.T @ jacobian + smoothness * fdmatrix.T @ fdmatrix\n",
    "        gradient = jacobian.T @ residuals - smoothness * fdmatrix.T @ fdmatrix @ parameters\n",
    "        deltap = np.linalg.solve(hessian, gradient)\n",
    "        new_parameters = parameters + deltap\n",
    "        predicted = cc.forward_model(new_parameters, basin_boundaries, density, x)\n",
    "        residuals = data - predicted\n",
    "        current_goal = np.linalg.norm(residuals)**2\n",
    "        if current_goal > goal_function[-1]:\n",
    "            break\n",
    "        parameters = new_parameters\n",
    "        goal_function.append(current_goal)\n",
    "    return parameters, goal_function"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we check if it works on our noisy data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "estimates = []\n",
    "for i in range(5):\n",
    "    noise = np.random.normal(loc=0, scale=1, size=data.size)\n",
    "    noisy_data = data + noise\n",
    "    estimated, goal_function = basin2d_smooth_inversion(\n",
    "        x, noisy_data, basin_boundaries, density, initial=np.full(30, 1000), smoothness=1e-5\n",
    "    )\n",
    "    estimates.append(estimated)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ax = cc.plot_prisms(depths, basin_boundaries)\n",
    "for estimated in estimates:\n",
    "    cc.plot_prisms(estimated, basin_boundaries, edgecolor=\"#0000ff66\", ax=ax)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## **Question time!**\n",
    "\n",
    "**What happens when the regularization paramater is extremely high?** Try to predict what the answer would be and then execute the code to check your reasoning.\n",
    "\n",
    "Hint: what is the smoothest possible model?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## **Your turn**\n",
    "\n",
    "**Can our regularized model recover a non-smooth geometry?** For example, real sedimentary basins often have [faults](https://en.wikipedia.org/wiki/Fault_(geology)) running through them, causing sharp jumps in the sediment thickness (up or down). \n",
    "\n",
    "To answer this question:\n",
    "\n",
    "1. Use the modified model depths below (the `depths` array) that introduce a shift up or down by 1-2 km in a section of the model of about 5-10 km.\n",
    "2. Generate new noisy data with this new model\n",
    "3. Invert the noisy data and try to find a model that:\n",
    "   1. Fits the data\n",
    "   2. Is stable (doesn't vary much if we change the noise)\n",
    "   3. Recovers the sharp boundary"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fault_model = np.copy(depths)\n",
    "fault_model[45:55] -= 2000"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cc.plot_prisms(fault_model, basin_boundaries)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now fill out the rest of the code below!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## **Question time!**\n",
    "\n",
    "**What would happen if we used a \"sharpness\" regularization?** Would we be able to recover the faults? What about the smoother parts of the model? \n",
    "\n",
    "One type of sharpness regularization is called \"total-variation regularization\" and it [has been used for this problem in the past](https://doi.org/10.1190/1.3524286)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Extra thinking points\n",
    "\n",
    "* What happens if we get the density wrong?\n",
    "* What are the sources of uncertainty in our final solution? Is it just the noise in the data?\n",
    "* How much does the solution depend on the inital estimate?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## **Bonus:** Optimizing code\n",
    "\n",
    "The code we wrote is not the greatest and it does take a while to run even for these really small 2D problems. There are ways in which we can make the code fast. But before we do any of that, **we need to know where our code spends most of its time**. Otherwise, we could spend hours optimizing a part of the code that is already really fast.\n",
    "\n",
    "This can be done with tools called **profilers**, which measure the time spent in each function of your code. This is also why its very important to **break up your code into functions**. In a Jupyter notebook, you can run the standard Python profiler by using the `%%prun` cell magic:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%prun \n",
    "basin2d_smooth_inversion(\n",
    "        x, noisy_data, basin_boundaries, density, initial=np.full(30, 1000), smoothness=1e-5\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `tottime` column is the amount of time spent on the function itself (not counting functions called inside it) and `cumtime` is the total time spent in the function, including function calls inside it. \n",
    "\n",
    "We can see from the profiling that the majority of the computation is spend in forward modelling, in particular for building the Jacobian. So if we can optimize `make_jacobian` that will have the biggest impact on performance of all.\n",
    "\n",
    "To start let's measure the computation time of `make_jacobian` with the `%%timeit` magic:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit\n",
    "make_jacobian(np.full(30, 1000), basin_boundaries, density, x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Alright, now we can try to do better.\n",
    "\n",
    "For many of these problems, the biggest return on investment is **not** parallelization or going to C/Fortran. **The largest improvements come from better maths/physics**. Here, we can take advantage of potential-field theory to cut down on the computation time of the Jacobian. \n",
    "\n",
    "We'll use the fact that the difference in gravity values produced by two models is the same as the gravity value produced by the difference in the models. Meaning that $\\delta g = g(m_1) - g(m_2) = g(m_1 - m_2)$. This way, we can reduce by more than half the number of forward modelling operations we do in the finite-difference computations.\n",
    "\n",
    "So instead of calculating the entire basin model with and without a small step in a single parameter, we can only calculate the effect of that small step."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_jacobian_fast(parameters, basin_boundaries, density, x):\n",
    "    \"\"\"\n",
    "    Calculate the Jacobian matrix by finite differences.\n",
    "    \"\"\"\n",
    "    jacobian = np.empty((x.size, parameters.size))\n",
    "    delta = 10\n",
    "    boundaries = cc.prism_boundaries(parameters, basin_boundaries)\n",
    "    for j in range(jacobian.shape[1]):\n",
    "        jacobian[:, j] = (\n",
    "            (\n",
    "                # Replace with a single forward modelling of a single prism\n",
    "                cc.prism_gravity(x, boundaries[j], boundaries[j + 1], parameters[j], parameters[j] + delta, density)\n",
    "            ) \n",
    "            / delta\n",
    "        )\n",
    "    return jacobian"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, we check if the results are still correct."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.allclose(\n",
    "    make_jacobian(np.full(30, 1000), basin_boundaries, density, x),\n",
    "    make_jacobian_fast(np.full(30, 1000), basin_boundaries, density, x)\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can measure the time again:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit\n",
    "make_jacobian_fast(np.full(30, 1000), basin_boundaries, density, x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This one change gave use 2 orders of magnitude improvement in the function that makes up most of the computation time. **Now that is time well spent!**\n",
    "\n",
    "We can measure how much of a difference this makes for the inversion as a whole by making a new function with our fast Jacobian matrix calculation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fast_basin2d_smooth_inversion(x, data, basin_boundaries, density, initial, smoothness, max_iterations=10):\n",
    "    \"\"\"\n",
    "    Solve the regularized inverse problem using the Gauss-Newton method.\n",
    "    \"\"\"\n",
    "    parameters = initial.astype(np.float64).copy() \n",
    "    predicted = cc.forward_model(parameters, basin_boundaries, density, x)\n",
    "    residuals = data - predicted\n",
    "    goal_function = [np.linalg.norm(residuals)**2]\n",
    "    fdmatrix = finite_difference_matrix(parameters.size)\n",
    "    for i in range(max_iterations):       \n",
    "        # Swap out the slow jacobian for the fast one\n",
    "        jacobian = make_jacobian_fast(parameters, basin_boundaries, density, x)\n",
    "        hessian = jacobian.T @ jacobian + smoothness * fdmatrix.T @ fdmatrix\n",
    "        gradient = jacobian.T @ residuals - smoothness * fdmatrix.T @ fdmatrix @ parameters\n",
    "        deltap = np.linalg.solve(hessian, gradient)\n",
    "        new_parameters = parameters + deltap\n",
    "        predicted = cc.forward_model(new_parameters, basin_boundaries, density, x)\n",
    "        residuals = data - predicted\n",
    "        current_goal = np.linalg.norm(residuals)**2\n",
    "        if current_goal > goal_function[-1]:\n",
    "            break\n",
    "        parameters = new_parameters\n",
    "        goal_function.append(current_goal)\n",
    "    return parameters, goal_function"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And now we can measure the computation time for both."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit \n",
    "basin2d_smooth_inversion(\n",
    "        x, noisy_data, basin_boundaries, density, initial=np.full(30, 1000), smoothness=1e-5\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit \n",
    "fast_basin2d_smooth_inversion(\n",
    "        x, noisy_data, basin_boundaries, density, initial=np.full(30, 1000), smoothness=1e-5\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**We changed 3 lines of code and achived a factor of 10 speedup.** Again, this could only be done because we first profiled the code and then focused on finding a fundamentally better way of calculating. "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:non-linear-gravity-inversion]",
   "language": "python",
   "name": "conda-env-non-linear-gravity-inversion-py"
  },
  "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.9.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}