%matplotlib inline
import numpy as np
import json
import cPickle as pickle
from IPython.display import Image
from fatiando.gravmag.euler import Classic, ExpandingWindow, MovingWindow
from fatiando.gravmag import fourier, sphere, polyprism
from fatiando.mesher import PolygonalPrism, Sphere
from fatiando.vis import mpl, myv
from fatiando import utils
import fatiando
print("Using Fatiando a Terra version {0}".format((fatiando.version)))
mpl.rc('font', size=9)

x, y, z, tf = np.loadtxt('data/synthetic_data.txt', unpack=True)
with open('data/metadata.json') as f:
    metadata = json.load(f)
    area = metadata['area']
    bounds = metadata['bounds']
    shape = metadata['shape']
    inc = metadata['inc']
    dec = metadata['dec']

metadata

mpl.figure(figsize=(3.33, 2.7))
mpl.axis('scaled')
mpl.contourf(y, x, tf, shape, 60, cmap=mpl.cm.RdBu_r)
mpl.colorbar().set_label('nT')
mpl.m2km()
mpl.xlabel('East (km)')
mpl.ylabel('North (km)')
mpl.savefig('fig/data.png', dpi=600)

with open('data/model.pickle') as f:
    model = pickle.load(f)
    dike, sill, batholith = model

scene = myv.figure(size=(1000, 600))
myv.polyprisms([dike], linewidth=1, color=(0, 0, 1))
myv.polyprisms([sill], linewidth=1, color=(0, 1, 0))
myv.polyprisms([batholith], linewidth=1, color=(1, 0, 0))
ax = myv.axes(myv.outline(bounds), ranges=[b*0.001 for b in bounds], nlabels=3, fmt='%.0f')
ax.axes.x_label, ax.axes.y_label, ax.axes.z_label = 'x', 'y', 'z'
ax.axes.font_factor = 2
myv.wall_north(bounds)
myv.wall_bottom(bounds)
scene.scene.camera.position = [-37282.427190894407, -35382.676458226633, -36760.719551179718]
scene.scene.camera.focal_point = [13286.06353429469, 14199.525065081265, 5611.7451127001987]
scene.scene.camera.view_angle = 19.2
scene.scene.camera.view_up = [0.37050322249789591, 0.35546774432369316, -0.85812006436401433]
scene.scene.camera.clipping_range = [42859.55137481264, 132512.66650687021]
scene.scene.camera.compute_view_plane_normal()
scene.scene.render()
myv.savefig('fig/model.png')
myv.show()
Image(filename='fig/model.png')

print('Dike:')
print(' top (m): {}'.format(dike.z1))
print(' bottom (m): {}'.format(dike.z2))
print(' magnetization (SI): {}'.format(dike.props['magnetization']))

print('Sill:')
print(' top (m): {}'.format(sill.z1))
print(' bottom (m): {}'.format(sill.z2))
print(' magnetization (SI): {}'.format(sill.props['magnetization']))

print('Batholith:')
print(' top (m): {}'.format(batholith.z1))
print(' bottom (m): {}'.format(batholith.z2))
print(' magnetization (SI): {}'.format(batholith.props['magnetization']))

tf_si = utils.nt2si(tf)
xderiv = fourier.derivx(x, y, tf_si, shape)
yderiv = fourier.derivy(x, y, tf_si, shape)
zderiv = fourier.derivz(x, y, tf_si, shape)

mpl.figure(figsize=(10, 4))
mpl.subplots_adjust()
levels = 30
cbargs = dict(orientation='horizontal')

ax = mpl.subplot(1, 3, 1)
mpl.axis('scaled')
mpl.title('x derivative')
mpl.contourf(y, x, xderiv, shape, levels, cmap=mpl.cm.RdBu_r)
mpl.colorbar(**cbargs).set_label('T/m')
mpl.m2km()
mpl.xlabel('East (km)')
mpl.ylabel('North (km)')

