Hydrogen orbitals in Julia

orbitals.jl

# Cribbed from http://commons.wikimedia.org/wiki/User:Geek3/hydrogen

function LegendreP(l, m, x)
    m = abs(m)
    a = 1
    if m > 0
        fac = 1
        c = (1 + x) * (1 - x)
        for i = 1:m
            a *= c * fac / (fac + 1)
            fac += 2
        end
    end
    a = sqrt((2 * m + 1) * a / (4 * pi))
    if m & 1 != 0
        a = -a
    end
    if l == m
        return a
    else
        b = a * x * sqrt(2 * m + 3)
        if l == m + 1
            return b
        else
            fac1 = sqrt(2 * m + 3)
            pn = 0
            for j = m+2:l
                fac = sqrt((4.0 * j * j - 1.0) / (j * j - m * m))
                pn = fac * (b * x - a / fac1)
                fac1 = fac
                a = b
                b = pn
            end
            return pn
        end
    end
end

function LaguerreL(n, k, x)
    L1 = 0.0
    res = 1.0
    for i = 0:n-1
        L0 = L1
        L1 = res
        res = ((2 * i + 1 + k - x) * L1 - (i + k) * L0) / (i + 1)
    end
    res
end

function Rnl(n, l, r)
    rho = 2 * r / n
    a = 4
    for i = n-l:n+l
        a /= i
    end
    a = sqrt(a)
    a /= n * n
    a * exp(-0.5 * rho) * rho^l * LaguerreL(n-l-1, 2*l+1, rho)
end

function Psi(n, l, m, x, y, z)
    r = sqrt(x^2 + y^2 + z^2)
    pabs = if r == 0
        Rnl(n, l, r) * LegendreP(l, m, 0)
    else
        Rnl(n, l, r) * LegendreP(l, m, z/r)
    end
    p0 = (n + l + m + 1) * pi
    if pabs < 0
        pabs = -pabs
        p0 += pi
    end
    if m != 0
        rxy = sqrt(x^2 + y^2)
        if rxy != 0
            p0 += m * atan2(y,x)
        end
    end
    pphase = p0 % (2*pi)
    (pabs, pphase)
end

function Psi_sqr(n, l, m, x, y, z)
    r = sqrt(x^2 + y^2 + z^2)
    pabs = Rnl(n, l, r) * if r == 0
        LegendreP(l, m, 0)
    else
        LegendreP(l, m, z/r)
    end
    pabs * pabs
end

using PyPlot
using Color

color_map = ColorMap("spring-2015", [
    color("rgb(156,198,217)"),
    color("rgb(124,159,192)"),
    color("rgb(229,133,142)"),
    HSL(24, 0.93, 0.64),
        HSL(24, 0.93, 0.70),
        HSL(24, 0.93, 0.78),
        HSL(24, 0.93, 0.94),
])

d = 70
r = 2500
for n=6:7, l=0:n-1, m=0:l
  print("Generating n=$n l=$l m=$m...\n")
  #n = 2; l = 1; m = 0
  Z = [Psi_sqr(n,l,m, a, 0, b) for a=linspace(-d,d,r), b=linspace(-d,d,r)]
  maxZ = maximum(Z)
  f = (x) -> log(maxZ/200+x)
  fig = figure(figsize=(50, 50), dpi=800)
  ax = fig[:add_axes]([0,0,1,1])
  plt.axis("off")
  ax[:imshow](map(f,Z), cmap=color_map)
  #ax = fig[:add_subplot](2,2,2)
  #im = ax[:imshow](Z, cmap=color_map)
  savefig("test/test-$n-$l-$m.png")
end
return