%  #===================================================================#
%  #         ---- Ett FEM-program för plana strukturer ----            #
%  #      ::: Små deformationer & linjär elasticitet antas :::         #
%  #                                                                   #
%  #   Av    : Carl Dahlberg, Institutionen för Hållfasthetslära, KTH  #
%  #   Datum : 20-Feb-2015                                             #
%  #   Comments translated to engling by Erik Olsson 2015-10-05
%  #===================================================================#
function [U,F,Sigma,numNodalDof] = SMALL_FE_PROG(varargin)
% ---
% ***  1 : Reading and checking the input from the user
% ---
if isempty(varargin)
	% Read input from an external file. Not presently implemented
	[Nodes,Elements,ElemType,MtrlProp,CrossSection,...
		PrescribedU,PointLoads,VolumeLoads] = READ_INPUT_FROM_FILE();
else
	% Input from the function call
	Nodes        = varargin{1};     % X och Y coordinates for the nodes        : [ numNodes , 2 ]
	Elements     = varargin{2};     % Matrix containing nodes of each elements : [ numElem , {antal noder/element} ]
	ElemType     = varargin{3};     % Element types                            :
	MtrlProp     = varargin{4};     % Material parameters for the elements     : [ numElem ]
	CrossSection = varargin{5};     % Cross-section properties                 : [ numElem ]
	PrescribedU  = varargin{6};     % Prescribed DoFs
	PointLoads   = varargin{7};     % Prescribed point loads
	VolumeLoads  = varargin{8};     % Distributed volume loads
end
numNodes = size(Nodes,1);
numElem  = size(Elements,1);
% Controll of input
[numNodalDof,numNodesElem,erFlag] = CHECK_INPUT(...
	numNodes,Elements,ElemType,MtrlProp,PrescribedU,PointLoads);
if erFlag
	% If erFlag =/= 0, terminate the program
	return
end
% ---
% ***  2 : Setting up vectors and matrices for storing the solution
% ---
numDof = numNodalDof * numNodes;  % Total number of DoFs
U = zeros(numDof,1);              % Deformation of the DOFs (displacements, rotations)
F = zeros(numDof,1);              % Nodal forces            (forces, moments)
Kmtrx = zeros(numDof,numDof);     %
% ---
% ***  3 : Loop over all elements
% ---
for e = 1 : numElem
	% --
	% **  3.1 : Establish an index-index verctor describing the mapping between
%               local and global equation numbers
	% --
	numDofElem = numNodesElem*numNodalDof;     % Antal (globala) frihetsgrader för elementet
	eqnNumLocal2Global = zeros(numDofElem,1);  % Lokalt -> Globalt ekvationsnummer (för assemblering)
	for i = 1 : numNodesElem
		n = Elements(e,i);                     % i:te nodnummret i element nummer e
		for j = 1 : numNodalDof
			eqnNumLocal2Global( (i-1)*numNodalDof + j ) = (n-1)*numNodalDof + j;
		end
	end
	% --
	% **  3.2 : ????
	% --
	ElemNodes = Nodes(Elements(e,:),:);
	[Tmtrx,len] = TRANSFORMATION_MATRIX(ElemType,ElemNodes);
	% --
	% **  3.3 : Read the local stiffness matrix for the element
	% --
	[ke] = ELEMENT_STIFFNESS_MATRIX(ElemType,MtrlProp(e),CrossSection(e),len);
	% --
	% **  3.4 : ?????
	% --
	Ke = Tmtrx' * ke * Tmtrx;
	% --
	% **  3.5 : Calculate contributions from possible volume forces
	% --
	[Fv] = VOLUME_FORCE(ElemType,VolumeLoads(e,:),Tmtrx,CrossSection(e),len);
	% --
	% **  3.6 : Assembly to global stiffness matrix
	% --
	Kmtrx(eqnNumLocal2Global,eqnNumLocal2Global) = ...
		Kmtrx(eqnNumLocal2Global,eqnNumLocal2Global) + Ke;
	% --
	% **  3.7 : ??????
	% --
	F(eqnNumLocal2Global) = F(eqnNumLocal2Global) + Fv;
