% D3Q19 lattice-Boltzmann method for solving Stokes equation,
% BGK and TRT collision operators,
% developed by Siarhei Khirevich, 1.Jun.2020, KAUST

clearvars; close all

%% ------------- INPUT PARAMETERS -------------
TRT = true; % use 'true' for TRT and 'false' for BGK collision
viscosity = 0.5; % can be modified for TRT; for BGK the user input value is ignored
magic = 3/16; % 3/16 (3/8) gives exact velocity profile in horizontal (diagonal)
              % channel; 1/8 provides exact average flow rate in horizontal channel

body_force = [1e-6 0.0 0.0]; % body force vector
total_iterations = 5000; % total LBM iterations
eps = 1e-8; % simulation termination accuracy in X-velocity
              
horiz_diag_channel_check = 0; % 0 for custom geometry to be read from file
                              % 1 for horizontal open channel
                              % 2 for diagonal channel

switch horiz_diag_channel_check
  case 0
    geo_dim = [100 100 100]; % size of computational domain
    geo = read_dis_file('input_100.dis', geo_dim); % specify path to your binary geometry file
  case 1
    geo_dim = [1 10 1]; % size of computational domain for horizontal channel,
                        % 2nd dimension is channel aperture
     % array storing geometry, must have only FLUID and SOLID values
  case 2
    geo_dim = [10 10 1]; % size of computational domain for horizontal channel, 1st and 2nd dims must be equal
    
  otherwise
    error(' "horiz_diag_channel_check" must be 0, 1, or 2');
end
%% ------------- END OF INPUT PARAMETERS -------------

FLUID = 0; % two voxel types are supported
SOLID = 1; 

% Initialization of computational domain; it is periodic along 3 directions X, Y, Z
% geo(:, :, :) = FLUID;
% geo( , , ) = SOLID; % specify your geometry here introducing SOLID voxels

if TRT % for TRT collision viscosity (and magic) can be considered as input parameters
  s_plus = 1.0 / (3.0 * viscosity + 0.5);
  s_minus = 1.0 / (magic / (1.0 / s_plus - 0.5) + 0.5);
else % for BGK if magic is fixed, viscosity is fixed too
  tau = sqrt(magic) + 0.5; % relaxation time for BGK (single-relaxation time) collision
  viscosity = (tau - 0.5) / 3.0;
end

%% Poiseuille flow, first part                              
switch horiz_diag_channel_check
  case 1 % horizontal channel, 2D test case
    % let us introduce a single line of solid voxels along 2-nd dimension, due to periodic 
    % boundary conditions this is enough to mimic flow between two parallel plates
    geo = zeros(geo_dim, 'int8');
    geo(:, :, :) = FLUID;
    geo(:, 1, :) = SOLID;
    channel_width = geo_dim(2) - 1;

  case 2 % diagonal channel, 2D test case, deals with (geo_dim(1) x geo_dim(1)) 
         % square matrix only. geo_dim(1) must exceed 2. For correct location 
         % of zero-velocity boundary see Fig. 1 in doi:10.1051/jp2:1994123 or
         % Fig. 1 in doi:10.1016/j.jcp.2014.10.038
    if (geo_dim(1) ~= geo_dim(2) || geo_dim(1) < 3)
      error('For diagonal channel test, X and Y domain dimensions must be equal, and exceed 2');
    end
    geo = zeros(geo_dim, 'int8');
    %body_force(2) = body_force(1); % let us apply forcing along the channel center line
    body_force(3) = 0;
    geo(:, :, :) = FLUID;    
    % this initializes diagonal (45 degrees) channel geometry in square matrix
    % First start with corner:
    geo([1 2], end, :) = SOLID; geo(1, end-1, :) = SOLID; geo(2, 1, :) = SOLID;
    % and the rest of geometry for arbitrary geo_dim(1), keeping wall thickness of 2 voxels
    for i = 3:geo_dim(1)
      geo(i, [i-2 i-1], :) = SOLID;
    end

    length_unit = 2^(-0.5); % let us introduce basic unit length to operate with it below; 
                            % It equals half diagonal of a pixel (not
                            % voxel), case is two-dimensional, sqrt(2) is in play
    channel_width = length_unit * (geo_dim(1) - 2);
    channel_width_in_units = (geo_dim(1) - 2);


end

