% Movie of all hydrogen orbitals (n,l,m), up to nmax, determined by the user. 
% Adapted from code by Peter van Alem. See 
% https://ch.mathworks.com/matlabcentral/fileexchange/64274-hydrogen-orbitals 

close all

vid = VideoWriter('hydrogenOrbitals.mp4','MPEG-4');
vid.Quality = 100;
vid.FrameRate = 30;
open(vid);

% Generate a sinusoidal ramp up and down of transparency when transitioning
% between orbitals 
x1 = (0:pi/15:14*pi/15);
fade1 = (-cos(x1)+1)/2; % Ramp up opacity
fade2 = ones(1,1); % One fully opaque image 
fade3 = (cos(x1+pi/15)+1)/2; % Ramp down opacity 
fade = [fade1 fade2 fade3]; % Concatenate 

% quantum numbers
% n
% 0 <= l < n
% -l <= m <= l
nmax = 4; % Maximum value of principle quantum number. Note that the number 
% of orbitals for a given n = n^2, and that the total number of orbitals 
% between n = 1 and nmax increases as nmax(nmax+1)(2nmax+1)/6, so for 
% nmax = 4, there are 30 orbitals, for nmax = 5, there are 55, etc, etc 

nOpaque = 20; % Number of fully opaque frames for a given orbital 

for n= 1:nmax
    for l = 0:n-1
        for m = -l:l
            for alphaVal = fade
                close all
                % plotting parameters
                probabilitydensity = 1.25e-5;
                a = 1;  % Bohr radius
                % angular part (Condon-Shortley)
                SphericalYlm = @(l,m,theta,phi)(-1)^m*sqrt((2*l+1)/(4*pi)* ...
                    factorial(l-abs(m))/factorial(l+abs(m)))* ...
                    AssociatedLegendre(l,m,cos(theta)).*exp(1i*m*phi);
                % real basis
                if (m < 0)
                    Y = @(l,m,theta,phi)sqrt(2)*(-1)^m* ...
                        imag(SphericalYlm(l,abs(m),theta,phi));
                elseif (m == 0)
                    Y = @(l,m,theta,phi)SphericalYlm(l,m,theta,phi);
                else
                    Y = @(l,m,theta,phi)sqrt(2)*(-1)^m*real(SphericalYlm(l,m,theta,phi));
                end
                % radial part
                R = @(n,l,r)sqrt((2/(a*n))^3*factorial(n-l-1)/(2*n*factorial(n+l))).* ...
                    exp(-r/(a*n)).*(2*r/(a*n)).^l*1/factorial(n-l-1+2*l+1) .* ...
                    AssociatedLaguerre(n-l-1,2*l+1,2*r/(a * n));
                % wave function
                psi = @(n,l,m,r,theta,phi)R(n,l,r).*Y(l,m,theta,phi);
                % setting the grid
                border = 32; 
                % "accuracy" slows program approximately with the cube of its value! 
                % Bump up to 500 only for n = 2, l = 0, due to artefact of
                % coloring of inner sphere. Just being super OCD here 
                if (n == 2) && (l ==0) 
                    accuracy = 500; % Set to 200 if you're not bothered by artefact 
                else
                    accuracy = 200; 
                end
                raster = linspace(-border,border,accuracy);
                [x, y, z] = ndgrid(raster,raster,raster);
                % conversion Cartesian to spherical coordinates
                r = sqrt(x.^2+y.^2+z.^2);
                theta = acos(z./r);
                phi = atan2(y,x);
                % plot orbital,  - and + wave function phase
                colors = sign(psi(n,l,m,r,theta,phi));
                isosurface(psi(n,l,m,r,theta,phi).^2,probabilitydensity,colors);
                alpha(0.8*alphaVal);
                colormap([0 0 1; 1 0.5 0])
                material shiny
                
                title(['$ n = ', num2str(n), ', l = ', num2str(l), ...
                    ', m = ', num2str(m) '$'], 'interpreter', 'latex', ...
                    'FontSize', 20)
                set(gcf,'color','w');
                set(gca,'CameraViewAngle',45,'Projection','perspective');
                camzoom(4)
                axis equal
                axis vis3d
                axis off % Comment out if you want axes shown. I find them unnecessary 

                % I find labels ugly so commented them out. You can
                % uncomment as you see fit 
                % xticklabels('');
                % yticklabels('');
                % zticklabels('');
                % xlabel('$x$', 'interpreter', 'latex', 'FontSize', 20)
                % ylabel('$y$', 'interpreter', 'latex', 'FontSize', 20)
                % zlabel('$z$', 'interpreter', 'latex', 'FontSize', 20)
                % Store the frame
                if (alphaVal == 1) % When full opaque, record nOpaque frames  
                    for ii = 1:nOpaque
                        frame = getframe(gcf);
                        writeVideo(vid,frame);
                    end
                else
                    frame = getframe(gcf);
                    writeVideo(vid,frame);
                end                
            end
        end
    end
end
% Output the movie as an mpg file
close(vid);

% Functions 
% _________
function Anm = AssociatedLaguerre(n,m,x)
Anm = 0;
for i = 0 : n
    Anm = Anm+factorial(m+n)*nchoosek(m+n,n-i)/factorial(i)*(-x).^i;
end 
end

function Alm = AssociatedLegendre(l,m,x)
Alm = 0;
for r = 0:floor(1/2 * l-1/2 * abs(m))
    Alm = Alm + (-1)^r * nchoosek(l-2*r,abs(m))*nchoosek(l,r) * ...
        nchoosek(2*l-2*r,l)*x.^(l-2*r-abs(m));
end
Alm = (1-x.^2).^(abs(m)/2).*(factorial(abs(m))/2^l*Alm);
end