ax = mpl.subplot(1, 3, 2)
mpl.axis('scaled')
mpl.title('y derivative')
mpl.contourf(y, x, yderiv, shape, levels, cmap=mpl.cm.RdBu_r)
mpl.colorbar(**cbargs).set_label('T/m')
mpl.m2km()
mpl.xlabel('East (km)')
mpl.ylabel('North (km)')

ax = mpl.subplot(1, 3, 3)
mpl.axis('scaled')
mpl.title('z derivative')
mpl.contourf(y, x, zderiv, shape, levels, cmap=mpl.cm.RdBu_r)
cb = mpl.colorbar(**cbargs).set_label('T/m')
mpl.m2km()
mpl.xlabel('East (km)')
mpl.ylabel('North (km)')

euler_s1 = [MovingWindow(Classic(x, y, z, tf_si, xderiv, yderiv, zderiv, i), 
                         windows=(50, 50), size=(1000, 1000), keep=0.3).fit()
            for i in [1, 2, 3]]

euler_s3 = [MovingWindow(Classic(x, y, z, tf_si, xderiv, yderiv, zderiv, i), 
                         windows=(50, 50), size=(3000, 3000), keep=0.3).fit()
            for i in [1, 2, 3]]

euler_s5 = [MovingWindow(Classic(x, y, z, tf_si, xderiv, yderiv, zderiv, i), 
                         windows=(50, 50), size=(5000, 5000), keep=0.3).fit()
            for i in [1, 2, 3]]

# Get the maximum value of z that we got from all runs.
# We'll use this to set the limits of the color bars
maxz = max(p[2] for euler in [euler_s1, euler_s3] for e in euler for p in e.estimate_)
labels = ['(a)', '(b)', '(c)', '(d)', '(e)', '(f)']
fig = mpl.figure(figsize=(7, 4.8))
mpl.subplots_adjust(left=0.06, top=0.98, bottom=0.05, wspace=0, hspace=0.02)
for j, euler in enumerate([euler_s1, euler_s3]):
    for i, e in enumerate(euler):
        xc, yc, zc = e.estimate_.T
        xc, yc, zc = xc[zc > 0], yc[zc > 0], zc[zc > 0]
        ax = mpl.subplot(2, 3, 3*j + i + 1)
        mpl.axis('scaled')
        mpl.text(6000, 23000, labels[3*j + i], fontsize=12, fontweight='bold')
        mpl.title('SI={0} | Window={1:.0f} km'.format(e.euler.model['structural_index'], e.size[0]*0.001))
        #for body in model:
        #    mpl.polygon(body, xy2ne=True, style='-b', linewidth=1.5)
        plot = mpl.scatter(yc, xc, s=30, c=zc*0.001, cmap=mpl.cm.RdYlGn_r, vmin=0, vmax=maxz*0.001)
        mpl.set_area(area)
        mpl.m2km()
        if i in [1, 2, 4, 5]:
            ax.set_yticklabels([])
        if j == 1:
            mpl.xlabel('East (km)')
            ax.set_xticklabels(['{:.0f}'.format(t*0.001) for t in ax.get_xticks()[:-1]])
        if j in [0]:
            ax.set_xticklabels([])
        if i in [0, 3]:
            mpl.ylabel('North (km)')
fig.colorbar(plot, cax=mpl.axes([0.91, 0.09, 0.03, 0.85]), orientation='vertical').set_label('Depth (km)')
mpl.savefig('fig/euler-solutions-low.png')
mpl.savefig('fig/euler-solutions.png', dpi=600)

# Make the figure for the 5km window separately because it won't go in the tutorial
maxz = max(p[2] for euler in [euler_s5] for e in euler for p in e.estimate_)
fig = mpl.figure(figsize=(12, 3.5))
for i, e in enumerate(euler_s5):
        xc, yc, zc = e.estimate_.T
        xc, yc, zc = xc[zc > 0], yc[zc > 0], zc[zc > 0]
        ax = mpl.subplot(1, 3, i + 1)
        mpl.axis('scaled')
        mpl.title('SI={0} | Window={1:.0f} km'.format(e.euler.model['structural_index'], e.size[0]*0.001))
        plot = mpl.scatter(yc, xc, s=30, c=zc*0.001, cmap=mpl.cm.RdYlGn_r, vmin=0, vmax=maxz*0.001)
        mpl.set_area(area)
        mpl.m2km()
        mpl.xlabel('East (km)')
        mpl.ylabel('North (km)')
