EvaluationEfficiencyOptimiz.../davisAPI/prediction.py
2024-12-01 18:49:24 +04:00

190 lines
7.5 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

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_speed=[]
wind_direction=[]
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_speed.append(wind_speed.magnitude.item())
wind_direction.append(wind_dir.magnitude.item())
return wind_speed,wind_direction
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!")