%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% MATLAB code which solves the Strongly Nonlinear version of Density      %
% Functional Theory (DFT) for the Face-Centred cubic (FCC) crystal in the %
% case of Generalised Exponential of power 4 (GEM-4) interations          %
%                                                                         %
% The paper this work is based on appears in:                             %
%   P. Subramanian et. al. "The density distribution in soft matter       %
%   crystals and quasicrystals", Phys. Rev. Lett, ???????                 %
% and its supplementary material, which you should cite if you use        %
% this code or an altered version.                                        %
%                                                                         %
% This code was written by Daniel Ratliff and Priya Subramanian, so for   %
% questions on this code or corrections/optimisations, please don't       %
% hesitate to drop us a line at danieljamesratliff@gmail.com              %
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

%% 0. Assigning variables for computational domain and branch specifications

%We start by setting the box parameters, including the size of the cell
% (Lx,Ly,Lz) and the resolution (Nx,Ny,Nz). We also save a copy of the
% initial size of the cell (Lx0,Ly0,Lz0), since the box size will change
% under minimization. We choose double the natural periodic cell for an FCC
% crystal and the resolution to be 12 points per unit

Lx =  2*sqrt(3);        Ly = 2*sqrt(3);         Lz = 2*sqrt(3);
Lx0 = Lx;               Ly0 = Ly;               Lz0 = Lz;
Nx = round(12*Lx);      Ny = round(12*Ly);      Nz = round(12*Lz);

%Add these to a structure, Boxparams, so that it can be passed a lot easier into
%functions

Boxparams.Lx0 = Lx; Boxparams.Lx = Lx;  Boxparams.Nx = Nx;
Boxparams.Ly0 = Ly; Boxparams.Ly = Ly;  Boxparams.Ny = Ny;
Boxparams.Lz0 = Lz; Boxparams.Lz = Lz;  Boxparams.Nz = Nz;

% Here we choose the order of the Strongly Nonlinear theory (SNLT) - namely, how
% many modes are to be considered in the expansion (1) in the Supplementary
% material
order = 3;

%We also load a table of pair potential values for the GEM-4 potential from
%the GEM-4_data file. Note: this could be generated in this code if one so
%pleased
GEM4 = load('GEM4_data.mat','mat');
VectorStructure.potential = GEM4.mat;

% We now set the solution branch details. The branch is computed following
% the amplitude of the |k|=1 mode phi_1, as this is monotonic unlike the chemical
% potential mu, meaning that arclength continuation is not required.
% Therefore, we set the following:

stepsize = 0.01;        %spacing of phi_1 values on the branch
nosteps = 100;          %how many steps along the branch we go
startval = 0.005;       %The starting point on the branch

% Therefore, the branch will go from phi_1 = startval to 
% phi_1 =startval+(nosteps-1)*stepsize

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

%% 1. Defining and generating the Reciprocal Lattice Vectors (RLV)

%In this section, we will be creating the set of reciprocal lattice vectors
%that the SNLT will utilise. The theory underpinning this section can be
%found in the supplementary material in paragraph 3 of the implementation.

% We start by selecting a basis set of vectors, the Principle Lattice
% Vectors (PLVs). In this code, the FCC crystal is selected and has the
% basis as written in equation (S2) (including the bulk mode represented by the 0 vector). 
% For book-keeping, we also assign the "order" of each vector according to
% the theory
% So, each row contains [x-component  y-component   x-component   order] 

RLV =  [0                 0                0             0 ;
    1       1       1        1 ;
    1       1      -1        1 ;
    1      -1      -1        1 ;
    1      -1       1        1 ;
    -1      -1      -1        1 ;
    -1      -1       1        1 ;
    -1       1      1        1 ;
    -1       1      -1        1 ];


%Now, add all the combinations of these, and to do this we need to extract
%the first modes to add at each order, these will then be used to generate
%the subsequent orders of the RLV lattice, as defined in the supplementary
%material

Order1Mat = permute(RLV(2:end,:),[3,2,1]);

