from Dynearthsol import Dynearthsol
import matplotlib.pyplot as plt 
import matplotlib.tri as tri
import numpy as np 
import triangle as tr
import sys 

G=6.67e-11
<<<<<<< HEAD
modelname="RC17"
modelinfo=modelname+".info"
nframe=9999
height=3000 
left=30000
=======
pi=3.1415926
modelname="mud50_02"
modelinfo=modelname+".info"
nframe=8600
height=3000
left=27000
>>>>>>> ies/master
right=110000

def area(Ecoord):
	ab=Ecoord[:,1,:]-Ecoord[:,0,:]
	ac=Ecoord[:,2,:]-Ecoord[:,0,:]
<<<<<<< HEAD
	_vec =np.hstack([ab,ac]).reshape(-1,2,2)
	area=np.abs(np.linalg.det(_vec))/2
=======
	area=np.abs(ab[:,0]*ac[:,1]-ab[:,1]*ac[:,0])/2
>>>>>>> ies/master
	return area


def gravity(pij,Ecenter,area,rho):
	dv=pij-Ecenter
	dr=np.sum(dv**2,axis=1)
	g0=2*G*np.dot(dv[:,1]*area/dr,rho)
	return g0

def freeair(vec_pij,Ecoord,rho):
	nsmaple=vec_pij.shape[0]
	g=np.empty(shape=nsmaple)
	A0=area(Ecoord)
	Ecenter=np.average(Ecoord,axis=1)
	for p0 in range(nsmaple):
		g[p0]=gravity(vec_pij[p0],Ecenter,A0,rho)
	return g

<<<<<<< HEAD
def bouguer(vx,vy,coord1,conn,rho):
	Ecoord=coord1[conn]
	nsmaple=len(vx)
	g=np.empty(shape=nsmaple)
	Ecenter=np.average(Ecoord,axis=1)
	mask=Ecenter[:,1]<0
	Ecenter=Ecenter[mask]
	Ecoord=Ecoord[mask]
	rho=rho[mask]
	A0=area(Ecoord)
	for p0 in range(nsmaple):
		g[p0]=gravity([vx[p0],vy[p0]],Ecenter,A0,rho)
	# g=g+0.386*vec_pij[:,1]
	return g

# def talwani(coord,conn,pij,rho):
# 	Ecoord= coord[conn]
# 	phi1=Ecoord[:,1,:]-Ecoord[:,0,:]
# 	phi2=Ecoord[:,2,:]-Ecoord[:,1,:]
# 	phi3=Dcoord[:,0,:]-Ecoord[:,2,:]
# 	dij=Ecoord[:,0,:]-pij
# 	theta=np.tan(dij[:,1]/dij[:,0])
# 	dg=2*G*rho
# 	return 


=======



def bouguer(vx,vy,coord,conn,rho,bcflag):
	coord1=coord.copy() 
	coord1[coord[:,1]>0,1]=0
	Ecoord=coord1[conn]
	Ecenter=np.average(Ecoord,axis=1)
	A0=area(Ecoord)+0.1

	
	# get X0 side element 
	sideTag=FindSideElement(conn, bcflag,1)
	rho1=rho[sideTag]
	conn1=conn[sideTag]
	bc1=(bcflag[conn1]&1).astype(bool)
	xz1=coord1[conn1[bc1]]

	# get X1 side element
	sideTag=FindSideElement(conn, bcflag,2)
	rho2=rho[sideTag]
	conn2=conn[sideTag]
	bc2=(bcflag[conn2]&2).astype(bool)
	xz2=coord1[conn2[bc2]]

	# start calculation
	nsmaple=len(vx)
	g=np.zeros(shape=nsmaple)
	for p0 in range(nsmaple):
		g[p0]+=gravity([vx[p0],vy[p0]],Ecenter,A0,rho)
		g[p0]+=SideCorrection(vx[p0],vy[p0],xz1,rho1,mod="-inf")
		g[p0]+=SideCorrection(vx[p0],vy[p0],xz2,rho2, mod="inf")
	return g



