From b858844059cb15ea0314bd9b53710bba95800037 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?H=C3=A5vard=20Haugen?= <havardhh@met.no>
Date: Mon, 13 Jan 2025 17:17:54 +0200
Subject: [PATCH] Altitude as dataloader property

---
 bris/__main__.py        |  8 +++-----
 bris/data/datamodule.py | 22 ++++++++++++++++++++++
 config/n320.yaml        | 26 +++++++++++++++++++-------
 3 files changed, 44 insertions(+), 12 deletions(-)

diff --git a/bris/__main__.py b/bris/__main__.py
index dd1c08f..4884726 100644
--- a/bris/__main__.py
+++ b/bris/__main__.py
@@ -37,12 +37,10 @@ def main():
         config=config,
         checkpoint_object=checkpoint,
     )
-
+    
     # Assemble outputs
-    run_name = "legendary_gnome"
-    workdir = "testdir"
-    # TODO: Figure out what the leadtimes are based on the config
-    # TODO: Get this from the config
+    run_name = config.run_name #"legendary_gnome"
+    workdir = config.hardware.paths.workdir #"testdir"
     num_members = 1
 
     # Get outputs and required_variables of each decoder
diff --git a/bris/data/datamodule.py b/bris/data/datamodule.py
index 52c5dbb..b91c163 100644
--- a/bris/data/datamodule.py
+++ b/bris/data/datamodule.py
@@ -288,6 +288,28 @@ def longitudes(self) -> tuple:
             return (self.data_reader.longitudes,)
         else:
             return self.data_reader.longitudes
+        
+    @cached_property
+    def altitude(self) -> tuple:
+        """
+        Retrives altitude from geopotential height in the datasets
+        """
+        name_to_index = self.data_reader.name_to_index
+        if isinstance(name_to_index, tuple):
+            altitude = ()
+            for i, n2i in enumerate(name_to_index):
+                if 'z' in n2i.keys():
+                    altitude += (self.data_reader[0][i][n2i['z'],0,:] / 9.81 ,)
+                else:
+                    altitude += (None,)
+        else:
+            if 'z' in name_to_index.keys():
+                altitude = (self.data_reader[0][name_to_index['z'],0,:] / 9.81 ,)
+            else:
+                altitude = (None,)
+        
+        return altitude
+
 
     @cached_property
     def field_shape(self) -> tuple:
diff --git a/config/n320.yaml b/config/n320.yaml
index 07e316e..b70b6d0 100644
--- a/config/n320.yaml
+++ b/config/n320.yaml
@@ -3,8 +3,8 @@ defaults:
   - override hydra/hydra_logging: none
   - _self_
   
-start_date: 2024-01-01T00:00:00
-end_date: 2024-01-02T00:00:00
+start_date: 2023-05-31T18:00:00
+end_date: 2023-06-01T12:00:00 #2024-06-01T01:00:00
 
 checkpoint_path: /pfs/lustrep4/scratch/project_465001383/aifs/experiments/reference/o96_pretrain/checkpoint/3fb85b72ba9647d5ae4fe86ce1b5a885/inference-last.ckpt
 
@@ -14,6 +14,8 @@ frequency: 6h
 
 deterministic: True
 
+run_name: n320_at_ref
+
 global_dataset: ${hardware.paths.data}${hardware.files.global_dataset}
 
 # If the user wants to release GPU cache and memory
@@ -43,6 +45,7 @@ dataloader:
 hardware:
   paths:
     data: /pfs/lustrep4//scratch/project_465001383/aifs/dataset/ERA5/
+    workdir: /pfs/lustrep4//scratch/project_465001383/haugenha/anemoi-training-ref-updated/run-anemoi/lumi/workdir/
   files:
     global_dataset: aifs-od-an-oper-0001-mars-n320-2023-2024-6h-v2.zarr
 
@@ -68,12 +71,21 @@ routing:
     domain: 0
     outputs:
       - netcdf:
-          filename_pattern: era_pred_%Y%m%dT%HZ.nc
-          variables: [2t]
-#      - verif:
-#          filename_pattern: verif/2t/n320_at_ref.nc
-#          variable: 2t
+          filename_pattern: /pfs/lustrep4/scratch/project_465001383/haugenha/anemoi-training-ref-updated/run-anemoi/lumi/predictions/n320_pred_%Y%m%dT%HZ.nc
+          variables: [2t, msl]
+      - verif:
+          filename: /pfs/lustrep4/scratch/project_465001383/haugenha/anemoi-training-ref-updated/run-anemoi/lumi/verif/2t/n320_at_ref.nc
+          variable: 2t
+          units: degC
+          obs_sources:
+            - verif:
+                filename: /pfs/lustrep4/scratch/project_465001383/aifs/old_project_misc/verification/202306_202405/t2m/AIFS_legendary_gnome_roll4.nc
+#      - verif: 
+#          filename_pattern: verif/mslp/n320_at_ref.nc
+#          variable: mslp
 #          obs:
+#            - netcdf:
+#                path: /pfs/lustrep4/scratch/project_465001383/aifs/old_project_misc/verification/202306_202405/mslp/AIFS_legendary_gnome_roll4.nc      
 
   
 reorder: ['10u',