% The following loop performs the following iteration scheme, starting with
% the PLVs:
% 1) To each lattice vector labelled with the highest order value,
% add each of the nonzero PLVs to it
% 2) from all of these results, retain only the unique results which differ
% from the set of lattice vectors you started with
% 3) Add these new lattice to the list of lattice vectors, and label these
% with an order value one higher than the previous order
%
% This iteration continues until the set of lattice vectors of the desired
% order are obtained
for order_ind = 1:order-1
    
    %Find the modes of the previous order
    [rowind,~] = find(RLV(:,end)==order_ind);
    
    %Add to each of these the nonzero PLVs, to get the next order terms. The 8
    %here is because there are 8 non-zero RLV vectors 
    Sum_Mat = repmat(RLV(rowind,:),[1,1,8])+Order1Mat;
    
    %Reshape these into the same format as the RLV list
    Sum_Mat_reshaped = reshape(permute(Sum_Mat,[2,1,3]),size(Sum_Mat,2),[])';
    
    %A lot of these sums will cause repeats of a) previous lattice modes
    %and b) of modes computed at this order. Remove these using unique to
    %remove the repeats (ignoring the order values)
    ScrubMat = unique(cat(1,RLV(:,1:end-1),Sum_Mat_reshaped(:,1:end-1)),'rows','stable');
    
    %Then, skim off the old ones (which will have all be unique by
    %construction) and extract those at the bottom of the scrubbed matrix -
    %these will be your next order terms
    Mode_append = ScrubMat(size(RLV,1)+1:end,:);
    
    Mode_append = uniquetol(Mode_append,'ByRows',true);
    
    %Label these as the next order terms
    Mode_append(:,end+1) = order_ind+1;
    
    %Add these to the RLV list
    RLV = cat(1,RLV,Mode_append);
    
end

%Normalise the RLVs post-computation
RLV(:,1:3) = RLV(:,1:3)/sqrt(3);

%Clear everything that isn't needed going forward
clear  Order1Mat rowind Sum_Mat Sum_Mat_reshaped ScrubMat Mode_append

%The resulting list of RLVs is added to a structure which will be used to
%pass vector and vector lists into functions, called VectorStructure
VectorStructure.RLV = RLV;

%% 2. Preallocation

% The remainder of the components of VectorStructure, namely the wavenumbers,
%  the required values of the pair potential for the  Strongly Nonlinear 
% equations and are generated using the
% Box_dependent_updates function, since certain parts of the problem will
% need to be recomputed as the box size changes - namely, the active
% wavenumbers and thus the relevant values of the pair potential
VectorStructure = Box_dependent_Updates(VectorStructure,1);

%Extract these results
PairPotValues = VectorStructure.PairPotValues;

Waveno_Vect = VectorStructure.Waveno_Vect;

%Extract the points which are, in principle, the first order points
[phi_index,~] = find(Waveno_Vect(:,end)==1,1);

%The aim now is to generate the sums appearing in (S1) of the supplementary
%material. We will do this by creating a 4D array whose j-th page (i.e.
%the last index) is the j-th sum, which will be multiplied by the
%appropriate amplitude. So, we start by creating an empty array
CosVector = zeros(Nx,Ny,Nz,size(Waveno_Vect,1));

%Create a spatial grid for this calculation
[X,Y,Z] = ndgrid((0:Nx-1)*Lx*2*pi/Nx,(0:Ny-1)*Ly*2*pi/Ny,(0:Nz-1)*Lz*2*pi/Nz);

%Now, extract the wavenumbers from the RLV and use them to generate the
%desired cosines, for the sum defined in (S1):
%M = phi0+phi1*(cosines of waveno. 1)+phisqrt(4/3)*(cosines of waveno.
%sqrt(4/3))+...
Wavevectors = struct('Kset', cell(1, size(Waveno_Vect,1)));

for ii=1:size(Waveno_Vect,1)
    
    %Find wavenumbers of a given size, using the unique wavenumber vector
    %generated a few lines ago. As these end up being nonintegers, we
    %identify things within the SAME TOLERANCE as uniquetol (10^-4)
    
    Wavevectors(ii).Kset(:,1:4) =...
        RLV(abs(RLV(:,1).^2+RLV(:,2).^2+RLV(:,3).^2 - Waveno_Vect(ii,1))<1E-4,:);
    
    Wavevectors(ii).Kset(:,5) = Waveno_Vect(ii,1);
    
    for jj = 1:size(Wavevectors(ii).Kset,1)
        
        %For each of these wavevectors, we add cos(k.*r)
        %to the current page of the 4D array, so the result on each page is
        %the sum of complex exponentials for wavevectors of a given length
        %and thus the sums like those appearing in (S1)
        
        CosVector(:,:,:,ii) = CosVector(:,:,:,ii)+...
            cos(Wavevectors(ii).Kset(jj,1)*X+ ...
            Wavevectors(ii).Kset(jj,2)*Y+...
            Wavevectors(ii).Kset(jj,3)*Z);
        
    end
end

%Add this collection of cosines to VectorStructure

VectorStructure.CosVector = CosVector;

%again, clear everything not needed going forward
clear X Y Z

