% =========================================================
% 热传导方程三种格式对比
% Explicit / Implicit / Crank-Nicolson
% =========================================================
clear; clc; close all;

% ---------------- 参数 ----------------
L = 1.0;
Nx = 100;
dx = L/(Nx-1);

alpha = 1;

% 计算格式
method = 'explicit';   % 'explicit' / 'implicit' / 'CN'

% 显式稳定性控制
sigma = 0.51;  % >0.5 不稳定（显式）
dt = sigma * dx^2 / alpha;

T_end = 0.5;
Nt = round(T_end/dt);

x = linspace(0,L,Nx);
t_all = (0:Nt)*dt;

% ---------------- 初始条件 ----------------
U = zeros(Nt+1, Nx);

u = sin(pi*x);   

U(1,:) = u;

% =========================================================
% 确定系数
% =========================================================

if strcmp(method,'implicit') || strcmp(method,'explicit')
    r = alpha*dt/dx^2;
elseif strcmp(method,'CN')
    r = alpha*dt/(2*dx^2);
end


% =========================================================
% 构造矩阵（隐式 / CN）
% =========================================================


if strcmp(method,'implicit') || strcmp(method,'CN')
    
    % 构造三对角矩阵【同学可借助AI理解】
    main = (1+2*r)*ones(Nx,1);
    off  = -r*ones(Nx,1);    
    A = spdiags([off main off],[-1 0 1],Nx,Nx);
    
    % 边界
    A(1,:) = 0; A(1,1) = 1;
    A(end,:) = 0; A(end,end) = 1;
end

if strcmp(method,'CN')
    
    mainB = (1-2*r)*ones(Nx,1);
    offB  = r*ones(Nx,1);
    
    B = spdiags([offB mainB offB],[-1 0 1],Nx,Nx);
    
    B(1,:) = 0; B(1,1)=1;
    B(end,:) = 0; B(end,end)=1;
end

%% ---------------- 图像 ----------------
figure('Color','w','Position',[100 100 1300 600]);

subplot(1,2,1)
h_num = plot(x,u,'b-','LineWidth',2); hold on;
h_exact = plot(x,u,'r--','LineWidth',1.5);
legend('数值解','解析解');
xlabel('x'); ylabel('T');
grid on;
axis([0 L -1 1]);

title(method);

subplot(1,2,2)
h_img = imagesc(x, t_all, U);
set(gca,'YDir','normal');
xlabel('x'); ylabel('t');
title('x-t 演化');
colorbar;

% =========================================================
% 时间推进
% =========================================================
for n = 1:Nt
    
    t = n*dt;
    
    switch method
        case 'explicit'
            
            u_new = u;
            u_new(2:end-1) = u(2:end-1) + r * ...
                (u(3:end) - 2*u(2:end-1) + u(1:end-2));
            
        case 'implicit'
            
            b = u';
            u_new = (A \ b)';
            

        case 'CN'
            
            b = B * u';
            u_new = (A \ b)';
    end
    
    % 边界条件
    u_new(1) = 0;
    u_new(end) = 0;
    
    u = u_new;
    
    % 存储
    U(n+1,:) = u;
    
    % 解析解
    u_exact = exp(-alpha*pi^2*t)*sin(pi*x);
    
    % 刷新
    if mod(n,5)==0
        
        subplot(1,2,1)
        set(h_num,'YData',u);
        set(h_exact,'YData',u_exact);
        title([method,' , t=',num2str(t,'%.3f')])
        
        subplot(1,2,2)
        set(h_img,'CData',U);
        
        drawnow limitrate
        % pause(0.1)
    end
end