<jupyter_start><jupyter_text>Building a non-linear gravity inversion from scratch (almost)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.  ImportsWe'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).<jupyter_code>import numpy as np
import matplotlib.pyplot as plt
import cheatcodes as cc<jupyter_output><empty_output><jupyter_text>This is a little trick to make the resolution of the matplotlib figures better for larger screens.<jupyter_code>plt.rc("figure", dpi=120)<jupyter_output><empty_output><jupyter_text>AssumptionsHere are some assumptions we'll work with:1. The basin is much larger in the y-dimension so we'll assume it's infinite (reducing the problem to 2D)1. The gravity disturbance is entirely due to the sedimentary basin1. The top of the basin is a flat surface at $z=0$1. The data are measured at a constant height of $z=1\ m$  Making synthetic dataFirst, we'll explore the forward modelling function and create some synthetic data.<jupyter_code>depths, basin_boundaries = cc.synthetic_model()

print(basin_boundaries)
print(depths)<jupyter_output>(0, 100000.0)
[  10.10105917   18.81092435   29.38998081   42.16839388   57.51774684
   75.85246428   97.63037844  123.35222519  153.55985009  188.83290654
  229.78383989  277.05097454  331.28955675  393.1606554   463.31788573
  542.39199706  630.97345373  729.59323554  838.70218841  958.64936255
 1089.65988011 1231.81297121 1385.02090173 1549.00958058 1723.30167502
 1907.20307243 2099.79350393 2299.92208393 2506.20842163 2717.04982321
 2930.63493217 3144.96395261 3357.87537423 3567.07887609 3770.19383914
 3964.79265602 4148.44780248 4318.78143888 4473.51615424 4610.52535757
 4727.88177002 4823.90248153 4897.18910944 4946.66173395 4971.58548322
 4971.58889307 4946.67346668 4897.21419639 4823.95117457 4727.97279851
 4610.69145219 4473.81290687 4319.30099926 4149.33939015 3966.2923388
 3772.66640115 3571.07472151 3364.2050895  3154.79214739 2945.59306685
 2739.36473798 2538.83917762 2346.69270631 2165.50377716 1997.6945141
 1845.45234455 1710.63077305 1594.63228888 1498.28126903 1421.699[...]<jupyter_text>Plot the model.<jupyter_code>cc.plot_prisms(depths, basin_boundaries)<jupyter_output><empty_output><jupyter_text>Forward model some gravity data at a set of locations.<jupyter_code>x = np.linspace(-5e3, 105e3, 60)
density = -300  # kg/m³
data = cc.forward_model(depths, basin_boundaries, density, x)

plt.figure(figsize=(9, 3))
plt.plot(x / 1000, data, ".k")
plt.xlabel("x [km]")
plt.ylabel("gravity disturbance [mGal]")
plt.show()<jupyter_output><empty_output><jupyter_text>Calculating the Jacobian matrixThe 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.<jupyter_code>def make_jacobian(parameters, basin_boundaries, density, x):
    """
    Calculate the Jacobian matrix by finite differences.
    """
    jacobian = np.empty((x.size, parameters.size))
    step = np.zeros_like(parameters)
    delta = 10
    for j in range(jacobian.shape[1]):
        step[j] += delta
        jacobian[:, j] = (
            (
                cc.forward_model(parameters + step, basin_boundaries, density, x)
                - cc.forward_model(parameters, basin_boundaries, density, x)
            ) 
            / delta
        )
        step[j] = 0
    return jacobian<jupyter_output><empty_output><jupyter_text>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.<jupyter_code>parameters = np.zeros(30) + 5000

jacobian = make_jacobian(parameters, basin_boundaries, density, x)

plt.figure()
plt.imshow(jacobian)
plt.colorbar(label="mGal/m")
plt.xlabel("columns")
plt.ylabel("rows")
plt.show()<jupyter_output><empty_output><jupyter_text>Solve the inverse problem 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.<jupyter_code>def basin2d_inversion(x, data, basin_boundaries, density, initial, max_iterations=10):
    """
    Solve the inverse problem using the Gauss-Newton method.
    """
    parameters = initial.astype(np.float64).copy() 
    predicted = cc.forward_model(parameters, basin_boundaries, density, x)
    residuals = data - predicted
    goal_function = [np.linalg.norm(residuals)**2]
    for i in range(max_iterations):       
        jacobian = make_jacobian(parameters, basin_boundaries, density, x)
        hessian = jacobian.T @ jacobian
        gradient = jacobian.T @ residuals
        deltap = np.linalg.solve(hessian, gradient)
        new_parameters = parameters + deltap
        predicted = cc.forward_model(new_parameters, basin_boundaries, density, x)
        residuals = data - predicted
        current_goal = np.linalg.norm(residuals)**2
        if current_goal > goal_function[-1]:
            break
        parameters = new_parameters
        goal_function.append(current_goal)
    return parameters, goal_function<jupyter_output><empty_output><jupyter_text>Now we can use this function to invert our synthetic data.<jupyter_code>estimated, goal_function = basin2d_inversion(
    x, data, basin_boundaries, density, initial=np.full(30, 1000),
)
predicted = cc.forward_model(estimated, basin_boundaries, density, x)<jupyter_output><empty_output><jupyter_text>Plot the observed vs predicted data so we can inspect the fit.<jupyter_code>plt.figure(figsize=(9, 3))
plt.plot(x / 1e3, data, ".k", label="observed")
plt.plot(x / 1e3, predicted, "-r", label='predicted')
plt.legend()
plt.xlabel("x [km]")
plt.ylabel("gravity disturbance [mGal]")
plt.show()<jupyter_output><empty_output><jupyter_text>Look at the convergence of the method.<jupyter_code>plt.figure()
plt.plot(goal_function)
plt.yscale("log")
plt.xlabel("iteration")
plt.ylabel("goal function (mGal²)")
plt.show()<jupyter_output><empty_output><jupyter_text>And finally see if our estimate is close to the true model.<jupyter_code>ax = cc.plot_prisms(depths, basin_boundaries)
cc.plot_prisms(estimated, basin_boundaries, edgecolor="blue", ax=ax)<jupyter_output><empty_output><jupyter_text>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!**  **Flex your coding muscles****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.<jupyter_code>noise = np.random.normal(loc=0, scale=1, size=data.size)
noisy_data = data + noise
estimated, goal_function = basin2d_inversion(
    x, noisy_data, basin_boundaries, density, initial=np.full(30, 1000),
)
predicted = cc.forward_model(estimated, basin_boundaries, density, x)
plt.figure(figsize=(9, 3))
plt.plot(x / 1e3, noisy_data, ".k", label="observed")
plt.plot(x / 1e3, predicted, "-r", label='predicted')
plt.legend()
plt.xlabel("x [km]")
plt.ylabel("gravity disturbance [mGal]")
plt.show()

ax = cc.plot_prisms(depths, basin_boundaries)
cc.plot_prisms(estimated, basin_boundaries, edgecolor="blue", ax=ax)<jupyter_output><empty_output><jupyter_text>We can go one step further and run several of these inversion in a loop for different random noise realizations.<jupyter_code>estimates = []
for i in range(5):
    noise = np.random.normal(loc=0, scale=1, size=data.size)
    noisy_data = data + noise
    estimated, goal_function = basin2d_inversion(
        x, noisy_data, basin_boundaries, density, initial=np.full(30, 1000),
    )
    estimates.append(estimated)
ax = cc.plot_prisms(depths, basin_boundaries)
for estimated in estimates:
    cc.plot_prisms(estimated, basin_boundaries, edgecolor="#0000ff66", ax=ax)<jupyter_output><empty_output><jupyter_text>The fact that **small changes in the data lead to large changes in the model** is the defining characteristic of an **unstable problem**.  **Question time!****Why does the inversion struggle to recover the deeper portions of the model in particular?**Hint: It's related to the physics of the forward modelling and the Jacobian matrix.  Regularization to the rescueTo deal with the instability issues we encountered, we will apply **first-order Tikhonov regularization** (aka "smoothness"). First thing we need to do is create the finite difference matrix $\bar{\bar{R}}$.<jupyter_code>def finite_difference_matrix(nparams):
    """
    Create the finite difference matrix for regularization.
    """
    fdmatrix = np.zeros((nparams - 1, nparams))
    for i in range(fdmatrix.shape[0]):
        fdmatrix[i, i] = -1
        fdmatrix[i, i + 1] = 1
    return fdmatrix
finite_difference_matrix(10)<jupyter_output><empty_output><jupyter_text>Now we can use this to make a new inversion function with smoothness.<jupyter_code>def basin2d_smooth_inversion(x, data, basin_boundaries, density, initial, smoothness, max_iterations=10):
    """
    Solve the regularized inverse problem using the Gauss-Newton method.
    """
    parameters = initial.astype(np.float64).copy() 
    predicted = cc.forward_model(parameters, basin_boundaries, density, x)
    residuals = data - predicted
    goal_function = [np.linalg.norm(residuals)**2]
    fdmatrix = finite_difference_matrix(parameters.size)
    for i in range(max_iterations):       
        jacobian = make_jacobian(parameters, basin_boundaries, density, x)
        hessian = jacobian.T @ jacobian + smoothness * fdmatrix.T @ fdmatrix
        gradient = jacobian.T @ residuals - smoothness * fdmatrix.T @ fdmatrix @ parameters
        deltap = np.linalg.solve(hessian, gradient)
        new_parameters = parameters + deltap
        predicted = cc.forward_model(new_parameters, basin_boundaries, density, x)
        residuals = data - predicted
        current_goal = np.linalg.norm(residuals)**2
        if current_goal > goal_function[-1]:
            break
        parameters = new_parameters
        goal_function.append(current_goal)
    return parameters, goal_function<jupyter_output><empty_output><jupyter_text>Now we check if it works on our noisy data.<jupyter_code>estimates = []
for i in range(5):
    noise = np.random.normal(loc=0, scale=1, size=data.size)
    noisy_data = data + noise
    estimated, goal_function = basin2d_smooth_inversion(
        x, noisy_data, basin_boundaries, density, initial=np.full(30, 1000), smoothness=1e-5
    )
    estimates.append(estimated)
ax = cc.plot_prisms(depths, basin_boundaries)
for estimated in estimates:
    cc.plot_prisms(estimated, basin_boundaries, edgecolor="#0000ff66", ax=ax)<jupyter_output><empty_output><jupyter_text>**Question time!****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.Hint: what is the smoothest possible model?  **Flex your coding muscles****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). To answer this question:1. Modify our model depths (the `depths` array) to introduce a shift up or down by 1-2 km in a section of the model of about 5-10 km.2. Generate new noisy data with this new model3. Invert the noisy data and try to find a model that:   1. Fits the data   2. Is stable (doesn't vary much if we change the noise)   3. Recovers the sharp boundary   Hint: Use `np.copy` to make a copy of the `depths` (to avoid overwriting it).<jupyter_code>fault_model = np.copy(depths)
fault_model[45:55] -= 2000
cc.plot_prisms(fault_model, basin_boundaries)
fault_data = cc.forward_model(fault_model, basin_boundaries, density, x)
estimates = []
for i in range(5):
    noise = np.random.normal(loc=0, scale=1, size=data.size)
    noisy_data = fault_data + noise
    estimated, goal_function = basin2d_smooth_inversion(
        x, noisy_data, basin_boundaries, density, initial=np.full(30, 1000), smoothness=1e-5
    )
    estimates.append(estimated)
ax = cc.plot_prisms(fault_model, basin_boundaries)
for estimated in estimates:
    cc.plot_prisms(estimated, basin_boundaries, edgecolor="#0000ff66", ax=ax)
ax.set_ylim(5.5, 0)<jupyter_output><empty_output><jupyter_text>**Question time!****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? 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).  Extra thinking points* What happens if we get the density wrong?* What are the sources of uncertainty in our final solution? Is it just the noise in the data?* How much does the solution depend on the inital estimate?  **Bonus:** Optimizing codeThe 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.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:<jupyter_code>%%prun 
basin2d_smooth_inversion(
        x, noisy_data, basin_boundaries, density, initial=np.full(30, 1000), smoothness=1e-5
)<jupyter_output><jupyter_text>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. 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.To start let's measure the computation time of `make_jacobian` with the `%%timeit` magic:<jupyter_code>%%timeit
make_jacobian(np.full(30, 1000), basin_boundaries, density, x)<jupyter_output>319 ms ± 69.2 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)<jupyter_text>Alright, now we can try to do better.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. 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.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.<jupyter_code>def make_jacobian_fast(parameters, basin_boundaries, density, x):
    """
    Calculate the Jacobian matrix by finite differences.
    """
    jacobian = np.empty((x.size, parameters.size))
    delta = 10
    boundaries = cc.prism_boundaries(parameters, basin_boundaries)
    for j in range(jacobian.shape[1]):
        jacobian[:, j] = (
            (
                # Replace with a single forward modelling of a single prism
                cc.prism_gravity(x, boundaries[j], boundaries[j + 1], parameters[j], parameters[j] + delta, density)
            ) 
            / delta
        )
    return jacobian<jupyter_output><empty_output><jupyter_text>First, we check if the results are still correct.<jupyter_code>np.allclose(
    make_jacobian(np.full(30, 1000), basin_boundaries, density, x),
    make_jacobian_fast(np.full(30, 1000), basin_boundaries, density, x)
)<jupyter_output><empty_output><jupyter_text>Now we can measure the time again:<jupyter_code>%%timeit
make_jacobian_fast(np.full(30, 1000), basin_boundaries, density, x)<jupyter_output>5.07 ms ± 739 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)<jupyter_text>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!**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.<jupyter_code>def fast_basin2d_smooth_inversion(x, data, basin_boundaries, density, initial, smoothness, max_iterations=10):
    """
    Solve the regularized inverse problem using the Gauss-Newton method.
    """
    parameters = initial.astype(np.float64).copy() 
    predicted = cc.forward_model(parameters, basin_boundaries, density, x)
    residuals = data - predicted
    goal_function = [np.linalg.norm(residuals)**2]
    fdmatrix = finite_difference_matrix(parameters.size)
    for i in range(max_iterations):       
        # Swap out the slow jacobian for the fast one
        jacobian = make_jacobian_fast(parameters, basin_boundaries, density, x)
        hessian = jacobian.T @ jacobian + smoothness * fdmatrix.T @ fdmatrix
        gradient = jacobian.T @ residuals - smoothness * fdmatrix.T @ fdmatrix @ parameters
        deltap = np.linalg.solve(hessian, gradient)
        new_parameters = parameters + deltap
        predicted = cc.forward_model(new_parameters, basin_boundaries, density, x)
        residuals = data - predicted
        current_goal = np.linalg.norm(residuals)**2
        if current_goal > goal_function[-1]:
            break
        parameters = new_parameters
        goal_function.append(current_goal)
    return parameters, goal_function<jupyter_output><empty_output><jupyter_text>And now we can measure the computation time for both.<jupyter_code>%%timeit 
basin2d_smooth_inversion(
        x, noisy_data, basin_boundaries, density, initial=np.full(30, 1000), smoothness=1e-5
)
%%timeit 
fast_basin2d_smooth_inversion(
        x, noisy_data, basin_boundaries, density, initial=np.full(30, 1000), smoothness=1e-5
)<jupyter_output>302 ms ± 4.95 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
{'fraction_non_alphanumeric': 0.09028469564305563, 'fraction_numerical': 0.056939128611125676, 'mean_word_length': 4.577192982456141, 'pattern_counts': {'":': 0, '<': 96, '<?xml version=': 0, '>': 99, 'https://': 3, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '25', 'dir': 'jupyter-structured-clean-dedup', 'id': '73873', 'license': 'MIT', 'n_tokens_mistral': 6779, 'n_tokens_neox': 5809, 'n_words': 1923, 'path': 'gravity-inversion-solution.ipynb', 'repo_name': 'tobi-ore/2020-aachen-inverse-problems', 'size': '19073'}