%% Initialization of LBM constants
n_links = 19; % const, for D3Q19 model
CS2 = 1.0/3.0; % const, speed of sound squared

% lattice directions, D3Q19, no [1 1 1]-type links
%       1 2  3 4  5 6  7   8 9 10 11 12 13 14 15 16 17 18 19
dirx = [0 1 -1 0  0 0  0   1 -1 1 -1 0  0   1 -1  1 -1  0  0];
diry = [0 0  0 1 -1 0  0   1 -1 0  0 1 -1  -1  1  0  0  1 -1];
dirz = [0 0  0 0  0 1 -1   0  0 1 -1 1 -1   0  0 -1  1 -1  1];
% please beware that directions are ordered in pairs: vector (x,y,z) is followed
% its opposite counterpart, (-x,-y,-z); this is necessary for correct bounce-back run
% using its current implementation

% indexes with opposite lattice directions, e. g. 2nd vector is antiparallel to 3rd one 
opposite_dir = int8([1, 3, 2, 5, 4, 7, 6, 9, 8, 11, 10, 13, 12, 15, 14, 17, 16, 19, 18]);
% indexes of nine vectors antiparallel to another nine vectors
antipar_dir = int8([2 4 6 8 10 12 14 16 18]);

% weights, such that w = CS2 * B in standard notation
B = [1.0, 6.0,  6.0,  6.0,  6.0,  6.0,  6.0, ...
         12.0, 12.0, 12.0, 12.0, 12.0, 12.0, ...
         12.0, 12.0, 12.0, 12.0, 12.0, 12.0];
B = 1./B;

%% Initialization of bounce-back arrays (for boundary condition)
direction_to_solid = zeros([n_links, geo_dim], 'int8');

% 'direction_to_solid' is needed for bounce back boundary condition later;
% check all fluid nodes, find the ones touching solid nodes via given lattice direction 
% and mark the the corresponting lattice direction + the opposite one; after first tries I was 
% not able to implement this in a fully matrix notation, hence please enjoy tons of for's)
% Please note that this implementation requires (below) special ordering of
% lattice direction vectors
for i = 1:geo_dim(1)
  for j = 1:geo_dim(2)
    for k = 1:geo_dim(3)
      if geo(i, j, k) == FLUID
        for l = 2:n_links
          if (geo( mod(i - 1 + dirx(l) + geo_dim(1), geo_dim(1)) + 1, ...
                   mod(j - 1 + diry(l) + geo_dim(2), geo_dim(2)) + 1, ...
                   mod(k - 1 + dirz(l) + geo_dim(3), geo_dim(3)) + 1) == SOLID)
             if direction_to_solid(l, i, j, k) == 0 % (opposite) direction was not ititialized previously 
               direction_to_solid(l, i, j, k) = 1;
               direction_to_solid(opposite_dir(l), i, j, k) = -1;
             else % direction to solid voxel was previously initialized; 
                  % it happens when FLUID voxel is touched by two SOLID ones 
                  % _along_ given direction, such as in slit channel of one pixel width.
                  % Here we will need later to reflect back both directions
               direction_to_solid(l, i, j, k) = 2;
               direction_to_solid(opposite_dir(l), i, j, k) = 3;
             end
          end
        end
      end
    end
  end
end


%% Allocation of LBM lattices
% Originally, this code was implemented in lat(l, x, y, z) layout, where 
% lat(l, ...) and lat(l+1, ...) are located next to each other in memory.
% However, performance check showed ~20-40% advantage of lat(x, y, z, l) layout
lat = zeros([geo_dim, n_links]); lat_after_collision = lat; lat_force = lat;
     
%initialize lattice, assuming density is 1.0
for l = 1:n_links
  lat(:, :, :, l) = CS2*B(l) - 0.5*B(l) * (body_force(1)*dirx(l) + body_force(2)*diry(l) + body_force(3)*dirz(l));
  lat_force(:, :, :, l) = B(l)*(body_force(1)*dirx(l) + body_force(2)*diry(l) + body_force(3)*dirz(l)); 
                              % lat_force can be multiplied by *(1 - 1/(2*tau)), then
                              % 1) vx = sum(..) + 0.5*body_force(1) in main loop and
                              % 2) vx at the end does not need "+ 0.5*body_force(1)"
end

if TRT
  lat_plus = lat; lat_minus = lat; lat_equi_plus = lat; lat_equi_minus = lat;