fig.colorbar(plot, cax=mpl.axes([0.91, 0.09, 0.03, 0.85]), orientation='vertical').set_label('Depth (km)')

scene = myv.figure(size=(1000, 600))
myv.polyprisms(model, opacity=0.3, linewidth=2)
myv.points([e for e in euler_s3[2].estimate_ if e[2] > 0])
ax = myv.axes(myv.outline(bounds), ranges=[b*0.001 for b in bounds], nlabels=3, fmt='%.0f')
ax.axes.x_label, ax.axes.y_label, ax.axes.z_label = 'x', 'y', 'z'
ax.axes.font_factor = 2
myv.wall_north(bounds)
myv.wall_bottom(bounds)
scene.scene.camera.position = [-10753.852296021096, -59325.007533753123, -22266.66194804628]
scene.scene.camera.focal_point = [16722.015964135739, 13784.075401721391, 4529.0411289607291]
scene.scene.camera.view_angle = 19.2
scene.scene.camera.view_up = [0.12859903242950413, 0.29830909042254711, -0.94576634293543571]
scene.scene.camera.clipping_range = [43526.915602369831, 131500.63847030781]
scene.scene.camera.compute_view_plane_normal()
scene.scene.render()
myv.savefig('fig/euler-solutions-3d.png')
scene.scene.camera.position = [61783.513706029873, -51534.690503972124, -18750.237814442207]
scene.scene.camera.focal_point = [14324.659515610345, 12669.434351699714, 2305.8434486142896]
scene.scene.camera.view_angle = 5.033164800000001
scene.scene.camera.view_up = [-0.14301336238889156, 0.21138117409951956, -0.96688426267807892]
scene.scene.camera.clipping_range = [41498.787075197979, 137750.00692720726]
scene.scene.camera.compute_view_plane_normal()
scene.scene.render()
myv.savefig('fig/euler-solutions-3d-zoom.png')
myv.show()
Image(filename='fig/euler-solutions-3d.png')

Image(filename='fig/euler-solutions-3d-zoom.png')

ideal_batholith = Sphere(18000, 12000, 2250, radius=2200, props=batholith.props)
ideal_batholith.center = [ideal_batholith.x, ideal_batholith.y, ideal_batholith.z]

scene = myv.figure(size=(800, 300))
myv.polyprisms([batholith], opacity=0.5, linewidth=2)
myv.points([ideal_batholith.center], color=(1, 0, 0), size=ideal_batholith.radius)
ax = myv.axes(myv.outline(bounds), ranges=[b*0.001 for b in bounds], nlabels=3, fmt='%.0f')
ax.axes.x_label, ax.axes.y_label, ax.axes.z_label = 'x', 'y', 'z'
ax.axes.font_factor = 2
myv.wall_north(bounds)
myv.wall_bottom(bounds)
scene.scene.camera.position = [-15337.932586926319, -59396.608310493262, -7022.645217992851]
scene.scene.camera.focal_point = [17482.870070805882, 15380.030101355769, 4903.8791935470208]
scene.scene.camera.view_angle = 9.830400000000001
scene.scene.camera.view_up = [0.064932345527732721, 0.12931668910473115, -0.98947510550961171]
scene.scene.camera.clipping_range = [40404.996563304114, 131956.0403000123]
scene.scene.camera.compute_view_plane_normal()
scene.scene.render()
myv.savefig('tmp/sphere-batholith.png')
myv.show()
Image(filename='tmp/sphere-batholith.png')

