%--------------------------------------------------------------------------
% This program is part of the paper "Three-dimensional electrohydrodynamic 
% flows by a central-moments-based lattice Boltzmann method".
% Copyright (C) 2025  Alessandro De Rosis (derosis.alessandro@icloud.com)
% 
% This program is free software; you can redistribute it and/or
% modify it under the terms of the GNU General Public License
% as published by the Free Software Foundation; either version 2
% of the License, or (at your option) any later version.
% 
% This program is distributed in the hope that it will be useful,
% but WITHOUT ANY WARRANTY; without even the implied warranty of
% MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
% GNU General Public License for more details.
% 
% This is the "D3Q19_CentralMoments.m" file.
% Author: Alessandro De Rosis (derosis.alessandro@icloud.com)
% Day: 1st Jun 2025
%--------------------------------------------------------------------------

clear all
clc
syms U V W omega k5_star k6_star k7_star k8_star k9_star Fx Fy Fz R real
syms h0 h1 h2 h3 h4 h5 h6 h7 h8 h9 h10 h11 h12 h13 h14 h15 h16 h17 h18 real

%% Define useful quantities
% Relaxation matrix
K = diag([1 1 1 1 1 omega omega omega omega omega 1 1 1 1 1 1 1 1 1]);
% Lattice directions and weights
w0 = 1/3;
ws = 1/18;
wl = 1/36;

cx = [0 -1 0 0 -1 -1 -1 -1 0 0 1 0 0 1 1 1 1 0 0];
cy = [0 0 -1 0 -1 1 0 0 -1 -1 0 1 0 1 -1 0 0 1 1];
cz = [0 0 0 -1 0 0 -1 1 -1 1 0 0 1 0 0 1 -1 1 -1];
w = [w0 ws ws ws wl wl wl wl wl wl ws ws ws wl wl wl wl wl wl];

% cx = [0 1 -1 0 0 0 0 1 -1 1 -1 1 -1 1 -1 0 0 0 0];
% cy = [0 0 0 1 -1 0 0 1 -1 -1 1 0 0 0 0 1 -1 1 -1];
% cz = [0 0 0 0 0 1 -1 0 0 0 0 1 -1 -1 1 1 -1 -1 1];
% w = [w0 ws ws ws ws ws ws wl wl wl wl wl wl wl wl wl wl wl wl];

