%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%% A 200 LINE TOPOLOGY OPTIMIZATION CODE FOR ELECTROMAGNETISM %%%%%%%%
% --------------------------- EXAMPLE GOAL ------------------------------ %
% Designs a 2D metalens with relative permittivity eps_r capable of       %
% monocromatic focusing of TE-polarized light at a point in space.        %
% --------------------- FIGURE OF MERIT MAXIMIZED ----------------------- %
%      Phi = |Ez|^2 in a "point" (in the center of a finite element)      %
% ------------------------- EQUATION SOLVED ----------------------------- %
%      \nabla * (\nabla Ez) + k^2 A Ez = F                                %
%      With first order absorping boundary condition:                     %
%      n * \nabla Ez = - i k Ez on boundaries                             %
%      and an incident plane wave propagating from bottom to top          %
% ---------------------- DOMAIN AND DISCRETIZATION ---------------------- %
% The equation is solved in a rectangular domain, discretized using       %
% quadrilateral bi-linear finite elements                                 %
% ----------------------------------------------------------------------- %
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%% Author: Rasmus E. Christansen, v2 April 2021 %%%%%%%%%%%%%%%%
% v2: updated method DENSITY_FILTER to correct sensitivity calculation    % 
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Disclaimer:                                                             %
% The authors reserves all rights but does not guaranty that the code is  %
% free from errors. Furthermore, we shall not be liable in any event      %
% caused by the use of the program.                                       %
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function [dVs,FOM]=top200EM(targetXY,dVElmIdx,nElX,nElY,dVini,eps_r,lambda,fR,maxItr)
% SETUP OF PHYSICS PARAMETERS
phy.scale = 1e-9; % Scaling finite element side length to nanometers
phy.eps_r = eps_r; % Relative permittivity
phy.k = 2*pi/(lambda*phy.scale); % Free-space wavenumber

% SETUP OF ALL INDEX SETS, ELEMENT MATRICES AND RELATED QUANTITIES
dis.nElX = nElX; % number of elements in x direction
dis.nElY = nElY; % number of elements in y direction
dis.tElmIdx = (targetXY(1)-1)*nElY+targetXY(2); % target index
dis.dVElmIdx = dVElmIdx; % design field element indices in model of physics
[dis.LEM,dis.MEM] = ELEMENT_MATRICES(phy.scale);
[dis]=INDEX_SETS_SPARSE(dis); % Index sets for discretized model

% SETUP FILTER AND THRESHOLDING PARAMETERS
filThr.beta = 5; % Thresholding sharpness
filThr.eta = 0.5; % Thresholding level
[filThr.filKer, filThr.filSca] = DENSITY_FILTER_SETUP( fR, nElX, nElY);

% INITIALIZE DESIGN VARIABLES, BOUNDS AND OPTIMIZER OPTIONS
dVs(length(dis.dVElmIdx(:))) = dVini; % Design variables
LBdVs = zeros(length(dVs),1); % Lower bound on design variables
UBdVs = ones(length(dVs),1); % Upper bound on design variables
options = optimoptions('fmincon','Algorithm','interior-point',...
    'SpecifyObjectiveGradient',true,'HessianApproximation','lbfgs',...
    'Display','off','MaxIterations',maxItr,'MaxFunctionEvaluations',maxItr);

% SOLVE DESIGN PROBLEM USING MATLAB BUILT-IN OPTIMIZER: FMINCON
FOM = @(dVs)OBJECTIVE_GRAD(dVs,dis,phy,filThr);
[dVs,~] = fmincon(FOM,dVs(:),[],[],[],[],LBdVs,UBdVs,[],options);

