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)
(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.69982545 1364.20386762 1324.23908356 1299.37606027 1286.37970806 1281.36048532 1280.00436832 1277.86656035 1270.70261956 1254.80221173 1227.28696229 1186.33602894 1131.31108729 1062.765288 982.33628338 892.53894688 796.4862351 697.57464821 599.17280554 504.34781951 415.65556423 335.00950302 263.63067355 202.0707732 150.29253113 107.78743462 73.71036175 47.0131077 26.56315486 11.23920733 0. ]
Plot the model.
cc.plot_prisms(depths, basin_boundaries)
<AxesSubplot:xlabel='x [km]', ylabel='depth [km]'>
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)
<AxesSubplot:xlabel='x [km]', ylabel='depth [km]'>
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.
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)
<AxesSubplot:xlabel='x [km]', ylabel='depth [km]'>
We can go one step further and run several of these inversion in a loop for different random noise realizations.
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)
The fact that small changes in the data lead to large changes in the model is the defining characteristic of an unstable problem.
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)
array([[-1., 1., 0., 0., 0., 0., 0., 0., 0., 0.], [ 0., -1., 1., 0., 0., 0., 0., 0., 0., 0.], [ 0., 0., -1., 1., 0., 0., 0., 0., 0., 0.], [ 0., 0., 0., -1., 1., 0., 0., 0., 0., 0.], [ 0., 0., 0., 0., -1., 1., 0., 0., 0., 0.], [ 0., 0., 0., 0., 0., -1., 1., 0., 0., 0.], [ 0., 0., 0., 0., 0., 0., -1., 1., 0., 0.], [ 0., 0., 0., 0., 0., 0., 0., -1., 1., 0.], [ 0., 0., 0., 0., 0., 0., 0., 0., -1., 1.]])
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).
fault_model = np.copy(depths)
fault_model[45:55] -= 2000
cc.plot_prisms(fault_model, basin_boundaries)
<AxesSubplot:xlabel='x [km]', ylabel='depth [km]'>
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)
(5.5, 0.0)
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
)
193062 function calls (189976 primitive calls) in 4.142 seconds Ordered by: internal time ncalls tottime percall cumtime percall filename:lineno(function) 146640 3.076 0.000 3.076 0.000 cheatcodes.py:132(kernel) 18330 0.957 0.000 4.033 0.000 cheatcodes.py:98(prism_gravity) 611 0.035 0.000 4.132 0.007 cheatcodes.py:84(forward_model) 611 0.017 0.000 0.048 0.000 function_base.py:23(linspace) 4340/1254 0.010 0.000 0.062 0.000 {built-in method numpy.core._multiarray_umath.implement_array_function} 10 0.006 0.001 4.011 0.401 <ipython-input-6-05b820af7edb>:1(make_jacobian) 611 0.005 0.000 0.005 0.000 {method 'reduce' of 'numpy.ufunc' objects} 1254 0.003 0.000 0.003 0.000 {built-in method numpy.array} 611 0.003 0.000 0.012 0.000 fromnumeric.py:70(_wrapreduction) 611 0.002 0.000 0.002 0.000 {built-in method numpy.arange} 611 0.002 0.000 0.053 0.000 cheatcodes.py:75(prism_boundaries) 621 0.002 0.000 0.009 0.000 numeric.py:75(zeros_like) 611 0.002 0.000 0.008 0.000 {method 'any' of 'numpy.generic' objects} 622 0.002 0.000 0.002 0.000 {built-in method numpy.zeros} 611 0.001 0.000 0.013 0.000 fromnumeric.py:2249(any) 10 0.001 0.000 0.002 0.000 linalg.py:314(solve) 611 0.001 0.000 0.001 0.000 {method 'reshape' of 'numpy.ndarray' objects} 621 0.001 0.000 0.003 0.000 <__array_function__ internals>:2(empty_like) 622 0.001 0.000 0.001 0.000 {method 'astype' of 'numpy.ndarray' objects} 611 0.001 0.000 0.051 0.000 <__array_function__ internals>:2(linspace) 611 0.001 0.000 0.004 0.000 <__array_function__ internals>:2(result_type) 1 0.001 0.001 4.141 4.141 <ipython-input-20-ea2ff521d745>:1(basin2d_smooth_inversion) 622 0.001 0.000 0.003 0.000 <__array_function__ internals>:2(copyto) 611 0.001 0.000 0.002 0.000 <__array_function__ internals>:2(ndim) 611 0.001 0.000 0.015 0.000 <__array_function__ internals>:2(any) 1222 0.001 0.000 0.004 0.000 _asarray.py:86(asanyarray) 621 0.001 0.000 0.011 0.000 <__array_function__ internals>:2(zeros_like) 611 0.001 0.000 0.001 0.000 fromnumeric.py:71(<dictcomp>) 611 0.001 0.000 0.001 0.000 fromnumeric.py:3075(ndim) 611 0.000 0.000 0.001 0.000 numeric.py:1816(isscalar) 611 0.000 0.000 0.006 0.000 _methods.py:53(_any) 631 0.000 0.000 0.000 0.000 {built-in method builtins.getattr} 611 0.000 0.000 0.000 0.000 {built-in method builtins.isinstance} 611 0.000 0.000 0.000 0.000 {method 'items' of 'dict' objects} 611 0.000 0.000 0.000 0.000 {built-in method _operator.index} 1 0.000 0.000 0.000 0.000 {method 'copy' of 'numpy.ndarray' objects} 611 0.000 0.000 0.000 0.000 function_base.py:18(_linspace_dispatcher) 622 0.000 0.000 0.000 0.000 multiarray.py:1043(copyto) 611 0.000 0.000 0.000 0.000 fromnumeric.py:3071(_ndim_dispatcher) 621 0.000 0.000 0.000 0.000 numeric.py:71(_zeros_like_dispatcher) 611 0.000 0.000 0.000 0.000 fromnumeric.py:2245(_any_dispatcher) 11 0.000 0.000 0.000 0.000 linalg.py:2363(norm) 621 0.000 0.000 0.000 0.000 multiarray.py:75(empty_like) 611 0.000 0.000 0.000 0.000 multiarray.py:634(result_type) 10 0.000 0.000 0.000 0.000 linalg.py:135(_commonType) 1 0.000 0.000 4.142 4.142 {built-in method builtins.exec} 20 0.000 0.000 0.000 0.000 linalg.py:107(_makearray) 11 0.000 0.000 0.000 0.000 <__array_function__ internals>:2(norm) 11 0.000 0.000 0.000 0.000 {built-in method numpy.empty} 10 0.000 0.000 0.002 0.000 <__array_function__ internals>:2(solve) 11 0.000 0.000 0.000 0.000 {method 'ravel' of 'numpy.ndarray' objects} 11 0.000 0.000 0.000 0.000 <__array_function__ internals>:2(dot) 31 0.000 0.000 0.000 0.000 _asarray.py:14(asarray) 41 0.000 0.000 0.000 0.000 linalg.py:112(isComplexType) 72 0.000 0.000 0.000 0.000 {built-in method builtins.issubclass} 20 0.000 0.000 0.000 0.000 linalg.py:125(_realType) 10 0.000 0.000 0.000 0.000 linalg.py:102(get_linalg_error_extobj) 10 0.000 0.000 0.000 0.000 linalg.py:200(_assert_stacked_square) 10 0.000 0.000 0.000 0.000 linalg.py:194(_assert_stacked_2d) 1 0.000 0.000 0.000 0.000 <ipython-input-18-0b2fd295018c>:1(finite_difference_matrix) 1 0.000 0.000 4.141 4.141 <string>:1(<module>) 1 0.000 0.000 0.000 0.000 numeric.py:268(full) 20 0.000 0.000 0.000 0.000 {method 'get' of 'dict' objects} 10 0.000 0.000 0.000 0.000 {method '__array_prepare__' of 'numpy.ndarray' objects} 11 0.000 0.000 0.000 0.000 linalg.py:2359(_norm_dispatcher) 11 0.000 0.000 0.000 0.000 multiarray.py:706(dot) 10 0.000 0.000 0.000 0.000 linalg.py:310(_solve_dispatcher) 10 0.000 0.000 0.000 0.000 {method 'append' of 'list' objects} 1 0.000 0.000 0.000 0.000 {method 'disable' of '_lsprof.Profiler' objects}
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)
319 ms ± 69.2 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
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)
)
True
Now we can measure the time again:
%%timeit
make_jacobian_fast(np.full(30, 1000), basin_boundaries, density, x)
5.07 ms ± 739 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
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
)
3.73 s ± 55.9 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%%timeit
fast_basin2d_smooth_inversion(
x, noisy_data, basin_boundaries, density, initial=np.full(30, 1000), smoothness=1e-5
)
302 ms ± 4.95 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
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.