cs = 1/sqrt(3);
cs2 = cs^2;
cs3 = cs^3;
cs4 = cs^4;
cs5 = cs^5;
cs6 = cs^6;
cs8 = cs^8;
cs10 = cs^10;
cs12 = cs^12;
T = sym(zeros(19,19));
M = zeros(19,19);
feq = sym(zeros(19,1));
geq = sym(zeros(19,1));
force_guo = sym(zeros(19,1));
Force = sym(zeros(19,1));
U2 = U*U;
V2 = V*V;
W2 = W*W;
u2 = U*U+V*V+W*W;
for i=1:length(cx)
    % build equilibrium populations
    if(cx(i)==0 && cy(i)==0 && cz(i)==0)
        feq(i) = R/3*(1-(U2+V2+W2)+3*(U2*V2+U2*W2+V2*W2));
    elseif(cx(i)~=0 && cy(i)==0 && cz(i)==0)
        feq(i) = R/18*(1+3*cx(i)*U+3*(U2-V2-W2)-9*cx(i)*(U*V2+U*W2)-9*(U2*V2+U2*W2));
    elseif(cx(i)==0 && cy(i)~=0 && cz(i)==0)
        feq(i) = R/18*(1+3*cy(i)*V+3*(-U2+V2-W2)-9*cy(i)*(U2*V+V*W2)-9*(U2*V2+V2*W2));
    elseif(cx(i)==0 && cy(i)==0 && cz(i)~=0)
        feq(i) = R/18*(1+3*cz(i)*W+3*(-U2-V2+W2)-9*cz(i)*(U2*W+V2*W)-9*(U2*W2+V2*W2));
    elseif(cx(i)~=0 && cy(i)~=0 && cz(i)==0)
        feq(i) = R/36*(1+3*(cx(i)*U+cy(i)*V)+3*(U2+V2)+9*cx(i)*cy(i)*U*V+...
                 9*(cy(i)*U2*V+cx(i)*U*V2)+9*U2*V2);
    elseif(cx(i)~=0 && cy(i)==0 && cz(i)~=0)
        feq(i) = R/36*(1+3*(cx(i)*U+cz(i)*W)+3*(U2+W2)+9*cx(i)*cz(i)*U*W+...
                 9*(cz(i)*U2*W+cx(i)*U*W2)+9*U2*W2);
    elseif(cx(i)==0 && cy(i)~=0 && cz(i)~=0)
        feq(i) = R/36*(1+3*(cy(i)*V+cz(i)*W)+3*(V2+W2)+9*cy(i)*cz(i)*V*W+...
                 9*(cz(i)*V2*W+cy(i)*V*W2)+9*V2*W2);
    end
    
    % Set the basis
    CX = cx(i)-U;
    CY = cy(i)-V;
    CZ = cz(i)-W;
    CX2 = CX^2;
    CY2 = CY^2;
    CZ2 = CZ^2;
    T(1,i) = 1;
    T(2,i) = CX;
    T(3,i) = CY;
    T(4,i) = CZ;
    T(5,i) = CX2+CY2+CZ2;
    T(6,i) = CX2-CY2;
    T(7,i) = CY2-CZ2;
    T(8,i) = CX*CY;
    T(9,i) = CX*CZ;
    T(10,i) = CY*CZ;
    T(11,i) = CX2*CY;
    T(12,i) = CX*CY2;
    T(13,i) = CX2*CZ;
    T(14,i) = CX*CZ2;
    T(15,i) = CY2*CZ;
    T(16,i) = CY*CZ2;
    T(17,i) = CX2*CY2;
    T(18,i) = CX2*CZ2;
    T(19,i) = CY2*CZ2;
    
    CX = cx(i);
    CY = cy(i);
    CZ = cz(i);
    CX2 = CX^2;
    CY2 = CY^2;
    CZ2 = CZ^2;
    M(1,i) = 1;
    M(2,i) = CX;
    M(3,i) = CY;
    M(4,i) = CZ;
    M(5,i) = CX2+CY2+CZ2;
    M(6,i) = CX2-CY2;
    M(7,i) = CY2-CZ2;
    M(8,i) = CX*CY;
    M(9,i) = CX*CZ;
    M(10,i) = CY*CZ;
    M(11,i) = CX2*CY;
    M(12,i) = CX*CY2;
    M(13,i) = CX2*CZ;
    M(14,i) = CX*CZ2;
    M(15,i) = CY2*CZ;
    M(16,i) = CY*CZ2;
    M(17,i) = CX2*CY2;
    M(18,i) = CX2*CZ2;
    M(19,i) = CY2*CZ2;
    
    A = U*cx(i)+V*cy(i)+W*cz(i);
    force_guo(i) = w(i)*( Fx*(3*(cx(i)-U)+9*A*cx(i))+...
                          Fy*(3*(cy(i)-V)+9*A*cy(i))+...
						  Fz*(3*(cz(i)-W)+9*A*cz(i)) );
 end