end
% ---
% ***  4 : Point forces are added to the force vector
% ---
for i = 1 : size(PointLoads,1)
	n = PointLoads(i,1);           % Nod
	d = PointLoads(i,2);           % Frihetsgrad (riktning)
	f = PointLoads(i,3);           % Kraft (eller moment)
	eqn = (n-1)*numNodalDof + d;   % Bidraget ska adderas till position "eqn" i F, dvs F(eqn)
	F(eqn) = F(eqn) + f;
end
% ---
% ***  5 : Establishes reduced equation system
% ---
[Kred,Fred,U,ActiveDof] = BOUNDARY_CONDITIONS(Kmtrx,F,U,PrescribedU,numNodalDof,numDof);
% ---
% ***  6 : ??????
% ---
Ured = Kred\Fred;
% ---
% ***  7 : Substitute the solution to the reduced system into U
% ---
U(ActiveDof) = Ured;
F = Kmtrx*U;
% ---
% ***  8 : Calculate the stresses in each element
% ---
[Sigma] = EVALUATE_STRESS(Nodes,ElemType,Elements,numElem,numNodalDof,numNodesElem,MtrlProp,CrossSection,U);
end %SMALL_FE_PROG
% ------------------------------------------------------------------------
%
%
%       |  |             Help functions                     |  |
%       V  V                                                V  V
%
%
% ------------------------------------------------------------------------
function [Tmtrx,len] = TRANSFORMATION_MATRIX(ElemType,ElemNodes)
% The element is a truss/beam and needs to be transformed from the local
% description to the global one
len = norm(ElemNodes(2,:)-ElemNodes(1,:));
l12 = (ElemNodes(2,1)-ElemNodes(1,1))/len;
l21 = -(ElemNodes(2,2)-ElemNodes(1,2))/len;
m12 = (ElemNodes(2,2)-ElemNodes(1,2))/len;
m21 = (ElemNodes(2,1)-ElemNodes(1,1))/len;
% Define Tmtrx based on the element type
switch ElemType
	case 'BAR2'
		Tn = [l12 m12]; Zn = [0 0];
		Tmtrx = [Tn Zn; Zn Tn];
	case 'BAR3'
		Tn = [l12 m12]; Zn = [0 0];
		Tmtrx = [Tn Zn Zn; Zn Tn Zn; Zn Zn Tn];
	case 'BEAM2'
		Tn = [l21 m21 0; 0 0 1]; Zn = [0 0 0; 0 0 0];
		Tmtrx = [Tn Zn; Zn Tn];
	case 'FRAME2'
		Tn = [l12 m12 0; l21 m21 0; 0 0 1]; Zn = [0 0 0; 0 0 0; 0 0 0];
		Tmtrx = [Tn Zn; Zn Tn];
end
end %TRANSFORMATION_MATRIX

% ------------------------------------------------------------------------
function [ke] = ELEMENT_STIFFNESS_MATRIX(ElemType,MtrlProp,CrossSection,len)
switch ElemType
	case 'BAR2'
		A = AREA(CrossSection);
		EA = MtrlProp*A;
		ke = EA/len*[ 1 -1
			-1  1];
	case 'BAR3'
		A = AREA(CrossSection);
		EA = MtrlProp*A;
		a = len/2;
		ke = EA/a^2*[ 7/24  1/24 -1/3
			1/24  7/24 -1/3
			-1/3  -1/3   2/3];
	case 'BEAM2'
		I = AREA_MOMENT_OF_INERTIA(CrossSection);
		EI = MtrlProp*I;
		a = len/2;
		ke = EI/(2*a^3)*[3    3*a   -3    3*a
			3*a  4*a^2 -3*a  2*a^2
			-3  -3*a    3   -3*a
			3*a  2*a^2 -3*a  4*a^2];
	case 'FRAME2'
		A = AREA(CrossSection);
		EA = MtrlProp*A;
		I = AREA_MOMENT_OF_INERTIA(CrossSection);
		EI = MtrlProp*I;
		a = len/2;
		ke = [EA/(2*a)  0             0             -EA/(2*a) 0             0
			0         3*EI/(2*a^3)  3*EI/(2*a^2)  0         -3*EI/(2*a^3) 3*EI/(2*a^2)
			0         3*EI/(2*a^2)  2*EI/a        0         -3*EI/(2*a^2) EI/a
			-EA/(2*a) 0             0             EA/(2*a)  0             0
			0         -3*EI/(2*a^3) -3*EI/(2*a^2) 0         3*EI/(2*a^3)  -3*EI/(2*a^2)
			0         3*EI/(2*a^2)  EI/a          0         -3*EI/(2*a^2) 2*EI/a];
