############################################
## Written in 2012 by
##   Andrew Poelstra <apoelstra@wpsoftware.net>
##
## To the extent possible under law, the author(s) have dedicated all
## copyright and related and neighboring rights to this software to
## the public domain worldwide. This software is distributed without
## any warranty.
##
## You should have received a copy of the CC0 Public Domain Dedication
## along with this software.
## If not, see <http://creativecommons.org/publicdomain/zero/1.0/>.
##
############################################

#####
##### Outputs the raw data for a sigma bond.
##### Run this file, and redirect the output to some file. Strip the
##### header and "ans = " to leave only numerical data, and feed it
##### to the utility 'render-me' to get a pretty picture.
#####


## CONSTANTS
# We will use Lorentz-Heaviside natural units -- c = hbar = e_0 = u_0 = 1
# In this system, one unit of space is roughly 1.97e-7 meters, or ~200nm
global hbar = 1;
global m    = 0.511000e6; # eV/c^2
global e    = 0.302822;   # sqrt(4*pi*alpha)
global k    = 1/(4*pi);
global a    = 137.036/m   # Bohr radius (for exact H solution)


## Converts natural-ordering index into 3D coordinates
function [x y z] = from_matrix_coords (i, M)
  h = 2 / M;
  # Create M^3 grid on [-1, 1-h]^3
  x = -1 + h * mod (floor ((i - 1) / M^2), M);
  y = -1 + h * mod (floor ((i - 1) / M), M);
  z = -1 + h * mod (i - 1, M);
  # Shift by half a grid unit (because our potentials will have
  # a pole at the origin and we want to avoid it).
  x += h/2;
  y += h/2;
  z += h/2;
  # SCALE
  x /= 20;
  y /= 20;
  z /= 20;
endfunction


## Checks if some index corresponds to a boundary
function rv = is_boundary (i, M)
  x = mod (floor ((i - 1) / M^2), M);
  y = mod (floor ((i - 1) / M), M);
  z = mod (i - 1, M);
  rv = (x == 0 || x == M-1 || y == 0 || y == M-1 || z == 0 || z == M);
endfunction


## Builds a matrix for solving the Schroedinger equatation
##       hbar^2
##     - ------ laplace(u) + Vu = Eu
##         2m
## On a MxMxM grid with fixed boundary values on all sides. Note that
## V is a function of position, while E is a constant.
##
## The mesh points are numbered by the natural ordering.
function rv = build_schrodinger_FD_matrix (M, V)
  global hbar;
  global m;

  rv = spalloc (M^3, M^3, 7 * M^3);
  h  = 2 / M;
  C  = hbar^2 / (2*m);

  for i = 1:M^3
    [x y z] = from_matrix_coords (i, M);
    # Boundary conditions
    if (is_boundary (i, M))
      rv (i, i) = 1;
    # Non-boundary values get delsq values
    else
      rv (i, i)     = -6*C/h^2 + V(x, y, z);
      rv (i, i + 1) = C/h^2;
      rv (i, i - 1) = C/h^2;
      rv (i, i + M) = C/h^2;
      rv (i, i - M) = C/h^2;
      rv (i, i + M^2) = C/h^2;
      rv (i, i - M^2) = C/h^2;
    endif
  endfor;
endfunction


## Samples a function f on the mesh
## This is used for comparing exact solutions
function rv = build_known_solution (M, f)
  rv = zeros (M^3, 1);

  for i = 1:M^3
    [x y z] = from_matrix_coords (i, M);
    rv(i) = f(x, y, z);
  endfor
endfunction


# CO2 molecule -- two point charges
potential_co2 = inline ('(-k*e^2/sqrt(x^2 + y^2 + z^2) + -k*e^2/sqrt((x + 5.90355e-4)^2 + y^2 + z^2))/2', 'x', 'y', 'z');

# Set M = 44, since that's the limit of what my laptop can do
A = build_schrodinger_FD_matrix (44, potential_co2);

# To get a pi bond, render the first 4 or 5 eigenvectors (there is
# a lot of duplication, I suspect for symmetry reasons) and extract
# the first by doing "V = V * [ 1; 0; 0; 0; 0 ]".
[V E] = eigs (A, 1, -15);   # -15eV is just a swag

# Normalize solution
V = V ./ norm (V, 2);
# Output for render-me
V .* conj(V)