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.
import numpy as np
import matplotlib.pyplot as plt
import cheatcodes as cc
This is a little trick to make the resolution of the matplotlib figures better for larger screens.
plt.rc("figure", dpi=120)
Here are some assumptions we'll work with:
First, we'll explore the forward modelling function and create some synthetic data.
depths, basin_boundaries = cc.synthetic_model()
print(basin_boundaries)
print(depths)
Plot the model.
cc.plot_prisms(depths, basin_boundaries)
Forward model some gravity data at a set of locations.
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()
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.
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
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.
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()
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.
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
Now we can use this function to invert our synthetic data.
estimated, goal_function = basin2d_inversion(
x, data, basin_boundaries, density, initial=np.full(30, 1000),
)
predicted = cc.forward_model(estimated, basin_boundaries, density, x)
Plot the observed vs predicted data so we can inspect the fit.
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()
Look at the convergence of the method.
plt.figure()
plt.plot(goal_function)
plt.yscale("log")
plt.xlabel("iteration")
plt.ylabel("goal function (mGal²)")
plt.show()
And finally see if our estimate is close to the true model.
ax = cc.plot_prisms(depths, basin_boundaries)
cc.plot_prisms(estimated, basin_boundaries, edgecolor="blue", ax=ax)
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!
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.
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.
To 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}}$.
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)
Now we can use this to make a new inversion function with smoothness.
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
Now we check if it works on our noisy data.
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)
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?
Can our regularized model recover a non-smooth geometry? For example, real sedimentary basins often have faults running through them, causing sharp jumps in the sediment thickness (up or down).
To answer this question:
depths
array) to introduce a shift up or down by 1-2 km in a section of the model of about 5-10 km.Hint: Use np.copy
to make a copy of the depths
(to avoid overwriting it).
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.
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.
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:
%%prun
basin2d_smooth_inversion(
x, noisy_data, basin_boundaries, density, initial=np.full(30, 1000), smoothness=1e-5
)
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:
%%timeit
make_jacobian(np.full(30, 1000), basin_boundaries, density, x)
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.
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
First, we check if the results are still correct.
np.allclose(
make_jacobian(np.full(30, 1000), basin_boundaries, density, x),
make_jacobian_fast(np.full(30, 1000), basin_boundaries, density, x)
)
Now we can measure the time again:
%%timeit
make_jacobian_fast(np.full(30, 1000), basin_boundaries, density, x)
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.
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
And now we can measure the computation time for both.
%%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
)
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.