def FindSideElement(conn,bcflag,flag):
	side=np.sum((bcflag[conn]&flag), axis=1)==flag*2
	return side



def SideCorrection(xi, yi, xz,rho, mod):
	## Turcotte  & Schubert, Eq. 5-106
	## gravity for infinite slab
	dz=np.abs(xz[0::2,1]-xz[1::2,1])
	XX=xz[::2,0]-xi
	ZZ=np.abs(yi-(xz[::2,1]+xz[1::2,1])/2)
	angle=np.arctan(XX/ZZ)
	if mod in "inf":
		grav=np.sum(2*G*dz*rho*(pi/2-angle))
	if mod == "-inf":
		grav=np.sum(2*G*dz*rho*(angle+pi/2))

	return np.abs(grav)

>>>>>>> ies/master
def main(): 
	mopen=Dynearthsol(modelname)
	rho=mopen.read_field(nframe,'density')
	conn=mopen.read_field(nframe,'connectivity')
	coord=mopen.read_field(nframe,'coordinate')
	bcflag=mopen.read_field(nframe,'bcflag').astype(np.int16)
<<<<<<< HEAD
	node32=coord[(bcflag&32).astype(bool)]



	###### create new vertice #######
	## extend two walls 
	xbtm=coord[bcflag==17]
	xbtm2=[[xbtm[0,0]-10000,xbtm[0,1]]]
	xtop=coord[bcflag==33]
	xtop2=[[xtop[0,0]-10000,xtop[0,1]]]
	x0=np.vstack([xbtm,xbtm2,xtop, xtop2])
	tridict=dict(vertices=x0)
	B=tr.triangulate(tridict,'qa1000000')
	n0=coord.shape[0]
	coord=np.vstack([coord,B['vertices']])
	conn=np.vstack([conn,B['triangles']+n0])	
	rho=np.hstack([rho,np.ones(shape=B['triangles'].shape[0])*1750])
	bcflag=np.hstack([bcflag,np.zeros(shape=B['vertices'].shape[0])]).astype(np.int16)

	xbtm=coord[bcflag==66]
	xbtm2=[[xbtm[0,0]+10000,xbtm[0,1]]]
	xtop=coord[bcflag==34]
	xtop2=[[xtop[0,0]+10000,xtop[0,1]]]
	x0=np.vstack([xbtm,xbtm2,xtop, xtop2])
	tridict=dict(vertices=x0)
	B=tr.triangulate(tridict,'qa1000000')
	n0=coord.shape[0]
	coord=np.vstack([coord,B['vertices']])
	conn=np.vstack([conn,B['triangles']+n0])
	rho=np.hstack([rho,np.ones(shape=B['triangles'].shape[0])*1750])

	################################
=======
>>>>>>> ies/master

	
	measureX=np.arange(left,right,500)
	measureY=np.ones(shape=measureX.shape)*height
<<<<<<< HEAD



	result=bouguer(measureX,measureY,coord,conn,rho)*1e5



	##### figure #####
	fig,ax=plt.subplots(figsize=(10,2.5))
	ax.set_title(modelname)
	ax.plot(measureX,result)
	ax.set_ylabel('Bouguer gravity anomaly (mgal)(blue)')
	ax2=ax.twinx()
	ax2.plot(node32[:,0],node32[:,1],'r')
	ax2.set_ylabel('topography (m)(red)')
	# ax2.set_aspect('equal', adjustable="datalim")
=======
	result=bouguer(measureX,measureY,coord,conn,rho,bcflag)*1e5


	##### figure #####
	f1=plt.subplot(212)
	plt.tripcolor(coord[:,0], coord[:,1], conn,rho,lw=0.3,)
	plt.axis("equal")
	plt.subplot(211,sharex=f1)
	plt.plot(measureX,result)
	plt.ylabel('Bouguer gravity anomaly (mgal)(blue)')
	plt.title(modelname)
>>>>>>> ies/master
	plt.show()

	
if __name__== "__main__":
<<<<<<< HEAD
	main()
=======
	main()
>>>>>>> ies/master
