heat_eqn_smooth.py

import numpy as np
import pandas as pd
import yfinance as yf

import matplotlib.pyplot as plt

def heat_smooth(
        ts: np.array,
        h: float,
        k: float,
        t_end: float=3.0,
    ) -> np.array:
    
    U = ts
    
    t = 0
    while t < t_end:
        
        U = (k/h**2)*(U[2:] + U[:-2]) + (1-2*k/h**2)*U[1:-1]
        
        U = np.hstack([
            np.array([ts[0]]),
            U,
            np.array([ts[-1]])
        ])
    
        t += k
    
    return U

df = yf.download('TSLA')
prices = df['Close'].values[-100:]
smooth = heat_smooth(prices, h=1, k=0.1, t_end=5)

plt.plot(prices)
plt.plot(smooth)
plt.legend(['Prices', 'Smothened Prices'])
plt.show()