main.py (1538B)
1 #!/usr/bin/python3.9 2 3 from dolfin import * 4 import numpy as np 5 import os 6 7 def main(): 8 mesh = UnitCubeMesh(25, 25, 25) 9 V = FunctionSpace(mesh, 'CG', 1) 10 11 m_e = 9.10e-31; m_k = 1.67e-27 12 hbar = 1.05e-34; k = 8.98e9 13 mu = m_e*m_k/(m_e + m_k) 14 e = 1.60e-19 15 16 # Potential for the Hydrogen Atom 17 # { 1/r if r > 0 18 # V = { 19 # { 0 else 20 pot_ = Expression('sqrt(x[0]*x[0]+x[1]*x[1]+x[2]*x[2])>0 ?\ 21 -1/sqrt(x[0]*x[0]+x[1]*x[1]+x[2]*x[2]) : 0',\ 22 degree=4) 23 pot = interpolate(pot_, V) 24 25 26 def boundary(x, on_boundary): 27 return on_boundary 28 29 bc = DirichletBC(V, Constant(0.), boundary) 30 31 u = TrialFunction(V) 32 v = TestFunction(V) 33 34 # Assemble system 35 a = hbar/(2*mu)*inner(grad(u), grad(v))*dx + e**2/k*pot*u*v*dx 36 m = u*v*dx 37 38 # Define Matrices 39 A = PETScMatrix() 40 M = PETScMatrix() 41 42 assemble(a, tensor=A) 43 assemble(m, tensor=M) 44 45 # Apply boundary on the system 46 bc.apply(A) 47 bc.apply(M) 48 49 # Create eigensolver 50 eigensolver = SLEPcEigenSolver(A, M) 51 eigensolver.parameters["spectrum"] = "smallest magnitude" 52 53 values = 30 54 eigensolver.solve(values) 55 56 u = Function(V) 57 58 59 if os.path.exists('./meshes') != True: 60 os.mkdir('./meshes') 61 62 f = File('./meshes/orbitals.pvd') 63 64 for i in range(values): 65 E, E_c, R, R_c = eigensolver.get_eigenpair(i) 66 67 u.vector()[:] = np.sqrt(R*R + R_c*R_c) 68 69 f << (u, i) 70 71 if __name__ == "__main__": 72 main()