import json
import cPickle as pickle
import numpy as np
from IPython.display import Image
from fatiando.mesher import PolygonalPrism
from fatiando.gravmag import polyprism
from fatiando  import gridder, utils
from fatiando.vis import myv, mpl
import fatiando
print("Using Fatiando a Terra version {0}".format((fatiando.version)))

area = [5000, 25000, 5000, 25000]
bounds = [0, 30000, 0, 30000, 0, 5000]

mpl.figure()
mpl.axis('scaled')
mpl.square(area)
dike_verts = mpl.draw_polygon(bounds[:4], mpl.gca(), xy2ne=True)

dike_verts

dike = PolygonalPrism(dike_verts, 0, 5000, {'magnetization': 10})

myv.figure(size=(600, 400))
myv.polyprisms([dike], linewidth=2)
myv.axes(myv.outline(bounds), ranges=[b*0.001 for b in bounds], nlabels=3, fmt='%.1f')
myv.wall_north(bounds)
myv.wall_bottom(bounds)
myv.savefig('tmp/model_dike.png')
myv.show()
Image(filename='tmp/model_dike.png')

mpl.figure()
mpl.axis('scaled')
mpl.square(area)
mpl.polygon(dike, xy2ne=True)
sill_verts = mpl.draw_polygon(bounds[:4], mpl.gca(), xy2ne=True)

sill_verts

sill = PolygonalPrism(sill_verts, 1000, 1500, {'magnetization': 10})

myv.figure(size=(600, 400))
myv.polyprisms([dike, sill], linewidth=2)
myv.axes(myv.outline(bounds), ranges=[b*0.001 for b in bounds], nlabels=3, fmt='%.1f')
myv.wall_north(bounds)
myv.wall_bottom(bounds)
myv.savefig('tmp/model_sill.png')
myv.show()
Image(filename='tmp/model_sill.png')

mpl.figure()
mpl.axis('scaled')
mpl.square(area)
mpl.polygon(dike, xy2ne=True)
mpl.polygon(sill, xy2ne=True)
batholith_verts = mpl.draw_polygon(bounds[:4], mpl.gca(), xy2ne=True)

batholith_verts

batholith = PolygonalPrism(batholith_verts, 500, 4000, {'magnetization': 2})

myv.figure(size=(600, 400))
myv.polyprisms([dike, sill, batholith], linewidth=2)
myv.axes(myv.outline(bounds), ranges=[b*0.001 for b in bounds], nlabels=3, fmt='%.1f')
myv.wall_north(bounds)
myv.wall_bottom(bounds)
myv.savefig('tmp/model_batholith.png')
myv.show()
Image(filename='tmp/model_batholith.png')

model = [dike, sill, batholith]

shape = (100, 100)
x, y, z = gridder.regular(area, shape, z=-300)
inc, dec = -15, 30

tf = utils.contaminate(polyprism.tf(x, y, z, model, inc, dec), 5)

mpl.figure()
mpl.axis('scaled')
for b in model:
    mpl.polygon(b, xy2ne=True)
mpl.contourf(y, x, tf, shape, 60, cmap=mpl.cm.RdBu_r)
cb = mpl.colorbar().set_label('nT')
mpl.m2km()
mpl.xlabel('East (km)')
mpl.ylabel('North (km)')
mpl.savefig('tmp/synthetic_data.png')
mpl.show()
Image(filename='tmp/synthetic_data.png')

np.savetxt('data/synthetic_data.txt', np.transpose([x, y, z, tf]))
This is what it looks like.
# Can run bash commands inside the IPython notebook
!head data/synthetic_data.txt

with open('data/metadata.json', 'w') as f:
    json.dump(dict(area=area, bounds=bounds, inc=inc, dec=dec, shape=shape), f)

!cat data/metadata.json

with open('data/model.pickle', 'w') as f:
    pickle.dump(model, f)

with open('data/dike.json', 'w') as f:
    json.dump(dike_verts.tolist(), f)
with open('data/sill.json', 'w') as f:
    json.dump(sill_verts.tolist(), f)
with open('data/batholith.json', 'w') as f:
    json.dump(batholith_verts.tolist(), f)