%-----------------------------------------------------------------
% (former Exercise #5 'Discretization Methods')
% Calculation of a truss structure with the Finite Element Method
% using linear shape functions. In addition to the matlab file
% "truss_structure.m", mass effects are considered and time
% integration is carried out using the Central Difference Method
% or Newmark's method.
% (c) 01/2007 by A. Schmidt, IAM, Uni Stuttgart
%-----------------------------------------------------------------

%-----------------------------------------------------------------
% remark: with this program an arbitrary 2d truss structure with
%         any number of nodes and elements which might have 
%         different Young's moduli, densities and cross section 
%         areas can be calculated.
%         Thus, the user has to input the location of all nodes,
%         the definition of all elements and the boundary 
%         conditions (forces and fixed nodes).
%-----------------------------------------------------------------
clear all


% model definition
% ----------------
% definition of all (global) nodes by their x- and y-coordinates
nod_x(1) = 0.0;
nod_y(1) = 1.0;
nod_x(2) = 0.0;
nod_y(2) = 0.0;
nod_x(3) = 1.0;
nod_y(3) = 1.0;
nod_x(4) = 1.0;
nod_y(4) = 0.0;
nod_x(5) = 2.0;
nod_y(5) = 0.0;

% definition of all (rod-) elements by their nodes
elem_nod1(1) = 1;
elem_nod2(1) = 2;
elem_nod1(2) = 1;
elem_nod2(2) = 3;
elem_nod1(3) = 2;
elem_nod2(3) = 3;
elem_nod1(4) = 2;
elem_nod2(4) = 4;
elem_nod1(5) = 3;
elem_nod2(5) = 4;
elem_nod1(6) = 3;
elem_nod2(6) = 5;
elem_nod1(7) = 4;
elem_nod2(7) = 5;


% definition of material
% ----------------------
% Young's modulus of all elements
E(1:7) = 2.1e+11;

% mass density of all elements
rho(1:7) = 7.8e+3;

% cross section area of all elements
A(1:7) = 1.0e-4;


% definitions for time integration
% --------------------------------
t_tot = 0.04;                       % total integration time
dt = 0.00015;                       % time step size
n_dt = floor(t_tot/dt);             % resulting number of time integration steps
alpha = 0.25;                       % Newmark Parameter alpha
delta = 0.6;                        % Newmark Parameter delta


% boundary conditions
% -------------------
% definition of all fixed nodes
fixed(1) = 1;
fixed(2) = 2;

% definition of external loads
nod_force(1) = 5; 
force_x(1:n_dt+1) = 0.0;
force_y(1:n_dt+1) = -100.0;


% definition of output variables
% ------------------------------
fac = 0.2;                          % enlarges the output window
scale = 1000.0;                     % scaling factor for animation

% Parameters
%-----------
n_nod = length(nod_x);              % total number of nodes 
n_elem = length(elem_nod1);         % total number of elements
n_fix = length(fixed);              % number of fixed nodes
n_forc = length(nod_force);         % number of nodes with ext. forces 
for k=1:1:n_elem                    % calculation of element lengths
  nod1 = elem_nod1(k);
  nod2 = elem_nod2(k);
  l(k) = sqrt((nod_x(nod2)-nod_x(nod1))^2+(nod_y(nod2)-nod_y(nod1))^2);
end

% window size (for output)
x_min = min(nod_x)-fac*(max(nod_x)-min(nod_x));
x_max = max(nod_x)+fac*(max(nod_x)-min(nod_x));
y_min = min(nod_y)-fac*(max(nod_y)-min(nod_y));
y_max = max(nod_y)+fac*(max(nod_y)-min(nod_y));

%-----------------------------------------------------------------
% Calculation part

% set up system stiffness matrix K and system mass matrix M
% ---------------------------------------------------------
K(2*n_nod,2*n_nod) = 0.0;
M(2*n_nod,2*n_nod) = 0.0;

for k=1:1:n_elem                     % for each element

  % calculation of the element stiffness matrix K_elem
  K_elem(4,4) = 0.0;
  K_elem(1,1) = 1.0;
  K_elem(1,3) = -1.0;
  K_elem(3,1) = -1.0;
  K_elem(3,3) = 1.0;
  
  K_elem = K_elem*E(k)*A(k)/l(k);    % K_elem in local co-system
  
  % calculation of the element mass matrix M_elem
  M_elem(1,1) = 2.0;
  M_elem(1,3) = 1.0;
  M_elem(3,1) = 1.0;
  M_elem(3,3) = 2.0;
  M_elem(2,2) = 2.0;
  M_elem(2,4) = 1.0;
  M_elem(4,2) = 1.0;
  M_elem(4,4) = 2.0;
  
  M_elem = M_elem*rho(k)*A(k)*l(k)/6.0;     % M_elem in local co-system

  % calculation of the transformation matrix T
  x1 = nod_x(elem_nod1(k));
  x2 = nod_x(elem_nod2(k));
  y1 = nod_y(elem_nod1(k));
  y2 = nod_y(elem_nod2(k));

  if (x2==x1)                        % angle gamma between local and global co-system 
    gamma = pi/2.0;
  else
    gamma = atan((y2-y1)/(x2-x1)); 
  end
  
  T(1,1) = cos(gamma);
  T(1,2) = sin(gamma);
  T(2,1) = -sin(gamma);
  T(2,2) = cos(gamma);
  T(3,3) = cos(gamma);
  T(3,4) = sin(gamma);
  T(4,3) = -sin(gamma);
  T(4,4) = cos(gamma);

  % transformation of the element stiffness matrix and mass matrix
  K_elem2 = T'*K_elem*T;
  M_elem2 = T'*M_elem*T;

  % assemble element stiffness matrix and mass matrix into system stiffness 
  % matrix K and system stiffnass matrix M
  index(1) = 2*elem_nod1(k)-1;
  index(2) = 2*elem_nod1(k);
  index(3) = 2*elem_nod2(k)-1;
  index(4) = 2*elem_nod2(k);

  for i=1:1:4
    for j=1:1:4
      K(index(i),index(j)) = K(index(i),index(j))+K_elem2(i,j);
      M(index(i),index(j)) = M(index(i),index(j))+M_elem2(i,j);
    end
  end
  
end


% calculation of the force vector f
% ---------------------------------
f(2*n_nod,n_dt+1) = 0.0;
for i=1:1:n_forc
  index = 2*nod_force(i)-1;
  f(index,:) = force_x(i,:);
  f(index+1,:) = force_y(i,:);
end


% reduce system by incorporating fixed nodes (K_red, M_red, f_red)
% ----------------------------------------------------------------
index1 = 1;
for i=1:1:n_nod
  if (ismember(i,fixed)==0)          % node i is not fixed
    index2 = 1;
    for j=1:1:n_nod
      if (ismember(j,fixed)==0)      % node j is not fixed
        K_red(index1:index1+1,index2:index2+1) = K(2*i-1:2*i,2*j-1:2*j);
        M_red(index1:index1+1,index2:index2+1) = M(2*i-1:2*i,2*j-1:2*j);
        index2 = index2+2;           % increase index2 by 2 for entry in next column
      end
    end  
    f_red(index1:index1+1,:) = f(2*i-1:2*i,:);
    index1 = index1+2;               % increase index1 by 2 for entry in next row
  end
end
dof = max(size(K_red));              % number of degrees of freedom in system


% time integration and animation
% ==============================

% (1) Central Difference Method
% -----------------------------
% needed quantities
u_red(dof,1) = 0.0;                  % initial condition (t=0) nodal displacements
u_red_gp(dof) = 0.0;                 % nodal displacements at temporal ghostpoint t=-dt
M_red_inv = inv(M_red);              % inverse of the mass matrix

% defining window for animation
figure(1)
plot(nod_x,nod_y,'bo')
axis([x_min x_max y_min y_max]);
hold on

% undeformed structure
plot(nod_x,nod_y,'bo')
axis([x_min x_max y_min y_max]);
hold on
for k=1:1:n_elem
  x(1) = nod_x(elem_nod1(k));
  x(2) = nod_x(elem_nod2(k));
  y(1) = nod_y(elem_nod1(k));
  y(2) = nod_y(elem_nod2(k));
  plot(x,y,'k')
end

% first time step using the temporal ghost point
time(2) = dt;
u_red(:,2) = dt^2*M_red_inv*(f_red(:,1)-K_red*u_red(:,1))+2.0*u_red(:,1)-u_red_gp(:);

% time stepping with animation
for i=2:1:n_dt
  time(i+1) = i*dt;
  u_red(:,i+1) = dt^2*M_red_inv*(f_red(:,i)-K_red*u_red(:,i))+2.0*u_red(:,i)-u_red(:,i-1);

  % output/animation
  % calculation of complete displacement vector
  index = 1;
  for j=1:1:n_nod
    if (ismember(j,fixed)==1)            % node i is fixed
      u(2*j-1) = 0.0;
      u(2*j) = 0.0;
    else                                 % node i is not fixed
      u(2*j-1) = u_red(index,i+1);
      u(2*j) = u_red(index+1,i+1);
      index = index+2;
    end
  end  

  figure(1)
  % undeformed structure
  plot(nod_x,nod_y,'o','LineWidth',2,'Color',[0.7 0.7 0.7])
  axis([x_min x_max y_min y_max]);
  hold on
  for k=1:1:n_elem
    x(1) = nod_x(elem_nod1(k));
    x(2) = nod_x(elem_nod2(k));
    y(1) = nod_y(elem_nod1(k));
    y(2) = nod_y(elem_nod2(k));
    plot(x,y,'LineWidth',2,'Color',[0.7 0.7 0.7])
  end

  % deformed structure
  for j=1:1:n_nod
    def_x(j) = nod_x(j)+scale*u(2*j-1);
    def_y(j) = nod_y(j)+scale*u(2*j);
  end
  plot(def_x,def_y,'ro','LineWidth',2)
  for k=1:1:n_elem
    x(1) = def_x(elem_nod1(k));
    x(2) = def_x(elem_nod2(k));
    y(1) = def_y(elem_nod1(k));
    y(2) = def_y(elem_nod2(k));
    plot(x,y,'LineWidth',2,'Color',[1 0 0])
  end
  hold off

  title(['Central Difference Method:  Deflection  ---  scale = ' num2str(scale) '  --- time = ' num2str(time(i+1))]) 
end


% (2) Newmark's Method
% --------------------
% needed quantities
u_red(dof,1) = 0.0;                  % initial condition (t=0) nodal displacements
v_red(dof,1) = 0.0;                  % initial condition (t=0) nodal velocities
a_red(dof,1) = 0.0;                  % initial condition (t=0) nodal accelerations
K_eff_inv = inv(M_red/alpha/dt^2+K_red);     % inverse of the effective stiffness matrix

% defining window for animation
figure(2)
plot(nod_x,nod_y,'bo')
axis([x_min x_max y_min y_max]);
hold on

% undeformed structure
plot(nod_x,nod_y,'bo')
axis([x_min x_max y_min y_max]);
hold on
for k=1:1:n_elem
  x(1) = nod_x(elem_nod1(k));
  x(2) = nod_x(elem_nod2(k));
  y(1) = nod_y(elem_nod1(k));
  y(2) = nod_y(elem_nod2(k));
  plot(x,y,'k')
end

% time stepping with animation
for i=1:1:n_dt
  time(i+1) = i*dt;
  u_red(:,i+1) = K_eff_inv*(f_red(:,i+1)+M_red*(u_red(:,i)/alpha/dt^2+v_red(:,i)/alpha/ ...
						dt+(1.0/2.0/alpha-1.0)*a_red(:,i))); 
  a_red(:,i+1) = 1.0/alpha/dt^2*(u_red(:,i+1)-u_red(:,i)) - 1.0/alpha/dt*v_red(:,i) - ...
      (1.0/2.0/alpha-1.0)*a_red(:,i); 
  v_red(:,i+1) = v_red(:,i) + dt*((1.0-delta)*a_red(:,i)+delta*a_red(:,i+1));
  
  % output/animation
  % calculation of complete displacement vector
  index = 1;
  for j=1:1:n_nod
    if (ismember(j,fixed)==1)            % node i is fixed
      u(2*j-1) = 0.0;
      u(2*j) = 0.0;
    else                                 % node i is not fixed
      u(2*j-1) = u_red(index,i+1);
      u(2*j) = u_red(index+1,i+1);
      index = index+2;
    end
  end  

  figure(2)
  % undeformed structure
  plot(nod_x,nod_y,'o','LineWidth',2,'Color',[0.7 0.7 0.7])
  axis([x_min x_max y_min y_max]);
  hold on
  for k=1:1:n_elem
    x(1) = nod_x(elem_nod1(k));
    x(2) = nod_x(elem_nod2(k));
    y(1) = nod_y(elem_nod1(k));
    y(2) = nod_y(elem_nod2(k));
    plot(x,y,'LineWidth',2,'Color',[0.7 0.7 0.7])
  end

  % deformed structure
  for j=1:1:n_nod
    def_x(j) = nod_x(j)+scale*u(2*j-1);
    def_y(j) = nod_y(j)+scale*u(2*j);
  end
  plot(def_x,def_y,'ro','LineWidth',2)
  for k=1:1:n_elem
    x(1) = def_x(elem_nod1(k));
    x(2) = def_x(elem_nod2(k));
    y(1) = def_y(elem_nod1(k));
    y(2) = def_y(elem_nod2(k));
    plot(x,y,'r','LineWidth',2)
  end
  hold off

  title(['Newmarks Method: Deflection  ---  scale = ' num2str(scale) '  ---  time = ' num2str(time(i+1))]) 
end