#############################################################################################################
#																											#
#	Title			:	Simulation of wave function using standard volume approach for square quantum wire  #
#						geometry mesh generated by Gmsh and visualized										#
#	Author			:	Alfred Ajay Aureate R																#
#	Roll No			:	EE10B052																			#
#	Project guide	:	Prof. Anil Prabhakar, Dept.of Electrical Engineering, IITM							#
#	Code location	:	/DDP_codes/sim_mesh/square_qw/sqw.py												#
#	Figure ref.		:	NA																					#
#	Date			:	19th May, 2015																		#
#																											#
#############################################################################################################

from fipy import *
import numpy as np
nx = 20
ny = nx
dx = 1./nx
dy = dx
L = dx * nx
mesh = Grid2D(dx=dx, dy=dy, nx=nx, ny=ny)							# generates square mesh
phi = CellVariable(mesh = mesh, value = 0.)
D = 1.
eq = ImplicitSourceTerm(coeff=1.28445) == DiffusionTerm(coeff=D)	# don't know why 1.28445 gives normalized solution, but it does!

#############################################################################################################
# setting the boundary conditions (shifted down by 1 later)													#
#############################################################################################################
valueTop = 1
valueBottom = 1
valueLeft = 1
valueRight = 1
phi.constrain(valueTop, mesh.facesTop)
phi.constrain(valueBottom, mesh.facesBottom)
phi.constrain(valueRight, mesh.facesRight)
phi.constrain(valueLeft, mesh.facesLeft)

phi1 = 1-phi														# shifted down by 1 for correctness of the expected boundary conditions
phi1.setName("$\psi$")
phis = abs(phi1*phi1).sum()*((dx*nx)*(dy*ny))						
phi2 = phi1/np.sqrt(phis)											# normalizing (temporarily)
phi2.setName("$\psi$")

#############################################################################################################
# Visualizing the ground state of the square quantum wire for about 10 seconds								#
#############################################################################################################
if __name__ == '__main__':
	viewer = Viewer(vars=phi2, datamin=0., datamax=0.1)
	viewer.plot()
timeStepDuration = 10 * 0.9 * dx**2 / (2 * D)
steps = 10
for step in range(steps):
	eq.solve(var=phi, dt=timeStepDuration)
	phi1 = 1-phi
	phis = abs(phi1*phi1).sum()*((dx*nx)*(dy*ny))				
	phi2 = phi1/np.sqrt(phis)										# normalizing (temporarily)
	if __name__ == '__main__':
		viewer.plot()

