1
2
3
4
5

conda env create --file=environment.yml
conda env create --file=environment_cpu.yml

conda activate tfm_env
pip install -e .

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93

import datetime
import yfinance as yf
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import time
from datetime import date
from timesfm import TimesFm
from huggingface_hub import login
import matplotlib.pyplot as plt

# Specify the stock codes, Shanghai stocks end with .ss, Shenzhen stocks end with .sz
start = date(2020, 1, 1)  # Create date object using date class
end = date(2024, 1, 1)  # Set the end date to January 1, 2024
codelist = ["000001.ss"]

# Add error retry mechanism for downloading data
for retry in range(3):  # Try up to 3 times
    try:
        data2 = yf.download(codelist, start=start, end=end)
        # Data preprocessing
        data2 = data2['Adj Close'].dropna()  # Use adjusted close price to remove missing values
        if not data2.empty:
            break  # Successfully downloaded and processed data, exit loop
    except Exception as e:
        print(f"Download failed, attempt {retry+1}. Error: {e}")
        if retry < 2:  # Wait before the final attempt
            time.sleep(5)  # Wait for 5 seconds before retrying

if data2.empty:
    raise ValueError("Data is empty. Please change the time range and try again or check your network connection.")

context_len = 512  # Set context length
horizon_len = 256  # Set prediction period length

if len(data2) < context_len:
    raise ValueError(f"Data length is less than context length ({context_len})")

context_data = data2[-context_len:]  # Use the most recent 512 days of data as context

# Initialize and import the TimesFM model
tfm = TimesFm(
    context_len=context_len,
    horizon_len=horizon_len,
    input_patch_len=32,
    output_patch_len=128,
    num_layers=20,
    model_dims=1280,
    backend='cpu',  # Change here, from 'gpu' to 'cpu'
)

# Log in to Hugging Face Hub, replace ***** with your Hugging token
login("*****")

tfm.load_from_checkpoint(repo_id="google/timesfm-1.0-200m")

# Prepare data
forecast_input = [context_data.values]
frequency_input = [0]  # Set data frequency (0 is high-frequency data)

# Run prediction
point_forecast, experimental_quantile_forecast = tfm.forecast(
    forecast_input,
    freq=frequency_input,
)

# Set chart size to 24*12 inches
plt.figure(figsize=(24, 12))

forecast_dates = pd.date_range(start=data2.index[-1] + pd.Timedelta(days=1), periods=horizon_len, freq='B')
forecast_series = pd.Series(point_forecast[0], index=forecast_dates)

# Additional part: get and plot actual price data from 2024.1.1 to current date
current_date = datetime.datetime.now().date()
data_recent = yf.download(codelist, start=date(2024, 1, 1), end=current_date)

if not data_recent.empty:
    data_recent = data_recent['Adj Close'].dropna()
    plt.plot(data_recent.index, data_recent.values, label="Actual Prices (2024-Now)")

# Create or update the chart (if there was a previous chart, update it here)
plt.plot(data2.index, data2.values, label="Actual Prices")
plt.plot(forecast_series.index, forecast_series.values, label="Forecasted Prices")
plt.xlabel("Date")
plt.ylabel("Price")
plt.title(f"Price Compare & Forecast for {codelist[0]}")
plt.legend()

# Save the chart to a file, ensure the size change is effective
plt.savefig(f'{codelist[0]}_compare.png', bbox_inches='tight') 

# Explicitly close the current chart
plt.close(fig='all')