Adressing first comments

neural_lam/constants.py

@@ -201,15 +201,10 @@
 EVAL_PLOT_VARS = ["T_2M"]
 STORE_EXAMPLE_DATA = True
 SELECTED_PROJ = ccrs.PlateCarree()
-EXAMPLE_FILE = "data/cosmo/samples/train/data.zarr"
-SAMPLE_GRIB = "/users/clechart/neural-lam/templates/lfff02180000"
+SAMPLE_GRIB = "neural-lam/templates/lfff02180000"
+SAMPLE_Z_GRIB = "neural-lam/templates/lfff02180000z"
 SAMPLE_Z_GRIB = "/users/clechart/neural-lam/templates/lfff02180000z"
-CHUNK_SIZE = 100
 EVAL_DATETIME = ["2020100215"]
-EVAL_PLOT_VARS = ["QV"]
-STORE_EXAMPLE_DATA = False
-COSMO_PROJ = ccrs.PlateCarree()
-SELECTED_PROJ = COSMO_PROJ
 POLLON = -170.0
 POLLAT = 43.0
 SMOOTH_BOUNDARIES = False

neural_lam/models/ar_model.py

@@ -880,7 +880,7 @@ def save_pred_as_grib(self, file_path, value_dir_path):
                 # here find the key of the cariable in constants.is_3D
                 #  and if == 7, assign a cut of 7 on the reshape. Else 1
                 if constants.IS_3D[variable]:
-                    shape_val = 13
+                    shape_val = len(constants.VERTICAL_LEVELS) 
                     vertical = constants.VERTICAL_LEVELS
                 else:
                     shape_val = 1
@@ -914,7 +914,7 @@ def save_pred_as_grib(self, file_path, value_dir_path):
                     replacement_data = np.load(file_path)
                     original_cut = replacement_data[
                         0, time_idx, :, min(value_range) : max(value_range) + 1
-                    ].reshape(582, 390, shape_val)
+                    ].reshape(constants.GRID_SHAPE[1], constants.GRID_SHAPE[0], shape_val)
                     cut_values = np.moveaxis(
                         original_cut, [-3, -2, -1], [-1, -2, -3]
                     )