Files
simpeg/simpegDCIP/Dev/DC2D_Demo_HTML_anim.py
T
D Fournier 78584ae49d Create pseudo-section simulation in Notebook.
Sub-functions added to BaseDC. Required for the simulation
2016-02-01 21:03:58 -08:00

217 lines
6.4 KiB
Python

import os
from SimPEG import *
import simpegDCIP as DC
import pylab as plt
from matplotlib import animation
from JSAnimation import HTMLWriter
import time
import re
#from readUBC_DC2DMesh import readUBC_DC2DMesh
#from readUBC_DC2DModel import readUBC_DC2DModel
#from readUBC_DC2DLoc import readUBC_DC2DLoc
#from convertObs_DC3D_to_2D import convertObs_DC3D_to_2D
#from readUBC_DC3Dobs import readUBC_DC3Dobs
#%%
home_dir = 'C:\\Users\\dominiquef.MIRAGEOSCIENCE\\ownCloud\\Research\\Modelling\\Synthetic\\Two_Sphere'
msh_file = 'Mesh_2D.msh'
mod_file = 'Model_2D.con'
obs_file = 'FWR_data3D.dat'
dsep = '\\'
# Forward solver
slvr = 'BiCGStab' #'LU'
# Preconditioner
pcdr = 'Jacobi' #'Gauss-Seidel'#
# Number of padding cells to remove from plotting
padc = 15
# Load UBC mesh 2D
mesh = DC.readUBC_DC2DMesh(home_dir + dsep + msh_file)
# Load model
model = DC.readUBC_DC2DModel(home_dir + dsep + mod_file)
# load obs file
[Tx,Rx,d,wd] = DC.readUBC_DC3Dobs(home_dir + dsep + obs_file)
[Tx, Rx] = DC.convertObs_DC3D_to_2D(Tx,Rx)
#%% Create system
#Set boundary conditions
mesh.setCellGradBC('neumann')
Div = mesh.faceDiv
Grad = mesh.cellGrad
Msig = Utils.sdiag(1./(mesh.aveF2CC.T*(1./model)))
A = Div*Msig*Grad
# Change one corner to deal with nullspace
A[0,0] = 1
A = sp.csc_matrix(A)
start_time = time.time()
if re.match(slvr,'BiCGStab'):
# Create Jacobi Preconditioner
if re.match(pcdr,'Jacobi'):
dA = A.diagonal()
P = sp.spdiags(1/dA,0,A.shape[0],A.shape[0])
# Create Gauss-Seidel Preconditioner
elif re.match(pcdr,'Gauss-Seidel'):
LD = sp.tril(A,k=0)
#LDinv = sp.linalg.splu(LD)
elif re.match(slvr,'LU'):
# Factor A matrix
Ainv = sp.linalg.splu(A)
print("LU DECOMP--- %s seconds ---" % (time.time() - start_time))
#%% Create SimPEG objects
# Create sub-mesh for plotting
hx = mesh.hx
hy = mesh.hy
hx_sub = hx[padc:-padc]
hy_sub = hy[padc:]
mesh_sub = Mesh.TensorMesh([hx_sub,hy_sub],(mesh.vectorNx[padc], mesh.vectorNy[padc]))
model_sub = model.reshape(mesh.nCy,mesh.nCx)
model_sub = mkvc(model_sub[padc:,padc:-padc].T)
xx = mesh_sub.vectorCCx
yy = mesh_sub.vectorCCy
#%% Solve
#txii = range(50,1950,100)
#jx_CC_sub = np.zeros((len(txii),mesh_sub.nCx,mesh_sub.nCy))
#jy_CC_sub = np.zeros((len(txii),mesh_sub.nCx,mesh_sub.nCy))
fig = plt.figure(figsize=(10,5))
axs = plt.axes(ylim = (yy[0],yy[-1]+mesh.hy[-1]*2), xlim = (xx[0],xx[-1]))#
plt.tight_layout(pad=0.4, w_pad=0.5, h_pad=1.0)
plt.ylim(yy[0],yy[-1]+mesh.hy[-1]*2)
plt.xlim(xx[0],xx[-1])
#im1 = axs.pcolormesh([],[],[], alpha=0.75,extent = (xx[0],xx[-1],yy[-1],yy[0]),interpolation='nearest',vmin=-1e-2, vmax=1e-2)
#im2 = axs.pcolormesh([],[],[],alpha=0.2,extent = (xx[0],xx[-1],yy[-1],yy[0]),interpolation='nearest',cmap='gray')
im1 = axs.pcolormesh(mesh_sub.vectorCCx,mesh_sub.vectorCCy,np.zeros((mesh_sub.nCy,mesh_sub.nCx)), alpha=0.75,vmin=-1e-2, vmax=1e-2)
im2 = axs.pcolormesh(mesh_sub.vectorCCx,mesh_sub.vectorCCy,np.zeros((mesh_sub.nCy,mesh_sub.nCx)), alpha=0.75,vmin=-1e-2, vmax=1e-2)
im3 = axs.streamplot(xx, yy, np.zeros((mesh_sub.nCy,mesh_sub.nCx)), np.zeros((mesh_sub.nCy,mesh_sub.nCx)),color='k')
im4 = axs.scatter([],[], c='r', s=200)
im5 = axs.scatter([],[], c='r', s=200)
#==============================================================================
# def init():
# im1.set_data([[],[],[]])
# im2.set_data([[],[],[]])
#
# return [im1]+[im2]
#==============================================================================
def animate(ii):
#for ii in range(len(txii)):
removeStream()
tx = np.asarray(np.c_[Tx[ii],np.ones(Tx[ii].shape[0])*mesh.vectorNy[-1]-1])
inds = Utils.closestPoints(mesh, tx )
RHS = mesh.getInterpolationMat( tx , 'CC').T*( [-1,1] / mesh.vol[inds] )
if re.match(slvr,'BiCGStab'):
if re.match(pcdr,'Jacobi'):
dA = A.diagonal()
P = sp.spdiags(1/dA,0,A.shape[0],A.shape[0])
# Iterative Solve
phi = sp.linalg.bicgstab(P*A,P*RHS, tol=1e-5)
phi = mkvc(phi[0])
elif re.match(slvr,'LU'):
#Direct Solve
phi = Ainv.solve(RHS)
j = -Msig*Grad*phi
j_CC = mesh.aveF2CCV*j
# Compute charge density solving div*grad*phi
Q = -mesh.faceDiv*mesh.cellGrad*phi
jx_CC = j_CC[0:mesh.nC].reshape(mesh.nCy,mesh.nCx)
jy_CC = j_CC[mesh.nC:].reshape(mesh.nCy,mesh.nCx)
#%% Grab only the core for presentation
jx_CC_sub = jx_CC[padc:,padc:-padc]
jy_CC_sub = jy_CC[padc:,padc:-padc]
Q_sub = Q.reshape(mesh.nCy,mesh.nCx)
Q_sub = Q_sub[padc:,padc:-padc]
J_rho = np.sqrt(jx_CC_sub**2 + jy_CC_sub**2)
lw = np.log10(J_rho/J_rho.min())
#axs.imshow(Q_sub,alpha=0.75,extent = (xx[0],xx[-1],yy[-1],yy[0]),interpolation='nearest',vmin=-1e-2, vmax=1e-2)
#axs.imshow(np.log10(model_sub.reshape(mesh_sub.nCy,mesh_sub.nCx)),alpha=0.2,extent = (xx[0],xx[-1],yy[-1],yy[0]),interpolation='nearest',cmap='gray')
global im1
im1 = axs.pcolormesh(mesh_sub.vectorCCx,mesh_sub.vectorCCy,Q_sub, alpha=0.75,vmin=-1e-2, vmax=1e-2)
global im2
im2 = axs.pcolormesh(mesh_sub.vectorCCx,mesh_sub.vectorCCy,np.log10(model_sub.reshape(mesh_sub.nCy,mesh_sub.nCx)), alpha=0.25)
global im3
im3 = axs.streamplot(xx, yy, jx_CC_sub, jy_CC_sub,color='k',linewidth = lw,density=0.5)
global im4
im4 = axs.scatter(tx[0,0],mesh.vectorNy[-1], c='r', s=75, marker='v' )
global im5
im5 = axs.scatter(tx[1,0],mesh.vectorNy[-1], c='b', s=75, marker='v' )
#plt.show()
#im1.set_array(Q_sub)
#im2.set_array(np.log10(model_sub.reshape(mesh_sub.nCy,mesh_sub.nCx)))
#im2.set_array(mesh_sub.vectorCCx, mesh_sub.vectorCCy,jx_CC_sub.T,jy_CC_sub.T)
#return [im1] + [im2]
#%% Create widget
def removeStream():
global im1
im1.remove()
global im2
im2.remove()
global im3
im3.lines.remove()
axs.patches = []
global im4
im4.remove()
global im5
im5.remove()
#def viewInv(msh,iteration):
#, linewidth=lw.T
#%%
#interact(viewInv,msh = mesh_sub, iteration = IntSlider(min=0, max=len(txii)-1 ,step=1, value=0))
# set embed_frames=True to embed base64-encoded frames directly in the HTML
anim = animation.FuncAnimation(fig, animate,
frames=len(Tx), interval=5)
anim.save(home_dir + '\\animation.html', writer=HTMLWriter(embed_frames=True))