Edit predictions

Code/python/notebooks/predict.py

@@ -2,6 +2,7 @@ import torch
 import pathlib
 import pandas as pd
 from aiRNN import dataloader, models, losses
+from aiRNN.preprocessors import denorm_coords
 import numpy as np
 import copy
 
@@ -63,16 +64,17 @@ with torch.no_grad():
         if X_c is not None:
             X_c = X_c.to(DEVICE)
         y_pred, _ = model(X_t, X_f, X_c, warm // step, pred // step)
-        results.append((y.cpu(), y_pred.cpu()))
-    y_true = torch.cat([r[0] for r in results], dim=0)
-    y_pred = torch.cat([r[1] for r in results], dim=0)
-    np_y_all = np.concatenate(
-        [y_true.numpy().reshape(-1, 3), y_pred.numpy().reshape(-1, 3)], axis=1
-    )
+        lat_pred, lon_pred, alt_pred = denorm_coords(y_pred[...,0], y_pred[...,1], y_pred[...,2])
+        lat_true, lon_true, alt_true = denorm_coords(y[...,0], y[...,1], y[...,2])
+        results.append((torch.stack([lat_true, lon_true, alt_true], dim=-1).cpu(), torch.stack([lat_pred, lon_pred, alt_pred], dim=-1).cpu()))
+    np_y_all = np.concatenate([
+        np.concatenate([t.numpy(), p.numpy()], axis=-1)
+        for t, p in results
+    ], axis=0)
     np.savetxt(
         f"predictions_{base_name}_wu{warm}_ps{pred}_aw{altitude_weight}.csv",
         np_y_all,
         delimiter=",",
-        header="true_x,true_y,true_z,pred_x,pred_y,pred_z",
+        header="true_lat,true_lon,true_alt,pred_lat,pred_lon,pred_alt",
         comments="",
     )