%With the initial cell parameters and the necessary quantities for
%the SNLT equations to be constructed defined, we can define the equations
%to solve as an anonymous function of the amplitudes and the phi_1
%amplitude which will remain fixed. The function which generates the SNLT
%equations can be found at the end of this script

SNLT_Eqtns_1 = @(phi_mu,phi_1)(SNLT_Eqtns(phi_mu,phi_1,VectorStructure,phi_index));

%Pre-allocate the save vectors for the solution branch outputs:
% solbranch - stores the amplitude and chemical potential at each branch
% point
% swell - stores the ratio between the crystal's ideal cell size and the
% original cell
% GP_Liquid - stores the value of the Grand potential for the liquid that
% corresponds to the branch's mu value
% GP_Crystal - stores the Grand potential for the crystal for the branch
% rho_bar - stores the average density of the crystal
% rho_liquid - stores the density of the liquid which corresponds to the
% value of mu at the branch point
solbranch = zeros(size(Waveno_Vect,1)+1,nosteps);
swell = zeros(1,nosteps);  
GP_Liquid = zeros(1,nosteps);       GP_Crystal = zeros(1,nosteps);
rho_bar = zeros(1,nosteps);      rho_liquid = zeros(1,nosteps);


% We set here some options for the solvers used in the branch part of the
% code, one for the SNLT system and one for finding the value of the liquid
% at a given value of mu
SNLT_options = optimoptions(@fsolve,'display','off');

rho_options = optimoptions(@fsolve,'display','off');

%% 3. Branch Computation for the SNLT system

% In this section, we compute the solution branch as determined by the SNLT
% approach - we do this using a loop which does the following
% 1) computes the solution for the amplitudes phi and chemical potential mu
% for a fixed amplitude phi_1 in the original periodic cell
% 2) Ideal box size is found by minimizing the Grand potential with respect
% to box size
% 3) Solution is recomputed in this box

tic

for loop_ind=1:nosteps
    
    %Fix a phi_1 amplitude determined by the loop
    phi1 = startval+(loop_ind-1)*stepsize;
    
    %Use this value of phi to fix the equations that we are solving for each branch
    SNLT_NoSwell = @(phi_mu)(SNLT_Eqtns_1(phi_mu,phi1));
    
   
    
    if loop_ind==1
        % The first initial guess uses the results from weakly nonlinear theory
        % to determine what the bulk mode phi_0 and the chemical potential
        % mu (denoted by phi_guess(end))
        phi_guess = zeros(size(Waveno_Vect,1)+1,1);
        
        phi_guess(1) = log(-1/PairPotValues(phi_index,2))-phi1;
        
        phi_guess(end) = log(-1/PairPotValues(phi_index,2))-...
            1/PairPotValues(phi_index,2)*PairPotValues(1,2)-...
            phi1*(PairPotValues(1,2)/PairPotValues(phi_index,2)-1);
        
    else
        phi_guess = phi_sol;
    end
    
    % Set the phi_1 amplitude the the fixed value
    phi_guess(phi_index) = phi1;
    
    % Use fsolve to solve the problem in the original box
    [phi_sol] = fsolve(SNLT_NoSwell,phi_guess,SNLT_options);
    
    % We have a solution, however we must now use this to find the ideal
    % box size for the crystal at this point of the branch. To do so, we
    % generate the crystal in physical space:
    
    M = 0;
    
    for ii=1:size(phi_sol,1)-1
        M = M+phi_sol(ii)*CosVector(:,:,:,ii);
    end
    
    % This is then fed into a routine (see repeated functions at the end of
    % this script) that, given a value for swell (namely, the ratio
    % between the ideal cell size that minimizes the Grand potential
    % and the original cell) computes the Grand potential
 
    
    GP_Func = @(swell)(GP_GEM4(M,phi_sol(end),Boxparams,swell));
    
    % Find the value of the swell which minimizes the above function. 
    % Here we use fminbd for ease and to subvert problems on the unstable
    % parts of the branch where other minimizer functions seem to fail
    
    SwellVal = fminbnd(GP_Func,0.99,1.05);
    
    % Store this swell value in a save vector
    swell(loop_ind) = SwellVal;
    
    % % Update the box using this swell value
    Boxparams.Lx = Lx0*SwellVal;    Boxparams.Ly = Ly0*SwellVal;
    Boxparams.Lz = Lz0*SwellVal;
    
    % Update the wavenumbers and relevant pair potential values for the SNLT
    % equations for the new cell size
    
    VectorStructure = Box_dependent_Updates(VectorStructure,SwellVal);
    
    % Define the SNLT equations to solve as a function of the phi_mu vector
    % for this new cell
    ToSolve = @(phi_mu)(SNLT_Eqtns(phi_mu,phi1,VectorStructure,phi_index));
    
    %Solve the SNLT using fsolve
    [phi_sol] = fsolve(ToSolve,phi_sol,SNLT_options);
    
    %Add the obtained amplitudes and chemical potential mu to save matrix
    solbranch(:,loop_ind) = phi_sol;
    
    %To extract physical and thermodynamic quantities, we construct the
    %crystal for this solution
    
    M = 0;
    
    for ii=1:size(phi_sol,1)-1
        M = M+phi_sol(ii)*CosVector(:,:,:,ii);
    end
    
    % Extract the average density for this crystal
    rho_bar(loop_ind) = sum(exp(M(:)))/numel(M);
    
    % Compute the Grand Potential
    GP_Crystal(loop_ind) = GP_GEM4(M,phi_sol(end),Boxparams,SwellVal);
    
    % Compute the density of the liquid at the obtained value of mu from the SNLT - on
    % the first loop, we use the density at onset, -1/Vhat(1), as an
    % initial guess and the previous liquid density on all subsequent loops
    if loop_ind==1
        
        rho_liquid(loop_ind) = fsolve(@(x)(log(x)+PairPotValues(1,2)*x)-solbranch(end,loop_ind),...
            -1/PairPotValues(phi_index,2),rho_options);
        
    else
        
        rho_liquid(loop_ind) = fsolve(@(x)(log(x)+PairPotValues(1,2)*x)-solbranch(end,loop_ind),...
            rho_liquid(loop_ind-1),rho_options);
        
    end
    
    %Calculate the Grand Potential of the liquid
    GP_Liquid(loop_ind) = rho_liquid(loop_ind)*(log(rho_liquid(loop_ind))-1)+...
                                    0.5*rho_liquid(loop_ind)^2*PairPotValues(1,2)-...
                                    solbranch(end,loop_ind)*rho_liquid(loop_ind);
    