T = simplify(T);
N = simplify(T*M^(-1)); %shift matrix
% N = [1,        0,        0,        0,             0,                 0,                 0,     0,     0,     0,    0,    0,    0,    0,    0,    0, 0, 0, 0;...
%     -U,        1,        0,        0,             0,                 0,                 0,     0,     0,     0,    0,    0,    0,    0,    0,    0, 0, 0, 0;...
%     -V,        0,        1,        0,             0,                 0,                 0,     0,     0,     0,    0,    0,    0,    0,    0,    0, 0, 0, 0;...
%     -W,        0,        0,        1,             0,                 0,                 0,     0,     0,     0,    0,    0,    0,    0,    0,    0, 0, 0, 0;...
%     U^2 + V^2 + W^2,     -2*U,     -2*V,     -2*W,             1,                 0,                 0,     0,     0,     0,    0,    0,    0,    0,    0,    0, 0, 0, 0;...
%     U^2 - V^2,     -2*U,      2*V,        0,             0,                 1,                 0,     0,     0,     0,    0,    0,    0,    0,    0,    0, 0, 0, 0;...
%     V^2 - W^2,        0,     -2*V,      2*W,             0,                 0,                 1,     0,     0,     0,    0,    0,    0,    0,    0,    0, 0, 0, 0;...
%     U*V,       -V,       -U,        0,             0,                 0,                 0,     1,     0,     0,    0,    0,    0,    0,    0,    0, 0, 0, 0;...
%     U*W,       -W,        0,       -U,             0,                 0,                 0,     0,     1,     0,    0,    0,    0,    0,    0,    0, 0, 0, 0;...
%     V*W,        0,       -W,       -V,             0,                 0,                 0,     0,     0,     1,    0,    0,    0,    0,    0,    0, 0, 0, 0;...
%     -U^2*V,    2*U*V,      U^2,        0,          -V/3,          -(2*V)/3,              -V/3,  -2*U,     0,     0,    1,    0,    0,    0,    0,    0, 0, 0, 0;...
%     -U*V^2,      V^2,    2*U*V,        0,          -U/3,               U/3,              -U/3,  -2*V,     0,     0,    0,    1,    0,    0,    0,    0, 0, 0, 0;...
%     -U^2*W,    2*U*W,        0,      U^2,          -W/3,          -(2*W)/3,              -W/3,     0,  -2*U,     0,    0,    0,    1,    0,    0,    0, 0, 0, 0;...
%     -U*W^2,      W^2,        0,    2*U*W,          -U/3,               U/3,           (2*U)/3,     0,  -2*W,     0,    0,    0,    0,    1,    0,    0, 0, 0, 0;...
%     -V^2*W,        0,    2*V*W,      V^2,          -W/3,               W/3,              -W/3,     0,     0,  -2*V,    0,    0,    0,    0,    1,    0, 0, 0, 0;...
%     -V*W^2,        0,      W^2,    2*V*W,          -V/3,               V/3,           (2*V)/3,     0,     0,  -2*W,    0,    0,    0,    0,    0,    1, 0, 0, 0;...
%     U^2*V^2, -2*U*V^2, -2*U^2*V,        0, U^2/3 + V^2/3, (2*V^2)/3 - U^2/3,     U^2/3 + V^2/3, 4*U*V,     0,     0, -2*V, -2*U,    0,    0,    0,    0, 1, 0, 0;...
%     U^2*W^2, -2*U*W^2,        0, -2*U^2*W, U^2/3 + W^2/3, (2*W^2)/3 - U^2/3, W^2/3 - (2*U^2)/3,     0, 4*U*W,     0,    0,    0, -2*W, -2*U,    0,    0, 0, 1, 0;...
%     V^2*W^2,        0, -2*V*W^2, -2*V^2*W, V^2/3 + W^2/3,   - V^2/3 - W^2/3, W^2/3 - (2*V^2)/3,     0,     0, 4*V*W,    0,    0,    0,    0, -2*W, -2*V, 0, 0, 1];
% N = simplify(N);
%N = eye(19,19);
%% Compute central moments
T = simplify(T);
%T = M;
K_eq = simplify(T*feq);
K_force_guo = simplify(T*force_guo);
K_force = subs(K_force_guo, U, 0);
K_force = subs(K_force, V, 0);
K_force = subs(K_force, W, 0);
Id = eye(19,19);
K_pre = sym(zeros(19,1));
syms k5_pre k6_pre k7_pre k8_pre k9_pre k24_pre k25_pre k26_pre real
K_pre(1) = R;
K_pre(6) = k5_pre;
K_pre(7) = k6_pre;
K_pre(8) = k7_pre;
K_pre(9) = k8_pre;
K_pre(10) = k9_pre;
%post-collision central moments
K_star = (Id-K)*K_pre + K*K_eq + (Id-K/2)*K_force
%post collision populations
syms k1 k2 k3 k4 k5 k6 k7 k8 k9 k10 k11 k12 k13 k14 k15 k16 k17 k18 real
K_sym = [R k1 k2 k3 k4 k5 k6 k7 k8 k9 k10 k11 k12 k13 k14 k15 k16 k17 k18];
for i=1:19
    if(K_star(i)~=sym(0))
        K_star(i) = K_sym(i);
    end
