From 0f5217a9632a24da358cbccbd6f781c8b315f91b Mon Sep 17 00:00:00 2001 From: Capucine Lechartre Date: Thu, 2 May 2024 17:14:50 +0200 Subject: [PATCH] works standalone --- neural_lam/models/ar_model.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/neural_lam/models/ar_model.py b/neural_lam/models/ar_model.py index 617ebe3c..e3b421ee 100644 --- a/neural_lam/models/ar_model.py +++ b/neural_lam/models/ar_model.py @@ -879,7 +879,12 @@ def save_pred_as_grib(self, file_path, value_dir_path): 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 + if constants.IS_3D[variable]: + shape_val = 13 + vertical = constants.VERTICAL_LEVELS + else: + shape_val = 1 + vertical = 1 # Find the value range to sample value_range = indices[variable] @@ -891,7 +896,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=grib_code, level=constants.VERTICAL_LEVELS + shortName=grib_code, level=vertical ) md = subset.metadata()