# Implementing Crank-Nicolson Method to Solve 1D Heat Diffusion

Hi all,

for my current project, I am trying to simulate the 1D heat diffusion equation numerically. Previously I have done this using the Crank-Nicholson method in python via (minimal code example):

``````import numpy as np

d_factor = Diffusion * dt / (2*ds * ds)

A_n = np.diagflat([-d_factor for b in range(steps-1)], -1) +\
np.diagflat([1.+d_factor]+[1.+2.*d_factor for b in range(steps-2)]+[1+d_factor]) +\
np.diagflat([-d_factor for b in range(steps-1)], 1)

B_n = np.diagflat([d_factor for b in range(steps-1)], -1) +\
np.diagflat([1.-d_factor]+[1.-2.*d_factor for b in range(steps-2)]+[1.-d_factor]) +\
np.diagflat([d_factor for b in range(steps-1)], 1)

N = N0

for k in range(1,len(t)-1):

N_new = np.linalg.solve(A_n,B_n.dot(N))

###  Boundary Conditions
N_new[0] = N0
N_new[-1] = Nd

N = N_new
``````

where `dt` and `ds` are the step size in time and space, `Diffusion` is a variable and `N0` is the initial state.

Is there a smart way that I can ‘translate’ this code into PyTensor to be able to use it in pm.model ? The issue is that `Diffusion` is a random variable and so `A_n` and `B_n` will need to be calcuated for each sampling step.
Any help would be highly appreciated.

Hi,

You want to use `pytensor.scan`, which is a differentiable loop. Docs here, some discussion and example using a scan together with a CustomDist (which you will want to do) here.

I don’t think this example needs a CustomDist, it’s just a deterministic Scan no?

Depends on the error structure of the model, you’re right. Can just start with a scan.

Hi,

could you provide a short code example of what you mean by that?