end
f_post_collision_onestep = collect(simplify(T\K_star), K_star);

%% Fei's two-steps approach
raw_moments = collect(simplify(N^(-1)*K_star), K_star)
syms r0 r1 r2 r3 r4 r5 r6 r7 r8 r9 r10 r11 r12 r13 r14 r15 r16 r17 r18 real
r = [r0 r1 r2 r3 r4 r5 r6 r7 r8 r9 r10 r11 r12 r13 r14 r15 r16 r17 r18]'; %symbolic raw moments
f_post_collision_twosteps = simplify(M\r)


%% Charge
syms KU KV KW omegaq Q...
    h0 h1 h2 h3 h4 h5 h6 h7 h8 h9 h10 h11 h12 h13 h14 h15 h16 h17 h18 real
heq = sym(zeros(19,1));
h = [h0 h1 h2 h3 h4 h5 h6 h7 h8 h9 h10 h11 h12 h13 h14 h15 h16 h17 h18]';
U2 = KU*KU;
V2 = KV*KV;
W2 = KW*KW;
u2 = U2+V2+W2;
for i=1:length(cx)
    % build equilibrium populations
    if(cx(i)==0 && cy(i)==0 && cz(i)==0)
        heq(i) = Q/3*(1-(U2+V2+W2)+3*(U2*V2+U2*W2+V2*W2));
    elseif(cx(i)~=0 && cy(i)==0 && cz(i)==0)
        heq(i) = Q/18*(1+3*cx(i)*KU+3*(U2-V2-W2)-9*cx(i)*(KU*V2+KU*W2)-9*(U2*V2+U2*W2));
    elseif(cx(i)==0 && cy(i)~=0 && cz(i)==0)
        heq(i) = Q/18*(1+3*cy(i)*KV+3*(-U2+V2-W2)-9*cy(i)*(U2*KV+KV*W2)-9*(U2*V2+V2*W2));
    elseif(cx(i)==0 && cy(i)==0 && cz(i)~=0)
        heq(i) = Q/18*(1+3*cz(i)*KW+3*(-U2-V2+W2)-9*cz(i)*(U2*KW+V2*KW)-9*(U2*W2+V2*W2));
    elseif(cx(i)~=0 && cy(i)~=0 && cz(i)==0)
        heq(i) = Q/36*(1+3*(cx(i)*KU+cy(i)*KV)+3*(U2+V2)+9*cx(i)*cy(i)*KU*KV+...
                 9*(cy(i)*U2*KV+cx(i)*KU*V2)+9*U2*V2);
    elseif(cx(i)~=0 && cy(i)==0 && cz(i)~=0)
        heq(i) = Q/36*(1+3*(cx(i)*KU+cz(i)*KW)+3*(U2+W2)+9*cx(i)*cz(i)*KU*KW+...
                 9*(cz(i)*U2*KW+cx(i)*KU*W2)+9*U2*W2);
    elseif(cx(i)==0 && cy(i)~=0 && cz(i)~=0)
        heq(i) = Q/36*(1+3*(cy(i)*KV+cz(i)*KW)+3*(V2+W2)+9*cy(i)*cz(i)*KV*KW+...
                 9*(cz(i)*V2*KW+cy(i)*KV*W2)+9*V2*W2);
    end
    
    % Set the basis
    CX = cx(i)-KU;
    CY = cy(i)-KV;
    CZ = cz(i)-KW;
    CX2 = CX^2;
    CY2 = CY^2;
    CZ2 = CZ^2;
    T(1,i) = 1;
    T(2,i) = CX;
    T(3,i) = CY;
    T(4,i) = CZ;
    T(5,i) = CX2+CY2+CZ2;
    T(6,i) = CX2-CY2;
    T(7,i) = CY2-CZ2;
    T(8,i) = CX*CY;
    T(9,i) = CX*CZ;
    T(10,i) = CY*CZ;
    T(11,i) = CX2*CY;
    T(12,i) = CX*CY2;
    T(13,i) = CX2*CZ;
    T(14,i) = CX*CZ2;
    T(15,i) = CY2*CZ;
    T(16,i) = CY*CZ2;
    T(17,i) = CX2*CY2;
    T(18,i) = CX2*CZ2;
    T(19,i) = CY2*CZ2;
    
    CX = cx(i);
    CY = cy(i);
    CZ = cz(i);
    CX2 = CX^2;
    CY2 = CY^2;
    CZ2 = CZ^2;
    M(1,i) = 1;
    M(2,i) = CX;
    M(3,i) = CY;
    M(4,i) = CZ;
    M(5,i) = CX2+CY2+CZ2;
    M(6,i) = CX2-CY2;
    M(7,i) = CY2-CZ2;
    M(8,i) = CX*CY;
    M(9,i) = CX*CZ;
    M(10,i) = CY*CZ;
    M(11,i) = CX2*CY;
    M(12,i) = CX*CY2;
    M(13,i) = CX2*CZ;
    M(14,i) = CX*CZ2;
    M(15,i) = CY2*CZ;
    M(16,i) = CY*CZ2;
    M(17,i) = CX2*CY2;
    M(18,i) = CX2*CZ2;
    M(19,i) = CY2*CZ2;
 end
