Heat Diffusion

Given an initial 2D temperature distribution and spatially-varying diffusion coefficients, compute the evolution of the temperature field over time using the following update equation:

$$\begin{align*} A_{i,j}^{k} =\ &A_{i,j}^{k-1} + \texttt{dt} \Big[\texttt{Drow}_{i, j} \cdot \big( A_{i,j+1}^{k-1} - 2A_{i,j}^{k-1} + A_{i, j-1}^{k-1} \big) \\ &\quad +\ \texttt{Dcol}_{i, j} \cdot \big( A_{i+1,j}^{k-1} - 2A_{i,j}^{k-1} + A_{i-1, j}^{k-1} \big)\big] \end{align*}$$

where:

• $A_{i,j}^k$ is the temperature at position $(i,j)$ (row $i$, column $j$) at time step $k$
• $\texttt{Drow}_{i,j}$ is the diffusion coefficient in the row direction (controls diffusion when row index for the tensor changes: $j \pm 1$)
• $\texttt{Dcol}_{i,j}$ is the diffusion coefficient in the column direction (controls diffusion when column index changes: $i \pm 1$)
• $\texttt{dt}$ is the time step size

Input

• A_initial: Initial temperature distribution of size $N \times N$ (given at $k=0$)
• Drow: Diffusion coefficients in row direction of size $N \times N$
• Dcol: Diffusion coefficients across the columns, of size $N \times N$
• N: Size of the grid
• T: Total physical simulation time
• dt: Time step size (0.5 seconds)

Output

• A_out: Temperature field at all time steps of size $\frac{T}{\texttt{dt}} \times N \times N$
  • The initial condition ($k=0$) should be copied to A_out[0, :, :]
  • Fill in time steps $k=1, 2, \ldots, \frac{T}{\texttt{dt}}$

Notes

Use zero-flux (Neumann) boundary conditions at all edges:

• At $i=0$ (first row): treat $A_{-1,j}^k = A_{0,j}^k$
• At $i=N-1$ (last row): treat $A_{N,j}^k = A_{N-1,j}^k$
• At $j=0$ (first column): treat $A_{i,-1}^k = A_{i,0}^k$
• At $j=N-1$ (last column): treat $A_{i,N}^k = A_{i,N-1}^k$