From a229d96eda4d6ade02c6d9219a03eb888197c68e Mon Sep 17 00:00:00 2001
From: Pavel_Sorokin <sorokin.zxcv@gmail.com>
Date: Sun, 1 Dec 2024 18:49:24 +0400
Subject: [PATCH] fixes

---
 davisAPI/prediction.py | 40 ++++++++++++++++++++++------------------
 1 file changed, 22 insertions(+), 18 deletions(-)

diff --git a/davisAPI/prediction.py b/davisAPI/prediction.py
index 492569c..01bf887 100644
--- a/davisAPI/prediction.py
+++ b/davisAPI/prediction.py
@@ -103,7 +103,7 @@ def create_batch(date, levels, downloads, download_path):
     return batch.regrid(res=0.1)
 
 
-def create_batch_random(levels: tuple[int], date: datetime):
+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")},
@@ -112,7 +112,7 @@ def create_batch_random(levels: tuple[int], date: datetime):
         metadata=Metadata(
             lat=torch.linspace(90, -90, 17),
             lon=torch.linspace(0, 360, 32 + 1)[:-1],
-            time=(date,),
+            time=date,
             atmos_levels=levels,
         ),
     )
@@ -138,24 +138,26 @@ def get_wind_speed_and_direction(prediction, batch: Batch, lat: float, lon: floa
 
     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()
 
-    u_scalar = u_values.item()
-    v_scalar = v_values.item()
+        print("u value:", u_scalar)
+        print("v value:", v_scalar)
 
-    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")
+        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_dir = metpy.calc.wind_direction(u_with_units, v_with_units)
+        wind_speed = metpy.calc.wind_speed(u_with_units, v_with_units)
 
-    wind_dir_text = wind_direction_to_text(wind_dir.magnitude)
-    # Вывод результата
-    print(f"Направление ветра: {wind_dir_text} ({wind_dir:.2f}°)")
-    print(f"Скорость ветра: {wind_speed:.2f} м/с")
-    return wind_dir.magnitude.item(), wind_speed.magnitude.item()
+        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):
@@ -169,12 +171,14 @@ def wind_direction_to_text(wind_dir_deg):
 
 def main():
     levels = (100,)
-    date = datetime(2024, 11, 5, 12)
 
+    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)
+    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