[파이썬] 주식 차분 Plot

'''###########################
plot animation (차분)
###########################'''
import pandas_datareader as pdr
from pandas import to_numeric
import datetime
from datetime import timedelta
import numpy as np
import pandas as pd
import warnings
from scipy.ndimage import gaussian_filter1d
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation

code = '069500'
name = 'KODEX200'

test_start = '2022-01-01'
test_end = '2023-12-31'
SIGMA = 4

# 데이터 생성
last_day = 352
df = pdr.DataReader(code, "naver", test_start, test_end)
df = df.apply(to_numeric)
df.columns = ['Open', 'High', 'Low', 'Close', 'Volume']
df['close_df'] = df[['Close']].diff(axis=0)
df = df.dropna()

# plt.figure(figsize=(15, 8))
fig, ax = plt.subplots(figsize=(12, 8))
ax.set_xlim(0, len(df))
ax.set_ylim(min(df['Close']), max(df['Close']))
ax.grid()
ax.legend()

x, y = [], []                                      # x, y, d(차분)
pt_close, = plt.plot([], label='Close')
pt_ga, = plt.plot([], label='Gaussian')
pt_diff, = plt.plot([], label='diff_Gaussian')

def update(frame):
    y.append(frame)
    if len(y) > 1:                                  # 차분을 계산하면 데이터가 하나 작아기기 때문에
        y_tmp = y[1:]                               # 차분에 따라 데이터수를 맞추기 위해서
        x.append(len(x))
        pt_close.set_data(x, y_tmp)
         
        g = gaussian_filter1d(y_tmp, sigma=SIGMA)   # gaussian 적용
        pt_ga.set_data(x, g)
        
        diff = []                                   # 차분 계산
        for i in range(len(y)):
            diff.append(y[i] - y[i-1])
        g_diff = gaussian_filter1d(diff, sigma=SIGMA+1)
        d_min = min(g_diff)
        d_max = max(g_diff)
        # print(len(y), '/', len(x), '/', len(y_tmp), '/', len(g_diff), '/', d_min, '/', d_max, end='')
        d = []
        for i in range(len(g_diff)):
            v = (g_diff[i] * ((g_diff[i] - d_min) / 20)) + 32000
            d.append(v)
        # print('/',len(d), '/', d[-1])
        pt_diff.set_data(x, d[1:])
    return frame,


%matplotlib qt5

# %matplotlib inline
ani = FuncAnimation(fig, update, frames=df['Close'].tolist(), interval=100)
# plt.show()
ani.save('stock_gaussian_diff.gif', writer='imagemagick')