end

toc

%% 4. Plots

% For nondimensional plotting, you need to use R^3 = 5.571858^3, emerging
% from the leengthscale of the GEM-4 pair potential
R_cubed = 5.571858^3;

%Figure 1 plots the nondimensional average density of the crystal and liquid densities
%against mu
figure(1);
hold on;
plot(solbranch(end,:),(rho_bar)*R_cubed,'r','linestyle','-');
plot(solbranch(end,:),rho_liquid*R_cubed,'k','linestyle','-');
hold off;
title('rho branch');
box on;
axis tight
xlabel('$\mu$','Interpreter','Latex');
ylabel('$\rho*R_s^3$','Interpreter','Latex');
drawnow;

% figure 2 plots the logarithm (base e) of the nondimensional average
% density and liquid density, again, against mu
figure(2);
hold on;
plot(solbranch(end,:),log(rho_bar*R_cubed),'r','linestyle','-','linewidth',2);
plot(solbranch(end,:),log(rho_liquid*R_cubed),'k','linewidth',1.5);
hold off;
title('log(rho) branch');
box on;
axis tight
xlabel('$\mu$','Interpreter','Latex');
ylabel('$\log(\rho*R_s^3)$','Interpreter','Latex');
drawnow;

% figure 3 plots the nondimensional Grand potential of the crystal and
% liquid phases against mu
figure(3);
hold on;
plot(solbranch(end,:),GP_Crystal*R_cubed,'r','linestyle','-','LineWidth',2);
plot(solbranch(end,:),GP_Liquid*R_cubed,'k','linewidth',1.5);
hold off;
box on;
axis tight
xlabel('$\mu$','Interpreter','Latex');
ylabel('$\frac{R_s^3\Omega}{V}$','Interpreter','Latex');
drawnow;

figure(4);
hold on;
plot(solbranch(end,:),(GP_Crystal-GP_Liquid)*R_cubed,'color','r','linestyle','-',...
    'LineWidth',2);
hold off;
set(gca,'xaxislocation','origin')
axis tight;
xlabel('$\mu$','Interpreter','Latex');
ylabel('$\frac{R_s^3(\Omega-\Omega_{liq})}{V}$','Interpreter','Latex');
drawnow;


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%              Repeatedly Used functions can be found here:              %
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

%% Box_Dependent_Updates:
function VectorStructure = Box_dependent_Updates(VectorStructure,swell)

RLV = VectorStructure.RLV;

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

%Extract the unique wavenumbers. These are done as squared
%wavenumber to keep things rational (and to save heartache later when things
%round badly, when unique operations break down, when it tries to feed a
%rational number into an index, etc.)
Waveno_Vect = uniquetol([(RLV(:,1)/swell).^2+(RLV(:,2)/swell).^2 + (RLV(:,3)/swell).^2,...
                                                RLV(:,4)],1E-5,'ByRows',true);

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%    2. GET PAIR POTENTIAL VALUES AT RLV!!! %%%%
%Make the pair potential in fourier space and steal the relevant values for the RLV points
%using this rather handy function I wrote

