fixing relationship to metadata

neural_lam/constants.py

@@ -143,6 +143,24 @@
     "V_10M": 0,
 }
 
+GRIB_NAME = {
+    "PP": 'pres',
+    "QV": 'q',
+    "RELHUM": 'r',
+    "T": 't',
+    "U": 'u',
+    "V": 'v',
+    "W": 'wz',
+    "CLCT": 'ccl',
+    "PMSL": 'prmsl',
+    "PS": 'sp',
+    "T_2M": '2t',
+    "TOT_PREC": 'tp',
+    "U_10M":'10u',
+    "V_10M": '10v'
+}
+
+
 # Vertical level weights
 # These were retrieved based on the pressure levels of
 # https://weatherbench2.readthedocs.io/en/latest/data-guide.html#era5

neural_lam/models/ar_model.py

@@ -876,7 +876,7 @@ def save_pred_as_grib(self, file_path, value_dir_path):
         final_data = earthkit.data.FieldList()
         # Loop through all the time steps and all the variables
         for time_idx, date_str in time_steps.items():
-            for variable in constants.PARAM_NAMES_SHORT:
+            for variable, grib_code in constants.GRIB_NAME.items():
                 # here find the key of the cariable in constants.is_3D
                 #  and if == 7, assign a cut of 7 on the reshape. Else 1
                 shape_val = 13 if constants.IS_3D[variable] else 1
@@ -891,7 +891,7 @@ def save_pred_as_grib(self, file_path, value_dir_path):
                 original_data = earthkit.data.from_source("file", sample_file)
 
                 subset = original_data.sel(
-                    shortName=variable.lower(), level=constants.VERTICAL_LEVELS
+                    shortName=grib_code, level=constants.VERTICAL_LEVELS
                 )
                 md = subset.metadata()