% FINAL BINARIZED DESIGN EVALUATION
filThr.beta = 1000;
disp('Black/white design evaluation:')
[FOM,~] = OBJECTIVE_GRAD(dVs(:),dis,phy,filThr);
end
%%%%%%%%%%%%%%% OBJECTIVE FUNCTION AND GRADIENT EVALUATION %%%%%%%%%%%%%%%%
function [FOM,sensFOM] = OBJECTIVE_GRAD(dVs,dis,phy,filThr)
% DISTRIBUTE MATERIAL IN MODEL DOMAIN BASED ON DESIGN FIELD
dFP(1:dis.nElY,1:dis.nElX) = 0; % Design field in physics, 0: air
dFP(dis.nElY:-1:ceil(dis.nElY*9/10),1:dis.nElX) = 1; % 1: material
dFP(dis.dVElmIdx(:)) = dVs; % Design variables inserted in design field

% FILTERING THE DESIGN FIELD AND COMPUTE THE MATERIAL FIELD
dFPS = DENSITY_FILTER(filThr.filKer,ones(dis.nElY,dis.nElX),filThr.filSca,dFP,ones(dis.nElY,dis.nElX));
dFPST = THRESHOLD( dFPS, filThr.beta, filThr.eta);
[A,dAdx] = MATERIAL_INTERPOLATION(phy.eps_r,dFPST,1.0); % Material field

% CONSTRUCT THE SYSTEM MATRIX
[dis,F] = BOUNDARY_CONDITIONS_RHS(phy.k,dis,phy.scale);
dis.vS = reshape(dis.LEM(:)-phy.k^2*dis.MEM(:)*(A(:).'),16*dis.nElX*dis.nElY,1);
S = sparse([dis.iS(:);dis.iBC(:)],[dis.jS(:);dis.jBC(:)],[dis.vS(:);dis.vBC(:)]);

% SOLVING THE STATE SYSTEM: S * Ez = F
[L,U,Q1,Q2] = lu(S); % LU - factorization
Ez = Q2 * (U\(L\(Q1 * F)));  Ez = full(Ez); % Solving

% FIGURE OF MERIT
P = sparse(dis.edofMat(dis.tElmIdx,:),dis.edofMat(dis.tElmIdx,:),1/4,...
    (dis.nElX+1)*(dis.nElY+1),(dis.nElX+1)*(dis.nElY+1)); % Weighting matrix
FOM = Ez' * P * Ez; % Solution in target element

% ADJOINT RIGHT HAND SIDE
AdjRHS = P*(2*real(Ez) - 1i*2*imag(Ez));

% SOLVING THE ADJOING SYSTEM: S.' * AdjLambda = AdjRHS
AdjLambda = (Q1.') * ((L.')\((U.')\((Q2.') * (-1/2*AdjRHS)))); % Solving

