Skip to content

Commit

Permalink
Altitude as dataloader property
Browse files Browse the repository at this point in the history
  • Loading branch information
havardhhaugen committed Jan 13, 2025
1 parent fefb3f4 commit b858844
Show file tree
Hide file tree
Showing 3 changed files with 44 additions and 12 deletions.
8 changes: 3 additions & 5 deletions bris/__main__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
22 changes: 22 additions & 0 deletions bris/data/datamodule.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
26 changes: 19 additions & 7 deletions config/n320.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -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
Expand Down Expand Up @@ -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

Expand All @@ -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',
Expand Down

0 comments on commit b858844

Please sign in to comment.