end
end %ELEMENT_STIFFNESS_MATRIX

% ------------------------------------------------------------------------
function [Fv] = VOLUME_FORCE(ElemType,VolumeLoads,Tmtrx,CrossSection,len)
A = AREA(CrossSection);
T = Tmtrx([1 2],[1 2]); % Enough for decribing the orientation between local (x,z) and global (X,Y)
K = T*VolumeLoads';
K_para = K(1); % Component parallel to the element
K_perp = K(2); % Component perpendicual to the element
switch ElemType
    case 'BAR2'
        fv = K_para*A*len/2*[1; 1];
    case 'BAR3'
        fv = K_para*A*len/6*[1; 1; 3];
    case 'BEAM2'
        fv = K_perp*A*len/2*[1; len/6; 1; -len/6];
    case 'FRAME2'
		fv = A*len/2*[K_para; K_perp; len*K_perp/6; K_para; K_perp; -len*K_perp/6];
end
Fv = Tmtrx'*fv;
end %VOLUME_FORCE

% ------------------------------------------------------------------------
function [Kred,Fred,U,ActiveDof] = BOUNDARY_CONDITIONS(Kmtrx,F,U,PrescribedU,numNodalDof,numDof)
Kred = Kmtrx;
Fred = F;
ActiveDof = 1:numDof;
% Looping over prescribed DOFs in order to move all known quantities to the
% right hand side in the equation system
numPresc = size(PrescribedU,1);
for i = 1:numPresc
	n = PrescribedU(i,1);        % Node
	d = PrescribedU(i,2);        % DOF
	eqn = (n-1)*numNodalDof + d; % Equation number
	Upresc = PrescribedU(i,3);
	Kcolumn = Kmtrx(:,eqn);
	Fred = Fred - Kcolumn*Upresc;
	U(eqn) = Upresc;
end
% Another loop over prescribed DOFs in order to reduce the equation system.
% Is made reversed over the prescribed DOFs in order to not remove the
% wrong lines/columns
for i = numPresc:-1:1
	n = PrescribedU(i,1);        % Node
	d = PrescribedU(i,2);        % DOF
	eqn = (n-1)*numNodalDof + d; % Equation number
	Kred(eqn,:) = [];    % Remove row in the stiffness materix K
	Kred(:,eqn) = [];    % Remove column in the stiffness materix K
	Fred(eqn) = [];      % Remove row in the force vector
	ActiveDof(eqn) = []; % Remove the equation number in the list of active equation numbers
end
end %BOUNDARY_CONDITIONS

% ------------------------------------------------------------------------
function A = AREA(CrossSection)
% Assumes a quadratic cross section
h = CrossSection;
A = h^2;
end %AREA

% ------------------------------------------------------------------------
function I = AREA_MOMENT_OF_INERTIA(CrossSection)
% Assumes a quadratic cross section
h = CrossSection;
I = h^4/12;
end %AREA_MOMENT_OF_INERTIA

% ------------------------------------------------------------------------
function [Sigma] = EVALUATE_STRESS(Nodes,ElemType,Elements,numElem,numNodalDof,numNodesElem,MtrlProp,CrossSection,U)
Sigma = zeros(numElem,1);     % Stress in each element

N_eval_points = 2;
XI = [-1/sqrt(3) 1/sqrt(3)];

