from datetime import datetime
from pathlib import Path
import metpy.calc
import numpy as np
import requests
import torch
import xarray as xr
from aurora import AuroraSmall, Batch, Metadata
from metpy.units import units
def get_download_paths(date):
"""Создает список путей для загрузки данных."""
download_path = Path("~/downloads/hres_0.1").expanduser()
downloads = {}
var_nums = {
"2t": "167", "10u": "165", "10v": "166", "msl": "151", "t": "130",
"u": "131", "v": "132", "q": "133", "z": "129", "slt": "043", "lsm": "172",
}
for v in ["2t", "10u", "10v", "msl", "z", "slt", "lsm"]:
downloads[download_path / date.strftime(f"surf_{v}_%Y-%m-%d.grib")] = (
f"https://data.rda.ucar.edu/ds113.1/"
f"ec.oper.an.sfc/{date.year}{date.month:02d}/ec.oper.an.sfc.128_{var_nums[v]}_{v}."
f"regn1280sc.{date.year}{date.month:02d}{date.day:02d}.grb"
)
for v in ["z", "t", "u", "v", "q"]:
for hour in [0, 6, 12, 18]:
prefix = "uv" if v in {"u", "v"} else "sc"
downloads[download_path / date.strftime(f"atmos_{v}_%Y-%m-%d_{hour:02d}.grib")] = (
f"https://data.rda.ucar.edu/ds113.1/"
f"ec.oper.an.pl/{date.year}{date.month:02d}/ec.oper.an.pl.128_{var_nums[v]}_{v}."
f"regn1280{prefix}.{date.year}{date.month:02d}{date.day:02d}{hour:02d}.grb"
)
return downloads, download_path
def download_data(downloads):
"""Скачивает файлы, если они отсутствуют в целевой директории."""
for target, source in downloads.items():
if not target.exists():
print(f"Downloading {source}")
target.parent.mkdir(parents=True, exist_ok=True)
response = requests.get(source)
response.raise_for_status()
with open(target, "wb") as f:
f.write(response.content)
print("Downloads finished!")
def load_surf(v, v_in_file, download_path, date):
"""Загружает переменные поверхностного уровня или статические переменные."""
ds = xr.open_dataset(download_path / date.strftime(f"surf_{v}_%Y-%m-%d.grib"), engine="cfgrib")
data = ds[v_in_file].values[:2]
data = data[None]
return torch.from_numpy(data)
def load_atmos(v, download_path, date, levels):
"""Загружает атмосферные переменные для заданных уровней давления."""
ds_00 = xr.open_dataset(
download_path / date.strftime(f"atmos_{v}_%Y-%m-%d_00.grib"), engine="cfgrib"
)
ds_06 = xr.open_dataset(
download_path / date.strftime(f"atmos_{v}_%Y-%m-%d_06.grib"), engine="cfgrib"
)
ds_00 = ds_00[v].sel(isobaricInhPa=list(levels))
ds_06 = ds_06[v].sel(isobaricInhPa=list(levels))
data = np.stack((ds_00.values, ds_06.values), axis=0)
data = data[None]
return torch.from_numpy(data)
def create_batch(date, levels, downloads, download_path):
"""Создает объект Batch с данными для модели."""
ds = xr.open_dataset(next(iter(downloads.keys())), engine="cfgrib")
batch = Batch(
surf_vars={
"2t": load_surf("2t", "t2m", download_path, date),
"10u": load_surf("10u", "u10", download_path, date),
"10v": load_surf("10v", "v10", download_path, date),
"msl": load_surf("msl", "msl", download_path, date),
},
static_vars={
"z": load_surf("z", "z", download_path, date)[0, 0],
"slt": load_surf("slt", "slt", download_path, date)[0, 0],
"lsm": load_surf("lsm", "lsm", download_path, date)[0, 0],
},
atmos_vars={
"t": load_atmos("t", download_path, date, levels),
"u": load_atmos("u", download_path, date, levels),
"v": load_atmos("v", download_path, date, levels),
"q": load_atmos("q", download_path, date, levels),
"z": load_atmos("z", download_path, date, levels),
},
metadata=Metadata(
lat=torch.from_numpy(ds.latitude.values),
lon=torch.from_numpy(ds.longitude.values),
time=(date.replace(hour=6),),
atmos_levels=levels,
),
)
return batch.regrid(res=0.1)
def create_batch_random(levels: tuple[int], date: tuple):
"""Создает объект Batch с рандомными данными для модели."""
return Batch(
surf_vars={k: torch.randn(1, 2, 17, 32) for k in ("2t", "10u", "10v", "msl")},
static_vars={k: torch.randn(17, 32) for k in ("lsm", "z", "slt")},
atmos_vars={k: torch.randn(1, 2, 4, 17, 32) for k in ("z", "u", "v", "t", "q")},
metadata=Metadata(
lat=torch.linspace(90, -90, 17),
lon=torch.linspace(0, 360, 32 + 1)[:-1],
time=date,
atmos_levels=levels,
),
)
def run_model(batch):
"""Инициализирует модель AuroraSmall и выполняет предсказание."""
model = AuroraSmall()
model.load_checkpoint("microsoft/aurora", "aurora-0.25-small-pretrained.ckpt")
model.eval()
model = model.to("cpu")
with torch.inference_mode():
prediction = model.forward(batch)
return prediction
def get_wind_speed_and_direction(prediction, batch: Batch, lat: float, lon: float):
target_lat = lat
target_lon = lon
lat_idx = torch.abs(batch.metadata.lat - target_lat).argmin()
lon_idx = torch.abs(batch.metadata.lon - target_lon).argmin()
u_values = prediction.atmos_vars["u"][:, :, :, lat_idx, lon_idx]
v_values = prediction.atmos_vars["v"][:, :, :, lat_idx, lon_idx]
wind_speeds=[]
wind_directions=[]
for i in range(u_values.numel()):
u_scalar = u_values.view(-1)[i].item() # Разворачиваем тензор в одномерный и берем элемент
v_scalar = v_values.view(-1)[i].item()
print("u value:", u_scalar)
print("v value:", v_scalar)
u_with_units = u_scalar * units("m/s")
v_with_units = v_scalar * units("m/s")
# Рассчитайте направление и скорость ветра
wind_dir = metpy.calc.wind_direction(u_with_units, v_with_units)
wind_speed = metpy.calc.wind_speed(u_with_units, v_with_units)
wind_speeds.append(wind_speed.magnitude.item())
wind_directions.append(wind_dir.magnitude.item())
return wind_speeds,wind_directions
def wind_direction_to_text(wind_dir_deg):
directions = [
"север", "северо-восток", "восток", "юго-восток",
"юг", "юго-запад", "запад", "северо-запад"
]
idx = int((wind_dir_deg + 22.5) // 45) % 8
return directions[idx]
def main():
levels = (100,)
date1 = datetime(2024, 11, 27, 12)
date2 = datetime(2024, 11, 28, 12)
date_tuple = (date1, date2,)
# downloads, download_path = get_download_paths(date)
# download_data(downloads) # Скачиваем данные, если их нет
# batch_actual = create_batch(date, levels, downloads, download_path)
batch_actual = create_batch_random(levels, date_tuple)
prediction_actual = run_model(batch_actual)
wind_speed_and_direction = get_wind_speed_and_direction(prediction_actual, batch_actual, 50, 20)
return wind_speed_and_direction
if __name__ == "__main__":
main()
print("Prediction completed!")