ルンゲ・クッタ法による拡散方程式におけるループがうまくいかない

Question

前提・実現したいこと

現在，Pythonで拡散方程式の数値解析をルンゲ・クッタ法で行うコードを作っています。
解こうとしている計算式がこちらになります。

コードの参考にしているサイトがこちらになります 。

https://math-fun.net/20191212/4034/

この式を解く過程で時間発展f(x,t+Δt)の行がうまく行きません。
この関数ｆのコード内での定義方法だとｊを関数に組み込めず関数として機能しないためにｊのループが上手くいかないのだと思いますが，どのように定義すればいいかがわからない状態です。何卒アドバイスをいただければ幸いです。

発生している問題・エラーメッセージ

いくつかためした中ででたエラーメッセージとしては以下になります
can't assign to function call
'function' object does not support item assignment 
object does not support item assignment.
など

該当のソースコード

途中までですが製作中のコードです

import math
import numpy as np
import matplotlib.pyplot as plt
 

# t = Tmax まで計算
Tmax = 2.0
# 時間ステップ
dt = 0.5
# ステップ数
nt = int(Tmax / dt)
# 領域幅
L = 1.5
# x方向のメッシュ幅
dx = 0.5
# x方向の分点数
nx = int(L / dx)
# x座標
x = np.arange(0, L, dx)
# 拡散係数
nu = 1
# グラフの格納場所
ims = []


def f(i,j):
    return np.cos(i*dx)

def g(i,j):
    return nu*(f(i+1,j)-2*f(i,j)+f(i-1,j))



for j in range(nt):
    for i in range(nx):
        k_1 = g(i,j)
        k_2 = g(i,j) + k_1/2 
        k_3 = g(i,j) + k_2/2 
        k_4 = g(i,j) + k_3

        f(i,j+1) = f(i,j) + 1/6*(k_1+2*k_2+2*k_3+k_4)*dt

Answer 1

Odeなどのモジュール利用により