for e = 1:numElem
	
	E = MtrlProp(e);
	
	ElemDofs = zeros(numNodalDof*numNodesElem,1);
	for i = 1:numNodesElem
		n = Elements(e,i); % i:th node nummer for element e
		for j = 1:numNodalDof
			ElemDofs((i-1)*numNodalDof + j) = (n-1)*numNodalDof + j;
		end
	end
	
	Ue = U(ElemDofs);
	
	ElemNodes = Nodes(Elements(e,:),:);
	[Tmtrx,len] = TRANSFORMATION_MATRIX(ElemType,ElemNodes);
	de = Tmtrx*Ue;
	switch ElemType
		case 'BAR2'
			% Constant stress - evaluation made in the centre point of the element
			xi = 1/2;
			[~,Bmtrx] = SHAPEFNC_AND_BMTRX(ElemType,xi,len);
			sigma_n = E*Bmtrx*de;
			Sigma(e) = sigma_n;
		case 'BAR3'
            % The stress can vary along the element - evaluation at several
            % points and choose the one with highest |stress|. Usually, the
            % stresses are evaluated in the Gauss points
			sigma_n = zeros(N_eval_points,1);
			for k = 1:N_eval_points
				xi = XI(k);
				[~,Bmtrx] = SHAPEFNC_AND_BMTRX(ElemType,xi,len);
				sigma_n(k) = E*Bmtrx*de;
			end
			k = find(sigma_n==max(abs(sigma_n)),1);
			Sigma(e) = sigma_n(k);
		case 'BEAM2'
			% Stress varies in both x- and z-dierction. It is chosen to
			% present the positive stress as a symmetric cross section is
			% assumed
			h = CrossSection(e);
			sigma_b = zeros(N_eval_points,1);
			for k = 1:N_eval_points
				xi = XI(k);
				[~,Bmtrx] = SHAPEFNC_AND_BMTRX(ElemType,xi,len);
				sigma_b(k) = E*abs(Bmtrx*de)*h/2;
			end
			k = find(sigma_b==max(abs(sigma_b)),1);
			Sigma(e) = sigma_b(k);
		case 'FRAME2'
            % The bending stress varies in both  x- and z-dierction. As an
            % overloaded normal stress from the normal deformation could be
            % present so is the largest summearized absolute value
            % presented with a sign in order to differentiate between
            % tenstion and compression
			h = CrossSection(e);
			sigma_b = zeros(N_eval_points,2);
			for k = 1:N_eval_points
				xi = XI(k);
				[~,Bmtrx] = SHAPEFNC_AND_BMTRX(ElemType,xi,len);
				sigma_b(k,:) = -E*Bmtrx([2 3 5 6])*de([2 3 5 6])*[h/2 -h/2];
			end
			sigma_n = E*Bmtrx([1 4])*de([1 4]); % Denna är konstant i elementet
			sigma_nb = sigma_n + sigma_b;
			[~,i] = max(abs(sigma_nb));
			sigma_nb = sigma_nb(i);
			[~,i] = max(abs(sigma_nb));
			Sigma(e) = sigma_nb(i);
	end
end
end %EVALUATE_STRESS

% ------------------------------------------------------------------------
function [Nmtrx,Bmtrx] = SHAPEFNC_AND_BMTRX(ElemType,xi,len)
% The elment type descides the matrices Nmtrx and Bmtrx
switch ElemType
	case 'BAR2'
		N1 = 1-xi;
		N2 = xi;
		Nmtrx = [N1 N2];
		dN1dxi = -1;
		dN2dxi = 1;
		Bmtrx = 1/len*[dN1dxi dN2dxi];
	case 'BAR3'
		N1 = xi*(xi-1)/2;
		N2 = xi*(xi+1)/2;
		N3 = 1-xi^2;
		Nmtrx = [N1 N2 N3];
		dN1dxi = (xi-1)/2 + xi/2;
		dN2dxi = (xi+1)/2 + xi/2;
		dN3dxi = -2*xi;
		Bmtrx = 1/len*[dN1dxi dN2dxi dN3dxi];
	case 'BEAM2'
		N1 = (2-3*xi+xi^3)/4;
		N2 = (1-xi-xi^2+xi^3)*len/4;
		N3 = (2+3*xi-xi^3)/4;
		N4 = (-1-xi+xi^2+xi^3)*len/4;
		Nmtrx = [N1 N2 N3 N4];
		B1 = 3*xi/(2*len^2);
		B2 = (3*xi-1)/(2*len);
		B3 = -3*xi/(2*len^2);
		B4 = (3*xi+1)/(2*len);
		Bmtrx = [B1 B2 B3 B4];
	case 'FRAME2'
		Nd1 = (1-xi)/2;
		Nd2 = (1+xi)/2;
		Nb1 = (2-3*xi+xi^3)/4;
		Nb2 = (1-xi-xi^2+xi^3)*len/4;
		Nb3 = (2+3*xi-xi^3)/4;
		Nb4 = (-1-xi+xi^2+xi^3)*len/4;
		Nmtrx = [Nd1 Nb1 Nb2 Nd2 Nb3 Nb4];
		Bd1 = -2/len;
		Bd2 = 2/len;
		Bb1 = 3*xi/(2*len^2);
		Bb2 = (3*xi-1)/(2*len);
		Bb3 = -3*xi/(2*len^2);
		Bb4 = (3*xi+1)/(2*len);
		Bmtrx = [Bd1 Bb1 Bb2 Bd2 Bb3 Bb4];