PairPotValues = zeros(size(Waveno_Vect,1),3);
Potential = VectorStructure.potential;

%Preallocate the data in column 1 (wavenumber squared) and column 3 (the
%order of the wavenumber on the lattice)

PairPotValues(:,1) = Waveno_Vect(:,1);     PairPotValues(:,3) = Waveno_Vect(:,2);

%The top element is always the zero wavenumber, so in that goes

PairPotValues(1,2) = Potential(1,2);

% Now, to get the rest, we extract from the potential in Fourier space:

for jj=2:size(Waveno_Vect,1)
    %Within the imported data, find the closest wavenumber to the one you
    %need
    
    
    GEM_ind = find((Potential(:,1).^2-Waveno_Vect(jj,1))>0,1);
    
    %...and allocate using linear interpolation, it's dat easy (geddit? because it's a .dat file?
    %huehuehue)
    
    PairPotValues(jj,2) = Potential(GEM_ind-1,2)+...
        (Potential(GEM_ind-1,1)-sqrt(Waveno_Vect(jj,1)))/...
        (Potential(GEM_ind-1,1)-Potential(GEM_ind,1))*...
        (Potential(GEM_ind,2)-Potential(GEM_ind-1,2));
    
end

%For neatness, put the expo structures into one vectorstructure, so that
%passing it into functions is more of a doddle
VectorStructure.Waveno_Vect = Waveno_Vect;      VectorStructure.PairPotValues = PairPotValues;

end

function SNLT = SNLT_Eqtns(phi_mu_vect,phi1_val,VectorStructure,phi_index)

%Unpack the parts of VectorStructure necessary for the computations here,
%namely the pair potential values and CosVector
PairPot = VectorStructure.PairPotValues;
CosVector = VectorStructure.CosVector;

%Do the fourier integral of the exponential of M:

% %Make M: Start with phi0 and then add the cosines, indexed by the grid -
% the -1 here is because the very last element of phi_mu_vect is, in fact,
% mu and not an amplitude
M = 0;
for ii=1:length(phi_mu_vect)-1
    M = M+phi_mu_vect(ii)*CosVector(:,:,:,ii);
end

%Exponentiate
exp_M = exp(M);

%Pre-allocate the array of the necessary Fourier Transform values of exp(M)
rho_hat_vector = zeros(length(PairPot),1);

% For each 
for ii=1:length(rho_hat_vector)
    
    FTerm = exp_M.*CosVector(:,:,:,ii);
    
    rho_hat_vector(ii) = sum(FTerm(:))/(numel(M))/CosVector(1,1,1,ii);
    
end

%Generate the equation (S2) for each of the modes
SNLT(2:length(phi_mu_vect)) =  phi_mu_vect(1:end-1)+PairPot(:,2).*rho_hat_vector;


%Add the chemical potential to the very first SNLT equation, to generate
%equation (S3) from the supplementary material
SNLT(2) = SNLT(2)-phi_mu_vect(end);

%The very first element of the vector SNLT is the constraint
% on the value of phi1
SNLT(1) =phi1_val -  phi_mu_vect(phi_index);
end


function GP_Phi = GP_GEM4(M,mu,Boxparams,swell)
% This subroutine calculates the Grand Potential
Lx = Boxparams.Lx0*swell;   Ly = Boxparams.Ly0*swell;   Lz = Boxparams.Lz0*swell;

Nx = Boxparams.Nx;  Ny = Boxparams.Ny;  Nz = Boxparams.Nz;

%Need a spatial grid, because who doesn't like a spatial grid?
[X,Y,Z] = ndgrid((0:Nx-1)*Lx*2*pi/Nx,(0:Ny-1)*Ly*2*pi/Ny,(0:Nz-1)*Lz*2*pi/Nz);

R2 = (X-Lx*pi).^2+(Y-Ly*pi).^2+(Z-Lz*pi).^2;

GEM4 = exp(-R2.^2/(5.571858)^4);

rho = exp(M);

Vhat = fftn(GEM4)*Lx*Ly*Lz*(2*pi)^3/(Nx*Ny*Nz);

convo_term = real(ifftn(fftn(rho).*Vhat));

GP_Term = (log(rho)-1+0.5*convo_term-mu).*rho;

% all things ready, now GP
GP_Phi = sum(GP_Term(:))/(Nx*Ny*Nz);

end

