srcs, recs = utils.connect_points(src_loc, rec_loc)
tts = ttime2d.straight(model, 'vp', srcs, recs)
tts, error = utils.contaminate(tts, 0.02, percent=True, return_stddev=True,
                               seed=seed)
# Make the mesh
mesh = SquareMesh(area, shape)
# and run the inversion
misfit = srtomo.SRTomo(tts, srcs, recs, mesh)
regularization = TotalVariation2D(10**-10, mesh.shape)
tomo = misfit + regularization
tomo = LCurve(misfit, regularization,
              [10**i for i in np.arange(-3, 3, 0.5)], jobs=8)
# Since Total Variation is a non-linear function, then the tomography becomes
# non-linear. So we need to configure fit to use the Levemberg-Marquardt
# algorithm, a gradient descent method, that requires an initial estimate
tomo.config('levmarq', initial=0.00001*np.ones(mesh.size)).fit()
mesh.addprop('vp', tomo.estimate_)

# Plot the L-curve annd print the regularization parameter estimated
mpl.figure()
mpl.title('L-curve: triangle marks the best solution')
tomo.plot_lcurve()
print "Estimated regularization parameter: %g" % (tomo.regul_param_)

# Calculate and print the standard deviation of the residuals
# Should be close to the data error if the inversion was able to fit the data
residuals = tomo.residuals()
print "Assumed error: %f" % (error)
print "Standard deviation of residuals: %f" % (np.std(residuals))

mpl.figure(figsize=(14, 5))
depths = (-1e-15*(xs - 50000)**4 + 8000 -
          3000*np.exp(-(xs - 70000)**2/(10000**2)))
depths -= depths.min()  # Reduce depths to zero
props = {'density': -300}
model = Polygon(np.transpose([xs, depths]), props)
x = np.linspace(0, 100000, 100)
z = -100*np.ones_like(x)
data = utils.contaminate(talwani.gz(x, z, [model]), 0.5, seed=0)

# Make the solver and run the inversion
misfit = PolygonalBasinGravity(x, z, data, 50, props, top=0)
regul = Smoothness1D(misfit.nparams)
# Use an L-curve analysis to find the best regularization parameter
lc = LCurve(misfit, regul, [10**i for i in np.arange(-10, -5, 0.5)], jobs=4)
initial = 3000*np.ones(misfit.nparams)
lc.config('levmarq', initial=initial).fit()

mpl.figure()
mpl.subplot(2, 2, 1)
mpl.plot(x, data, 'ok', label='observed')
mpl.plot(x, lc.predicted(), '-r', linewidth=2, label='predicted')
mpl.legend()
ax = mpl.subplot(2, 2, 3)
mpl.polygon(model, fill='gray', alpha=0.5)
mpl.polygon(lc.estimate_, style='o-r')
ax.invert_yaxis()
mpl.subplot(1, 2, 2)
mpl.title('L-curve')
lc.plot_lcurve()
mpl.show()