end
end %SHAPEFNC_AND_BMTRX

% ------------------------------------------------------------------------
function [numNodalDof,numNodesElem,erFlag] = CHECK_INPUT(...
	numNodes,Elements,ElemType,MtrlProp,PrescribedU,PointLoads)
erFlag = 0;
% The chosen element type will determine the number of inputs needed for
% different categories and how "big" the problem will be, both on element
% level and globally.
switch ElemType
	case 'BAR2'
		numNodalDof = 2;    % Two global DOFs/node: u_x and u_y
		numNodesElem = 2;   % Two nodes/element: xi=0 and xi=1
		numMtrlProp = 1;    % One material parameter: The Young's modulus 
	case 'BAR3'
		numNodalDof = 2;    % Two global DOFs/node: u_x and u_y
		numNodesElem = 3;   % Tree nodes/element: xi=-1, xi=1 and xi=0
		numMtrlProp = 1;    % One material parameter: The Young's modulus
	case 'BEAM2'
		numNodalDof = 3;    % Three global DOFs/node: u_x and u_y and theta_z
		numNodesElem = 2;   % Two nodes/element: xi=-1 and xi=1
		numMtrlProp = 1;    % One material parameter: The Young's modulus
	case 'FRAME2'
		numNodalDof = 3;    % Three global DOFs/node: u_x and u_y and theta_z
		numNodesElem = 2;   % Two nodes/element: xi=-1 and xi=1
		numMtrlProp = 1;    % One material parameter: The Young's modulus
	otherwise
		disp([' ERROR! Unknown element type, ' ElemType ', has been assigned.'])
		disp(' === Program is terminated. === ')
		disp(' ')
		erFlag = 1;
		return
end
% Controll the inputs so it agrees with the element type
if numMtrlProp ~= size(MtrlProp,2);
	disp([' ERROR! Wrong number of material parameters for the element type ' ElemType '.'])
	disp(['  You have written ' num2str(size(MtrlProp,2)) ', but it should be ' num2str(numMtrlProp) '.'])
	disp(' === Program is terminated. === ')
	disp(' ')
	erFlag = 1;
	return
end
if numNodesElem ~= size(Elements,2);
	disp([' FEROR! Wrong number of nodes for the element type ' ElemType '.'])
	disp(['  You have written ' num2str(size(Elements,2)) ', but it should be ' num2str(numNodesElem) '.'])
	disp(' === Program is terminated. === ')
	disp(' ')
	erFlag = 1;
	return
end
% Controll of the displacement BC given on the DOFs
if (any(PrescribedU(:,1)>numNodes)||any(PrescribedU(:,1)<1))
	disp(' FEROR! A DOF has been prescribed for a non-existing node.')
	disp(' === Program is terminated. === ')
	disp(' ')
	erFlag = 1;
	return
end
if (any(PrescribedU(:,2)>numNodalDof)||any(PrescribedU(:,2)<1))
	disp(' FEROR! A non existing DOF has been prescribed.')
	disp(' === Program is terminated. === ')
	disp(' ')
	erFlag = 1;
	return
end
% Controll of the point force BC
if isempty(PointLoads)
	% %No prescribed point forces -> no check
else
	if (any(PointLoads(:,1)>numNodes)||any(PointLoads(:,1)<1))
		disp(' FEROR! A point force is presctibed on a non-existing node.')
		disp(' === Program is terminated. === ')
		disp(' ')
		erFlag = 1;
		return
	end
	if (any(PointLoads(:,2)>numNodalDof)||any(PointLoads(:,2)<1))
		disp(' FEROR! A point force is presctibed in a direction (Fx, Fy, Mz) that not exist.')
		disp(' === Program is terminated. === ')
		disp(' ')
		erFlag = 1;
		return
	end
end
end %CHECK_INPUT