else
  lat_equi = lat;  
end

T = 200; % number of iterations to accumulate convergence values
t = 1; % counter for convergence check
exit_flag = false; % flag for simulation termination


output_interval_sec = 2.0; % in seconds, output simulation info interval
last_output = 0.0;
tic
%% Main LBM cycle
for its = 1:total_iterations
  % macro quantities, density
  den = sum(lat, 4);
  % and momentum, J in standard notation
  Jx = sum( lat(:, :, :, dirx > 0) - lat(:, :, :, dirx < 0), 4);
  Jy = sum( lat(:, :, :, diry > 0) - lat(:, :, :, diry < 0), 4);
  Jz = sum( lat(:, :, :, dirz > 0) - lat(:, :, :, dirz < 0), 4);

  % calculate mean X-momentum (velocity) value at FLUID nodes
  jx = reshape(Jx, geo_dim) + 0.5*body_force(1);
  jy = reshape(Jy, geo_dim) + 0.5*body_force(2);
  jz = reshape(Jz, geo_dim) + 0.5*body_force(3);
  jx = jx(geo == FLUID); jy = jy(geo == FLUID); jz = jz(geo == FLUID);
  j_current = norm( [mean(jx(:)) mean(jy(:)) mean(jz(:))] );
  
  if its > 1
    this_output = toc;
    relative_momentum_change = abs(1.0 -  j_current / jx_prev);
    if this_output - last_output > output_interval_sec
      disp([num2str(its) ': |Vi+1/Vi - 1|: ' ...
             num2str(relative_momentum_change, '%.1e') ', ' num2str(eps, '%.1e') ' required for ' num2str(T) ' its.'])
      last_output = this_output;
    end
    if relative_momentum_change < eps
      if t == T % if exit criterion is satisfied for T consequent iterations
        exit_flag = true;
      else
        t = t+1;
      end
    else % if exit criterion is not satisfied for at least one out of T iterations, reset t
      t = 1;
    end

    jx_prev = j_current;
  else
    jx_prev = j_current;
  end
  
  if its == total_iterations || exit_flag
    if its == total_iterations
      disp(['  ---> Total iterations limit exceeded, accuracy of ' num2str(eps, '%.2e') '  not reached']);
    else
      disp([' ---> Accuracy criterion of ' num2str(eps, '%.2e') ' is reached']);
    end
    break
  end

  % Collision
  if TRT % TRT collision
    lat_plus(:, :, :, 1) = lat(:, :, :, 1);
    lat_plus(:, :, :, antipar_dir) = 0.5*(lat(:, :, :, antipar_dir) + lat(:, :, :, opposite_dir(antipar_dir)));
    lat_plus(:, :, :, opposite_dir(antipar_dir)) = lat_plus(:, :,  :, antipar_dir);

    lat_minus(:, :, :, 1) = 0.0;
    lat_minus(:, :, :, antipar_dir) = 0.5*(lat(:, :, :, antipar_dir) - lat(:, :, :, opposite_dir(antipar_dir)));
    lat_minus(:, :, :, opposite_dir(antipar_dir)) = -lat_minus(:, :, :, antipar_dir);

    for l = 2:n_links
      lat_equi_plus(:, :, :, l) = CS2 * B(l)*den;
    end
    lat_equi_plus(:, :, :, 1) = den - 2.0*sum(lat_equi_plus(:, :, :, antipar_dir), 4); % e0

    % calculate linear equilibrium, Stokes only
    lat_equi_minus(:, :, :, 1) = 0.0;
    for l = 2:n_links
      lat_equi_minus(:, :, :, l) = B(l) * (Jx*dirx(l) + Jy*diry(l) + Jz*dirz(l));
    end

    lat_after_collision = lat - s_plus*(lat_plus - lat_equi_plus) - s_minus*(lat_minus - lat_equi_minus) + lat_force;
  else % BGK collision
      % calculate linear equilibrium (Stokes only, no Navier-- term)
      for l = 1:n_links
        lat_equi(:, :, :, l) = CS2*B(l) * (den + 3.0*(Jx*dirx(l) + Jy*diry(l) + Jz*dirz(l)));
      end

      % collision + forcing, BGK
      lat_after_collision = (1.0 - 1.0/tau)*lat + lat_equi/tau + lat_force;
  end
  
  % propagation, periodic boundary conditions
  for l = 1:n_links
    lat(:, :, :, l) = circshift(lat_after_collision(:, :, :, l), [dirx(l) diry(l) dirz(l) 0]); 
    % "0" in [... dirz(l) 0] is here because "lat_after_collision" is 4D array
  end

  % Bounce-back, lattice directions facing SOLID voxels are reflected back;
  % values are taken from "lat_after_collision" which is not affected by propagation above.
  % 'permute' here is needed for current bounce-back implementation, which
  % was designed for lat(l, x, y, z) layout and will not work with lat(x, y, z, l)
  lat = permute(lat, [4 1 2 3]); lat_after_collision = permute(lat_after_collision, [4 1 2 3]); 
  lat(direction_to_solid == -1) = lat_after_collision(direction_to_solid == 1);

  % case of two anti-parallel directions touching SOLID voxels simultaneously
  lat(direction_to_solid == 3) = lat_after_collision(direction_to_solid == 2);
  lat(direction_to_solid == 2) = lat_after_collision(direction_to_solid == 3);
  lat = permute(lat, [2 3 4 1]); lat_after_collision = permute(lat_after_collision, [2 3 4 1]); 
