secondlaw-ai / skyrim

🌎 🤝 AI weather models united
Apache License 2.0
158 stars 7 forks source link

Refactor Skyrim interface for enhanced forecasting #22

Closed m13uz closed 4 months ago

m13uz commented 4 months ago

Nothing changed for the Skyrim.predict() method

from skyrim.core import Skyrim
import datetime

model = Skyrim("pangu", ic_source="cds")
start_time = datetime.datetime(2024, 5, 1)

# Forecast with 2 steps
pred_da = model.forecast(start_time, n_steps=2)
print(pred_da.shape)  
# Output: 
# (3, 69, 721, 1440)

# Forecast with 8 steps for specific channels
pred_da = model.forecast(start_time, n_steps=8, channels=["u10m", "v10m"])
print(pred_da.shape)  
# Output: 
# (9, 2, 721, 1440)

print(pred_da.channel.values)
# Output: 
# ['u10m' 'v10m']

print(pred_da.time.values)  
# Output:
# ['2024-05-01T00:00:00.000000000', '2024-05-01T06:00:00.000000000',
#  '2024-05-01T12:00:00.000000000', '2024-05-01T18:00:00.000000000',
#  '2024-05-02T00:00:00.000000000', '2024-05-02T06:00:00.000000000',
#  '2024-05-02T12:00:00.000000000', '2024-05-02T18:00:00.000000000',
#  '2024-05-03T00:00:00.000000000']