%atom.m
%Steven Sahyun, November 9, 2015
%University of Wisconsin Physics Dept.
%sahyuns@uww.edu
%Creates a graph of atomic radial electron probability densities useful for plotting to a 3D printer

clear;
clf();
xmax=200;
ymax=200;
x = 1:1:xmax;
y = 1:1:ymax;
xcenter = xmax/2;
ycenter = ymax/2;
meshgrid(x,y);

% Make a blank matrix
x0 = 1+0*x;
y0 = 1+0*y;

z = x0'*y0;

%Principle quantum number n
n = 3;
l = 2;
m = 1;

%Bohr radius
a0 = 7.5;

%radius in ratio of bohr radiua a0
rmax = xmax;
%rmax = 5*(2*n-1)*a0; %set scaling for complete picture
rstep = 1;

r = 0:rstep:rmax;
%theta is 0 to 360 degrees
theta = 0:pi/180:2*pi;
%phi is 0 to 180 degrees
phi = 0:pi/180:pi;

%Fill vectors with 1
R = 1+0.*r;
P = 1+0.*theta;
F = 1+0.*phi;

%Atomic number (of protons)
Z1 = 1;

rho = 2*Z1*r/n;

% Schrodinger equation for radial wave function
if n == 1
    %n = 1, l=0, m=0
    R = R.*(2/(a0^(3/2))).*exp(-r/(a0));
    P = P.*1/sqrt(2);
    F = F.*1/sqrt(2*pi);
    %R = (1/2*sqrt(6)).*rho.*Z1^(3/2).*exp(-rho/2);
elseif n == 2
    if l == 0
        %n=2, l=0, m=0
        R = R.*(1/(2*sqrt(2)*a0^(3/2))).*(2-r/a0).*exp(-r/(2*a0));
        P = P.*1/sqrt(2);
        F = F.*1/sqrt(2*pi);
    else
        if m == 0
            %n = 2, l = 1, m=0
            R = R.*(1/(2*sqrt(6)*a0^(3/2))).*(r/a0).*exp(-r/(2*a0));
            P = P.*sqrt(6).*cos(theta)/2;
            F = F.*1/sqrt(2*pi);
        else
            %n = 2, l = 1, m = ±1
            R = R.*(1/(2*sqrt(2)*a0^(3/2))).*(2-r/a0).*exp(-r/(2*a0));
            P = abs(P.*sqrt(6).*sin(theta)/2);
            F = F.*exp(i.*phi)/sqrt(2*pi);
        end
    end
elseif n == 3
    if l == 0
        %n=3, l=0, m=0
        R = R.*(2/(81*sqrt(3)*a0^(3/2)))*(27-18*r/a0+2*(r/a0).^2).*exp(-r/(3*a0));
        P = P.*1/sqrt(2);
        F = F.*1/sqrt(2*pi);
    elseif l == 1
        if m == 0
            %n=3, l=1, m=0
            R = R.*(4/(81*sqrt(6)*a0^(3/2)))*(6-r/a0)*(r/a0).*exp(-r/(3*a0));
            P = abs(P.*sqrt(6).*cos(theta)/2);
            F = F.*1/sqrt(2*pi);
        else
            %n=3, l=1, m=±1
            R = R.*(4/(81*sqrt(6)*a0^(3/2)))*(6-r/a0)*(r/a0).*exp(-r/(3*a0));
            P = abs(P.*sqrt(3).*sin(theta)/2);
            F = F.*exp(i.*phi)/sqrt(2*pi)
        end
    else
        if m == 0
            %n=3, l=2, m=0
            R = R.*4./(81.*sqrt(30).*(a0).^(3/2)).*(r.^2/a0.^2).*exp((-r)./(3*a0));
            P = abs(P.*sqrt(10)./4.*(3.*(cos(theta)).^2.-1));
            F = F.*1./sqrt(2*pi);
        elseif m == 1
            %n=3, l=2, m=±1
            R = R.*4./(81.*sqrt(30).*(a0).^(3/2)).*(r.^2/a0.^2).*exp((-r)./(3*a0));
            P = abs(P.*sqrt(15).*sin(theta).*cos(theta)/2);
            F = F.*exp(i.*phi)/sqrt(2*pi);
        else
            %n=3, l=2, m=±2
            R = R.*4./(81.*sqrt(30).*(a0).^(3/2)).*(r.^2/a0.^2).*exp((-r)./(3*a0));
            P = abs(P.*sqrt(15).*sin(theta)^(2)/4);
            F = F.*exp(i.*phi*2)/sqrt(2*pi);
        end
    end
    R = R;
end

%Psi = F.*P.*R;

%Radial Probability
Rprob = 4*pi*r.^2.*R.^2;

%Radial Probability graph
figure(1)
plot(r, Rprob);
title(['Radial Probability for Hydrogen n = ' num2str(n), ', l = ', num2str(l), ', m = ', num2str(m)]);

%Radial Probability surface
figure(2)

zoffset = 0;

% Set the probability for each matrix grid element.
for mx = 1:xmax
   for my = 1:ymax 
    z(mx,my) = z(mx,my)*Rprob(1+round(sqrt(abs(mx-xcenter).^2+abs(my-ycenter).^2)))-zoffset;
    %if z(mx,my) < 0.1 % If potential goes below 0, set value to 0; this can be considered the Schwartzchild radius.
    %    z(mx, my) = 0;
    %end
   end
end

%Plot the surface
surf(x,y,z);

%Angular probability  Surface
figure(3)

%Hydrogen2S = surf2solid(x, y, 75*z, 'Elevation', -1);
%stlwrite('Hydrogen2S.stl',Hydrogen2S)

% Set the Y*R probability for each matrix grid element.
for mx = 1:xmax
   for my = 1:ymax 
       %P will be in degrees. atan(my/mx) is in radians
    %z1 = z0*R Radial component
    z(mx,my) = z(mx,my)*Rprob(1+round(sqrt(abs(mx-xcenter).^2+abs(my-ycenter).^2)));
    %z2 = z1*zt Theta component
    z(mx,my) = z(mx,my)*P(1+abs(round(360/(2*pi)*atan((my-(ycenter+0.001))/(mx-(xcenter+0.002))))))-zoffset;
    %Psi = z2*zp Phi compnent
    
    %if z(mx,my) < 0.1 % If potential goes below 0, set value to 0; this can be considered the Schwartzchild radius.
    %    z(mx, my) = 0;
    %end
   end
end

surf(x,y,z);

%zscale = 40; %this is to set the z-axis scale to have a good plot, 10x10x2  
%z = z*zscale;
zscale =2904;
z = z*zscale;



%convert to a 3D printable file; You need to have stlwrite.m and surf2solid.m by
%Sven Holcombe in the same directory as this file.
%See: http://jmumakerlab.blogspot.com/2013/11/exporting-stl-from-matlab.html
% http://www.mathworks.com/matlabcentral/fileexchange/42876-surf2solid-make-a-solid-volume-from-a-surface-for-3d-printing
% http://www.mathworks.com/matlabcentral/fileexchange/20922-stlwrite-filename--varargin-

%functions
%Gravityfield = surf2solid(x, y, z, 'Elevation', 0);
%stlwrite('Gravityfield.stl',Gravityfield)
Hydrogen3Dm1= surf2solid(x, y, z, 'Elevation', -4.15);
stlwrite('Hydrogen3Dm1.stl',Hydrogen3Dm1)