end
%% End of main LBM cycle
toc
elapsedTime = toc;
disp([num2str(sum(geo(:) == FLUID)*its/elapsedTime/10^6) ' MFLUPS, million(s) of fluid lattice updates/sec']);

% output velocity field with 0.5*(body_force) shift. reshape is needed to switch to 3D from 4D array
vx = reshape(Jx, geo_dim) + 0.5*body_force(1); 
vy = reshape(Jy, geo_dim) + 0.5*body_force(2);
vz = reshape(Jz, geo_dim) + 0.5*body_force(3);
vx(geo == SOLID) = 0.0; vy(geo == SOLID) = 0.0; vz(geo == SOLID) = 0.0; % explicitly set SOLID voxels to zero
disp(['Average velocity X, Y, Z: ', num2str(mean(vx(geo == FLUID)), '%.3e'), ' ', num2str(mean(vy(geo == FLUID)), '%.3e'), ' ', num2str(mean(vz(geo == FLUID)), '%.3e')])

[mx, my, mz] = meshgrid(1:geo_dim(2), 1:geo_dim(1), 1:geo_dim(3));
figure; qh = quiver3(mx - 0.5, my - 0.5, mz - 0.5, vy, vx, vz, 0.8); daspect([1 1 1])
color_arrows_quiver(qh);
title('Flow field vectors at centers of FLUID voxels');
xlabel('X'); ylabel('Y'); zlabel('Z');

delta_pressure = (sum(lat, 4) - 1.0) * CS2; delta_pressure(geo == SOLID) = 0.0;


%% Poiseuille flow, second part
if horiz_diag_channel_check == 1 || horiz_diag_channel_check == 2
  n_analyt_profile = 200;  % 200 points for analytical profile
  analyt_profile_step = channel_width / n_analyt_profile;
  vel_prof_plot = zeros(n_analyt_profile, 1);