% COMPUTING SENSITIVITIES
dis.vDS = reshape(-phy.k^2*dis.MEM(:)*(dAdx(:).'),16*dis.nElX*dis.nElY,1);
DSdx = sparse(dis.iElFull(:),dis.jElFull(:),dis.vDS(:)); % Constructing dS/dx
DSdxMulV = DSdx * Ez(dis.idxDSdx); % Computing dS/dx * Field values
DsdxMulV = sparse(dis.iElSens,dis.jElSens,DSdxMulV);
sens = 2*real(AdjLambda(dis.idxDSdx).' * DsdxMulV); % Computing sensitivites
sens = full(reshape(sens,dis.nElY,dis.nElX));

% FILTERING SENSITIVITIES
DdFSTDFS = DERIVATIVE_OF_THRESHOLD( dFPS, filThr.beta, filThr.eta);
sensFOM = DENSITY_FILTER(filThr.filKer,filThr.filSca,ones(dis.nElY,dis.nElX),sens,DdFSTDFS);

% EXTRACTING SENSITIVITIES FOR DESIGNABLE REGION
sensFOM = sensFOM(dis.dVElmIdx);

% FMINCON DOES MINIMIZATION
FOM = -FOM; sensFOM = -sensFOM(:);

% PLOTTING AND PRINTING
figure(1); % Field intensity, |Ez|^2
imagesc((reshape(Ez.*conj(Ez),dis.nElY+1,dis.nElX+1))); colorbar; axis equal;
figure(2); % Physical design field
imagesc(1-dFPST); colormap(gray); axis equal; drawnow;
disp(['FOM: ' num2str(-FOM)]); % Display FOM value
end

%%%%%%%%%%%%%%%%%%%%%%%%%% AUXILIARY FUNCTIONS %%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%% ABSORBING BOUNDARY CONDITIONS AND RIGHT HAND SIDE %%%%%%%%%%%%
function [dis,F] = BOUNDARY_CONDITIONS_RHS(waveVector,dis,scaling)
AbsBCMatEdgeValues = 1i*waveVector*scaling*[1/6 ; 1/6 ; 1/3 ; 1/3];
% ALL BOUNDARIES HAVE ABSORBING BOUNDARY CONDITIONS
dis.iBC = [dis.iB1(:);dis.iB2(:);dis.iB3(:);dis.iB4(:)];
dis.jBC = [dis.jB1(:);dis.jB2(:);dis.jB3(:);dis.jB4(:)];
dis.vBC = repmat(AbsBCMatEdgeValues,2*(dis.nElX+dis.nElY),1);
% BOTTOM BOUNDARY HAS INCIDENT PLANE WAVE
F = zeros((dis.nElX+1)*(dis.nElY+1),1); % System right hand side
F(dis.iRHS(1,:)) = F(dis.iRHS(1,:))+1i*waveVector;
F(dis.iRHS(2,:)) = F(dis.iRHS(2,:))+1i*waveVector;
F = scaling*F;
end
%%%%%%%%%%%%%%%%%%%%% CONNECTIVITY AND INDEX SETS %%%%%%%%%%%%%%%%%%%%%%%%%
function [dis]=INDEX_SETS_SPARSE(dis)
% INDEX SETS FOR SYSTEM MATRIX
nEX = dis.nElX; nEY = dis.nElY; % Extracting number of elements
nodenrs = reshape(1:(1+nEX)*(1+nEY),1+nEY,1+nEX); % Node numbering
edofVec = reshape(nodenrs(1:end-1,1:end-1)+1,nEX*nEY,1); % First DOF in element
dis.edofMat = repmat(edofVec,1,4)+repmat([0 nEY+[1 0] -1],nEX*nEY,1);
dis.iS = reshape(kron(dis.edofMat,ones(4,1))',16*nEX*nEY,1);
dis.jS = reshape(kron(dis.edofMat,ones(1,4))',16*nEX*nEY,1);
dis.idxDSdx = reshape(dis.edofMat',1,4*nEX*nEY);
% INDEX SETS FOR BOUNDARY CONDITIONS
TMP = repmat([[1:nEY];[2:nEY+1]],2,1);
dis.iB1 = reshape(TMP,4*nEY,1); % Row indices
dis.jB1 = reshape([TMP(2,:);TMP(1,:);TMP(3,:);TMP(4,:)],4*nEY,1); % Column indices
TMP = repmat([1:(nEY+1):(nEY+1)*nEX;(nEY+1)+1:(nEY+1):(nEY+1)*nEX+1],2,1);
dis.iB2 = reshape(TMP,4*nEX,1);
dis.jB2 = reshape([TMP(2,:);TMP(1,:);TMP(3,:);TMP(4,:)],4*nEX,1);
TMP = repmat([(nEY+1)*(nEX)+1:(nEY+1)*(nEX+1)-1;(nEY+1)*(nEX)+2:(nEY+1)*(nEX+1)],2,1);
dis.iB3 = reshape(TMP,4*nEY,1);
dis.jB3 = reshape([TMP(2,:);TMP(1,:);TMP(3,:);TMP(4,:)],4*nEY,1);
TMP = repmat([2*(nEY+1):nEY+1:(nEY+1)*(nEX+1);(nEY+1):nEY+1:(nEY+1)*(nEX)],2,1);
dis.iB4 = reshape(TMP,4*nEX,1);
dis.jB4 = reshape([TMP(2,:);TMP(1,:);TMP(3,:);TMP(4,:)],4*nEX,1);
dis.iRHS = TMP;
% INDEX SETS FOR INTEGRATION OF ALL ELEMENTS
ima0 = repmat([1,2,3,4,1,2,3,4,1,2,3,4,1,2,3,4],1,nEX*nEY).';
jma0 = repmat([1,1,1,1,2,2,2,2,3,3,3,3,4,4,4,4],1,nEX*nEY).';
addTMP = repmat(4*[0:nEX*nEY-1],16,1);
addTMP = addTMP(:);
dis.iElFull = ima0+addTMP;
dis.jElFull = jma0+addTMP;
% INDEX SETS FOR SENSITIVITY COMPUTATIONS
dis.iElSens = [1:4*nEX*nEY]';
jElSens = repmat([1:nEX*nEY],4,1);
dis.jElSens = jElSens(:);
end
%%%%%%%%%%%%%%%%%% MATERIAL PARAMETER INTERPOLATION %%%%%%%%%%%%%%%%%%%%%%%
function [A,dAdx] = MATERIAL_INTERPOLATION(eps_r,x,alpha_i)
A = 1 + x*(eps_r-1) - 1i * alpha_i * x .* (1 - x); % Interpolation
dAdx = (eps_r-1)*(1+0*x) - 1i * alpha_i * (1 - 2*x); % Derivative of interpolation
end
%%%%%%%%%%%%%%%%%%%%%%%%%%% DENSITY FILTER  %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function [xS]=DENSITY_FILTER(FilterKernel,FilterScalingA,FilterScalingB,x,func)
xS = conv2((x .* func)./FilterScalingA,FilterKernel,'same')./FilterScalingB;
end
function [ Kernel, Scaling ] = DENSITY_FILTER_SETUP( fR, nElX, nElY )
[dy,dx] = meshgrid(-ceil(fR)+1:ceil(fR)-1,-ceil(fR)+1:ceil(fR)-1);
Kernel = max(0,fR-sqrt(dx.^2+dy.^2)); % Cone filter kernel
Scaling = conv2(ones(nElY,nElX),Kernel,'same'); % Filter scaling
end
%%%%%%%%%%%%%%%%%%%%%%%%%%%% THRESHOLDING %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function [ xOut ] = THRESHOLD( xIn, beta, eta)
xOut = (tanh(beta*eta)+tanh(beta*(xIn-eta)))./(tanh(beta*eta)+tanh(beta*(1-eta)));
end
function [ xOut ] = DERIVATIVE_OF_THRESHOLD( xIn, beta, eta)
xOut = (1-tanh(beta*(xIn-eta)).^2)*beta./(tanh(beta*eta)+tanh(beta*(1-eta)));
end
%%%%%%%%%%%%%%%%%%%%%%%%%%% ELEMENT MATRICES %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function [LaplaceElementMatrix,MassElementMatrix] = ELEMENT_MATRICES(scaling)
% FIRST ORDER QUADRILATERAL FINITE ELEMENTS
aa=scaling/2; bb=scaling/2; % Element size scaling
k1=(aa^2+bb^2)/(aa*bb); k2=(aa^2-2*bb^2)/(aa*bb); k3=(bb^2-2*aa^2)/(aa*bb);
LaplaceElementMatrix = [k1/3 k2/6 -k1/6 k3/6 ; k2/6 k1/3 k3/6 -k1/6; ...
    -k1/6  k3/6  k1/3  k2/6; k3/6 -k1/6  k2/6  k1/3];
MassElementMatrix = aa*bb*[4/9 2/9 1/9 2/9 ; 2/9 4/9 2/9 1/9 ; ...
    1/9 2/9 4/9 2/9; 2/9 1/9 2/9 4/9];
end