Numeric solution to 3D hydrogen atom Schrodinger Equation. Reimplemented against https://github.com/quantum-visualizations/qmsolve , but with more efficiency.

hatom.jl

#!/usr/bin/env julia

# Constants

const k = 3.8099821161548593 # hbar^2 / (2*m_e) /(Å^2) / eV
const k_c = 14.39964547842567 # (e*e / (4 * np.pi * epsilon_0))  # measured in eV / Å
const m = 1.0  # mass of electron


using LinearAlgebra;
using SparseArrays;
using Arpack;

function mgrid(xs...)
    it = Iterators.product(xs...);
    ret = [];
    for i in 1:length(xs)
        push!(ret, getindex.(it, i));
    end
    return ret;
end

const xlen = 30.0;
const ylen = xlen;
const zlen = xlen;

const N = 50;

const dx = xlen / N;

# generate the grid
const xdat = range(-xlen/2, xlen/2, length=N);
const ydat = range(-ylen/2, ylen/2, length=N);
const zdat = range(-zlen/2, zlen/2, length=N);

gx, gy, gz = mgrid(xdat, ydat, zdat);

# Input the Potential
function potential(gx, gy, gz) ::AbstractArray
    r = sqrt.(gx.^2 + gy.^2 + gz.^2);
    r[r .< 0.0001] .= 0.0001;
    return -k_c ./ r;
end


# Construct Hamiltonian

# Kinetic part
Identity = sparse(1.0I, N, N);

T_ = sparse(-2.0I, N, N);
T_[diagind(T_,  1)] .= 1.0;
T_[diagind(T_, -1)] .= 1.0;

T_ .*= -k / (m * dx*dx)

T = kron(T_, kron(Identity, Identity)) + 
    kron(Identity, kron(T_, Identity)) + 
    kron(Identity, kron(Identity, T_));

# Potential part
P = potential(gx, gy, gz);
# P = reshape(P, N^3);
Pmin = minimum(P);

# Hamiltonian
H = T;
H[diagind(H)] .+= vec(P);


# Solve the equation Hψ = Eψ
λ, ϕ = eigs(H, nev=10, which=:LM, sigma=min(0, Pmin));

println(λ);

#= 
[-10.712968304076878, -3.4170976668414994, -3.4170976668414994, -3.417097666841471, -3.031393331772243, -1.5153353066901687, -1.5153353066901154, -1.5153353066900657, -1.4739896184375283, -1.4739896184374004]
 32.593136 seconds (1.10 M allocations: 2.621 GiB, 0.15% gc time, 1.15% compilation time)

Analytical Result:
[-13.6, -3.4, -3.4, -3.4, -3.4, -1.511 ...]
=#