end
switch horiz_diag_channel_check
  case 1 % horizontal channel

    vel_prof = zeros(channel_width, 1);

    % velocity profile and average flow rate from: 
    % https://en.wikipedia.org/wiki/Hagen%E2%80%93Poiseuille_equation#Plane_Poiseuille_flow
    for c = 1:channel_width
      vel_prof(c) = (c-0.5)*abs(channel_width - c + 0.5)/viscosity*0.5*body_force(1);
    end
    
    vx_len = size(vx, 2) - 1;
    figure; plot((1:vx_len)-0.5, vx(1, 2:end, 1), 'o'); hold on

    for c = 1:n_analyt_profile
      vel_prof_plot(c) = analyt_profile_step*c*abs(channel_width - c*analyt_profile_step)/viscosity*0.5*body_force(1);
    end

    vel_avg_analy_plot = channel_width^2*body_force(1) / (12.0 * viscosity);
    vel_avg_simu_plot = mean(vx(geo == FLUID));
    max_vel_all = max(max(vel_prof_plot(c)), max(vx(1, 2:end, 1)));
    
    % let us add relative velocity error (to analytics) for each voxel
    rel_error_x = (1:vx_len)-0.5;
    rel_error_y = 1.0 - vx(1, 2:end, 1)./vel_prof';
    
  case 2 % diagonal channel
    % analytical profile, similar to horiz_diag_channel_check = 1
    vel_prof = zeros(channel_width_in_units, 1);
    for c = 1:channel_width_in_units
      vel_prof(c) = length_unit^2*(c - 0.5)*abs(channel_width_in_units - c + 0.5)/viscosity*0.5*(body_force(1)*cos(pi/4) + body_force(2)*sin(pi/4));
    end
    
    vxy = (vx.^2 + vy.^2).^0.5;
    % there are 'size(vx, 1)-2' diagonals, for each one is enough one point
    vx_len = size(vx, 1) - 2;
    figure; plot( ((1:vx_len)-0.5)*length_unit, vxy(1, 1:vx_len, 1), 'o'); hold on
    
    for c = 1:n_analyt_profile
      vel_prof_plot(c) = analyt_profile_step*c*abs(channel_width_in_units*length_unit - c*analyt_profile_step)/viscosity*0.5*(body_force(1)*cos(pi/4) + body_force(2)*sin(pi/4));
    end
    channel_width = channel_width_in_units*length_unit;
    
    % relative velocity error
    rel_error_x = ((1:vx_len)-0.5)*length_unit;
    rel_error_y = 1.0 - vxy(1, 1:vx_len, 1)./vel_prof';
    
    % average and max(simu, analy) velocities
    vel_avg_analy_plot = (channel_width)^2*(body_force(1)*cos(pi/4) + body_force(2)*sin(pi/4)) / (12.0 * viscosity);
    vel_avg_simu_plot = mean(vxy(geo == FLUID));
    max_vel_all = max(max(vel_prof_plot(c)), max(vxy(1, 1:vx_len, 1)));
end

if horiz_diag_channel_check == 1 || horiz_diag_channel_check == 2
    xlabel('distance across channel, from left wall, in lattice units');
    if TRT; col_operator = 'TRT'; else col_operator = 'BGK'; end
    if horiz_diag_channel_check == 1 
      channel_type = 'Horizontal';
      flow_direction = 'X';
    else
      channel_type = 'Diagonal';
      flow_direction = 'X-Y';
    end
    title({[channel_type ' channel, \Lambda = ', num2str(magic), ', \nu = ', num2str(viscosity), ', ', ...
            col_operator, ' collision,'], ...
           ['body force: ', num2str(body_force(1), '%.2e'), ... 
            ', channel width: ', num2str(channel_width, '%.3f'), ' lattice unit(s),'], ...
           ['relative average velocity error: ', num2str(1.0 - vel_avg_simu_plot/vel_avg_analy_plot, '%.2e')]})
    
    if(verLessThan('matlab','9.0')) % if 'yyaxis' is not available, it was introduced in 2016a release
      [ax,h1,h2] = plotyy(analyt_profile_step*[1:n_analyt_profile], vel_prof_plot, rel_error_x, rel_error_y);
      set(h1, 'LineStyle', '-'); set(h2, 'LineStyle', 'none', 'Marker', 'x')
      set(get(ax(1),'Ylabel'), 'String', ['velocity along ', flow_direction,' direction, in lattice units']) 
      set(get(ax(2),'Ylabel'), 'String', 'relative error, 1 - v_{simu}/v_{analy}') 
      set(ax(2),'YLim', 2.0*[-max(abs(rel_error_y)) max(abs(rel_error_y))]) 
      set(ax(1),'YTick', 0 : max_vel_all/4.0 : max_vel_all); set(ax(1),'Box','off') 
      set(ax(2),'YTick', 2.0*(-max(abs(rel_error_y)) : max(abs(rel_error_y))/2.0 : max(abs(rel_error_y))) ) 
    else
      plot(analyt_profile_step*[1:n_analyt_profile], vel_prof_plot, 'r-');
      ylabel(['velocity along ', flow_direction,' direction, in lattice units']);
      yyaxis right
      plot(rel_error_x, rel_error_y, 'xg')
      ylim(2.0*[-max(abs(rel_error_y)) max(abs(rel_error_y))])
      ylabel('relative error, 1 - v_{simu}/v_{analy}');
      ax = gca; ax.YAxis(2).Color = 'g';
    end
    legend('simu. velocity', 'analy. velocity', 'rel. error')
end