tf_sphere = utils.nt2si(
    utils.contaminate(
        sphere.tf(x, y, z, [ideal_batholith], inc, dec) + polyprism.tf(x, y, z, [dike, sill], inc, dec),
        5))

mpl.figure(figsize=(10, 3.5))
mpl.subplot(1, 2, 1)
mpl.title('Polygonal prism')
mpl.axis('scaled')
mpl.contourf(y, x, tf_si, shape, 60, cmap=mpl.cm.RdBu_r)
mpl.colorbar().set_label('T')
mpl.m2km()
mpl.xlabel('East (km)')
mpl.ylabel('North (km)')
mpl.subplot(1, 2, 2)
mpl.title('Sphere')
mpl.axis('scaled')
mpl.contourf(y, x, tf_sphere, shape, 60, cmap=mpl.cm.RdBu_r)
mpl.colorbar().set_label('T')
mpl.m2km()
mpl.xlabel('East (km)')
mpl.ylabel('North (km)')

xderiv_sphere = fourier.derivx(x, y, tf_sphere, shape)
yderiv_sphere = fourier.derivy(x, y, tf_sphere, shape)
zderiv_sphere = fourier.derivz(x, y, tf_sphere, shape)

euler_sphere = MovingWindow(
    Classic(x, y, z, tf_sphere, xderiv_sphere, yderiv_sphere, zderiv_sphere, 3),
    windows=(50, 50), size=(3000, 3000), keep=0.3).fit()

xc, yc, zc = euler_sphere.estimate_.T
xc, yc, zc = xc[zc > 0], yc[zc > 0], zc[zc > 0]

fig = mpl.figure()
mpl.axis('scaled')
plot = mpl.scatter(yc, xc, s=30, c=zc*0.001, cmap=mpl.cm.RdYlGn_r, vmin=0)
mpl.colorbar().set_label('Depth (km)')
mpl.set_area(area)
mpl.m2km()
mpl.xlabel('East (km)')
mpl.ylabel('North (km)')

scene = myv.figure(size=(1000, 600))
myv.polyprisms([sill, dike], opacity=0.3, linewidth=2)
myv.points([ideal_batholith.center], color=(0, 0, 1), size=ideal_batholith.radius, opacity=0.3)
myv.points([e for e in euler_sphere.estimate_ if e[2] > 0])
ax = myv.axes(myv.outline(bounds), ranges=[b*0.001 for b in bounds], nlabels=3, fmt='%.0f')
ax.axes.x_label, ax.axes.y_label, ax.axes.z_label = 'x', 'y', 'z'
ax.axes.font_factor = 2
myv.wall_north(bounds)
myv.wall_bottom(bounds)
scene.scene.camera.position = [-10753.852296021096, -59325.007533753123, -22266.66194804628]
scene.scene.camera.focal_point = [16722.015964135739, 13784.075401721391, 4529.0411289607291]
scene.scene.camera.view_angle = 19.2
scene.scene.camera.view_up = [0.12859903242950413, 0.29830909042254711, -0.94576634293543571]
scene.scene.camera.clipping_range = [43526.915602369831, 131500.63847030781]
scene.scene.camera.compute_view_plane_normal()
scene.scene.render()
myv.savefig('tmp/euler-solutions-3d-sphere.png')
scene.scene.camera.position = [61783.513706029873, -51534.690503972124, -18750.237814442207]
scene.scene.camera.focal_point = [14324.659515610345, 12669.434351699714, 2305.8434486142896]
scene.scene.camera.view_angle = 5.033164800000001
scene.scene.camera.view_up = [-0.14301336238889156, 0.21138117409951956, -0.96688426267807892]
scene.scene.camera.clipping_range = [41498.787075197979, 137750.00692720726]
scene.scene.camera.compute_view_plane_normal()
scene.scene.render()
myv.savefig('tmp/euler-solutions-3d-sphere-zoom.png')
myv.show()
Image(filename='tmp/euler-solutions-3d-sphere.png')

Image(filename='tmp/euler-solutions-3d-sphere-zoom.png')