T = simplify(T);
N = simplify(T*M^(-1)); %shift matrix
% N = [1,        0,        0,        0,             0,                 0,                 0,     0,     0,     0,    0,    0,    0,    0,    0,    0, 0, 0, 0;...
%     -U,        1,        0,        0,             0,                 0,                 0,     0,     0,     0,    0,    0,    0,    0,    0,    0, 0, 0, 0;...
%     -V,        0,        1,        0,             0,                 0,                 0,     0,     0,     0,    0,    0,    0,    0,    0,    0, 0, 0, 0;...
%     -W,        0,        0,        1,             0,                 0,                 0,     0,     0,     0,    0,    0,    0,    0,    0,    0, 0, 0, 0;...
%     U^2 + V^2 + W^2,     -2*U,     -2*V,     -2*W,             1,                 0,                 0,     0,     0,     0,    0,    0,    0,    0,    0,    0, 0, 0, 0;...
%     U^2 - V^2,     -2*U,      2*V,        0,             0,                 1,                 0,     0,     0,     0,    0,    0,    0,    0,    0,    0, 0, 0, 0;...
%     V^2 - W^2,        0,     -2*V,      2*W,             0,                 0,                 1,     0,     0,     0,    0,    0,    0,    0,    0,    0, 0, 0, 0;...
%     U*V,       -V,       -U,        0,             0,                 0,                 0,     1,     0,     0,    0,    0,    0,    0,    0,    0, 0, 0, 0;...
%     U*W,       -W,        0,       -U,             0,                 0,                 0,     0,     1,     0,    0,    0,    0,    0,    0,    0, 0, 0, 0;...
%     V*W,        0,       -W,       -V,             0,                 0,                 0,     0,     0,     1,    0,    0,    0,    0,    0,    0, 0, 0, 0;...
%     -U^2*V,    2*U*V,      U^2,        0,          -V/3,          -(2*V)/3,              -V/3,  -2*U,     0,     0,    1,    0,    0,    0,    0,    0, 0, 0, 0;...
%     -U*V^2,      V^2,    2*U*V,        0,          -U/3,               U/3,              -U/3,  -2*V,     0,     0,    0,    1,    0,    0,    0,    0, 0, 0, 0;...
%     -U^2*W,    2*U*W,        0,      U^2,          -W/3,          -(2*W)/3,              -W/3,     0,  -2*U,     0,    0,    0,    1,    0,    0,    0, 0, 0, 0;...
%     -U*W^2,      W^2,        0,    2*U*W,          -U/3,               U/3,           (2*U)/3,     0,  -2*W,     0,    0,    0,    0,    1,    0,    0, 0, 0, 0;...
%     -V^2*W,        0,    2*V*W,      V^2,          -W/3,               W/3,              -W/3,     0,     0,  -2*V,    0,    0,    0,    0,    1,    0, 0, 0, 0;...
%     -V*W^2,        0,      W^2,    2*V*W,          -V/3,               V/3,           (2*V)/3,     0,     0,  -2*W,    0,    0,    0,    0,    0,    1, 0, 0, 0;...
%     U^2*V^2, -2*U*V^2, -2*U^2*V,        0, U^2/3 + V^2/3, (2*V^2)/3 - U^2/3,     U^2/3 + V^2/3, 4*U*V,     0,     0, -2*V, -2*U,    0,    0,    0,    0, 1, 0, 0;...
%     U^2*W^2, -2*U*W^2,        0, -2*U^2*W, U^2/3 + W^2/3, (2*W^2)/3 - U^2/3, W^2/3 - (2*U^2)/3,     0, 4*U*W,     0,    0,    0, -2*W, -2*U,    0,    0, 0, 1, 0;...
%     V^2*W^2,        0, -2*V*W^2, -2*V^2*W, V^2/3 + W^2/3,   - V^2/3 - W^2/3, W^2/3 - (2*V^2)/3,     0,     0, 4*V*W,    0,    0,    0,    0, -2*W, -2*V, 0, 0, 1];
% N = simplify(N);
%N = eye(19,19);
%% Compute central moments
h = [h0 h1 h2 h3 h4 h5 h6 h7 h8 h9 h10 h11 h12 h13 h14 h15 h16 h17 h18]';
Kq = diag([1 omegaq omegaq omegaq 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1]);
T = simplify(T);
Kq_eq = simplify(T*heq);
Kq_pre = sym(zeros(19,1));
syms k1q_pre k2q_pre k3q_pre real
Kq_pre(1) = Q;
Kq_pre(2) = k1q_pre;
Kq_pre(3) = k2q_pre;
Kq_pre(4) = k3q_pre;
%post-collision central moments
Kq_star = (Id-Kq)*Kq_pre + K*Kq_eq
%post collision populations
syms kq1 kq2 kq3 kq4 kq5 kq6 kq7 kq8 kq9 kq10 kq11 kq12 kq13 kq14 kq15 kq16 kq17 kq18 real
Kq_sym = [Q kq1 kq2 kq3 kq4 kq5 kq6 kq7 kq8 kq9 kq10 kq11 kq12 kq13 kq14 kq15 kq16 kq17 kq18];
for i=1:19
    if(Kq_star(i)~=sym(0))
        Kq_star(i) = Kq_sym(i);
    end
end
h_post_collision_onestep = collect(simplify(T\Kq_star), Kq_star);

%% Two-steps approach
raw_moments = collect(simplify(N^(-1)*Kq_star), Kq_star)
syms rq0 rq1 rq2 rq3 rq4 rq5 rq6 rq7 rq8 rq9 rq10 rq11 rq12 rq13 rq14 rq15 rq16 rq17 rq18 real
rq = [rq0 rq1 rq2 rq3 rq4 rq5 rq6 rq7 rq8 rq9 rq10 rq11 rq12 rq13 rq14 rq15 rq16 rq17 rq18]'; %symbolic raw moments
h_post_collision_twosteps = simplify(M\rq)