from IPython.display import Image
Image(filename='sketch.png')

import numpy as np
from fatiando.gravmag import talwani
from fatiando.mesher import Polygon
from fatiando.inversion.base import Misfit

class PolyBasin2D(Misfit):
    def __init__(self, x, z, gz, height, npoints, density):
        super(PolyBasin2D, self).__init__(
            data=gz, nparams=npoints, islinear=False,
            positional=dict(x=x, z=z), 
            model=dict(height=height, 
                       props={'density': density}))
    
    def p2vertices(self, p):
        x, h = self.positional['x'], self.model['height']
        n = self.nparams + 2
        verts = np.empty((n, 2))
        verts[:, 0] = np.linspace(x.min(), x.max(), n)[::-1]
        verts[:, 1] = np.concatenate([[h], p, [h]])
        return verts
    
    def _get_predicted(self, p):
        x, z = self.positional['x'], self.positional['z']
        verts = self.p2vertices(p)
        poly = Polygon(verts, self.model['props'])
        return talwani.gz(x, z, [poly])
        
    def _get_jacobian(self, p):    
        x, z = self.positional['x'], self.positional['z']
        verts = self.p2vertices(p)
        delta = np.array([0, 10])
        jac = np.zeros((self.ndata, self.nparams))
        for i in xrange(self.nparams):
            diff = Polygon([verts[i + 2], verts[i + 1] - delta, 
                            verts[i], verts[i + 1] + delta],
                           self.model['props'])
            jac[:, i] = talwani.gz(x, z, [diff])/(2 * delta[1])
        return jac
    
    def fit(self):
        super(PolyBasin2D, self).fit()
        self._estimate = Polygon(self.p2vertices(self.p_), 
                                 self.model['props'])
        return self

from fatiando import utils

x = np.linspace(0, 100000, 100)
z = -100*np.ones_like(x)
xs = np.linspace(0, 100000, 100)[::-1]
depths = (-1e-15*(xs - 50000)**4 + 8000
          - 3000*np.exp(-(xs - 70000)**2/(10000**2)))
depths -= depths.min()
density = -500
model = Polygon(np.transpose([xs, depths]), {'density':density})
gz = utils.contaminate(talwani.gz(x, z, [model]), 0.5, seed=0)

%matplotlib inline
from fatiando.vis import mpl

mpl.figure(figsize=(8, 5))
mpl.subplot(2, 1, 1)
mpl.plot(x, gz, 'ok')
mpl.subplot(2, 1, 2)
mpl.polygon(model, '-k', fill='grey')
mpl.gca().invert_yaxis()

misfit = PolyBasin2D(x, z, gz, 0, 30, density)
initial = 3000*np.ones(misfit.nparams)
misfit.config('levmarq', initial=initial).fit()

misfit.p_

misfit.estimate_

mpl.figure(figsize=(8, 5))
mpl.subplot(2, 1, 1)
mpl.plot(x, gz, 'ok')
mpl.plot(x, misfit.predicted(), '-r')
mpl.subplot(2, 1, 2)
mpl.polygon(model, '-k', fill='grey')
mpl.polygon(misfit.estimate_, 'o-r')
mpl.gca().invert_yaxis()

from fatiando.inversion.regularization import Smoothness1D
solver = misfit + 1e-8*Smoothness1D(misfit.nparams)
solver.config('levmarq', initial=initial).fit()

solver.estimate_

mpl.figure(figsize=(8, 5))
mpl.subplot(2, 1, 1)
mpl.plot(x, gz, 'ok')
mpl.plot(x, solver.predicted(), '-r')
mpl.subplot(2, 1, 2)
mpl.polygon(model, '-k', fill='grey')
mpl.polygon(solver.estimate_, 'o-r')
mpl.gca().invert_yaxis()

from fatiando.inversion.regularization import LCurve
lc = LCurve(misfit, Smoothness1D(misfit.nparams),
            [10**i for i in range(-15, -5)])
lc.config('levmarq', initial=initial).fit()

lc.estimate_

lc.regul_param_

lc.plot_lcurve()

mpl.figure(figsize=(8, 5))
mpl.subplot(2, 1, 1)
mpl.plot(x, gz, 'ok')
mpl.plot(x, lc.predicted(), '-r')
mpl.subplot(2, 1, 2)
mpl.polygon(model, '-k', fill='grey')
mpl.polygon(lc.estimate_, 'o-r')
mpl.gca().invert_yaxis()