Skip to content

Commit ed72bc2

Browse files
authored
Reproject ERA5 to EASE2 grid (#249)
* 💚 Update ruff check to also sort imports * ♻️ Specify an encoder for each input dataset in the config * ✅ Fix tests to use an encoder for each input dataset in the config * 🎉 Skeleton for ReprojectionEncoder * 🚚 Convert configure_trainer into build_trainer * 🚚 Inline add_callbacks and add_loggers * 🚚 Change run_directory from a property to a function * ✨ Add area attribute to DataSpace * ✨ Add latitude/longitude attributes to CommonDataModule * ✨ Add a SupportsLatLon protocol for passing lists of latitudes and longitudes * ✨ Pass latitude and longitude information to the ReprojectionEncoder * ✨ Add nearest neighbours calculation to ReprojectionEncoder * ✨ Add a reprojection filter that also updates the dataset geography * 🚚 Moved grid setup into a dedicated factory for easier addition of future grids * 🚚 Consolidate geography related functionality in geotools * 🚚 Renamed filters for consistency * 🐛 Fix signature for nearest_neighbour_indices * 👽 Add support for mars_grid which is required while creating a dataset * 👽 Expanded on GeographicMetadata.get to meet external requirements * 🔧 New dataset config files for regridded ERA5 * 👽 Fix override import for Python 3.11 compatibility * 🏷️ Use typedef shorthands for reprojection arguments * 🐛 Better implementation of Metadata.get that allows additional variables to pass through to the wrapped class * ⚡ Cache ReprojectFilter mappings at class-level to avoid recalculation on each GroupOfDates * 🐛 Add batch normalisation to the reprojection encoder * 🐛 Ensure that lats/lons are also set when loading from checkpoint * 🚨 Use override from typing_extensions for Python 3.11 compatibility * 🚨 Fixed linting/type errors * 👽 Change index lookup to torch.long to avoid indexing error * 🎨 Simplify lat/lon checks in SingleDataset * ⚡ Cache lat/lon checks in SingleDataset * 🔧 Fix template lookups in 25p0km weather configs * 🥅 Catch the unlikely scenario where different timepoints in the input data are using different projections * 🏷️ Fix return type of geography * 🔧 Fix dataset names to match filenames * ✅ Add test for mismatch between config filenames and YAML key/name attribute * 🐛 Ensure that mars_area correctly uses longitudes * 🥅 Catch case where no field can be found in a FieldList * 🏷️ Update return type of nearest_neighbour_indices * 🥅 Catch error when an encoder is missing for an input dataset * ⚡ Cache additional properties in CommonDataModule for speed * ✏️ Fix typo in comment * 🚚 Rename GeographicGrid arguments to avoid confusion * ✅ Add some test for ReprojectionEncoder * 🏷️ Fix shadow typing errors in GeographicMetadata * ➕ Add types-pyyaml * 🚚 Renamed the 25p0km configs to make it clear that these are the newest versions of the ERA5 data * ⚡ Add latitudes/longitudes as ignored hyperparameters of BaseModel * ⚡ Add latitudes/longitudes as parameters of BaseEncoder * ⚡ Update ReprojectionEncoder to read relevant latitudes/longitudes from BaseEncoder * ⚰️ Drop SupportsLatLon protocol * 📌 Update uv lock file * ✅ Simplify TestReprojectionEncoder * ✅ Added dummy latitudes/longitudes arguments to models in testing * 🔧 Add a config file to demonstrate ReprojectionEncoder usage * 📝 Better explanation of mismatching datasets in SingleDataset exceptions * 🔧 Distinguish imp log messages with emoji * 📝 Add new log messages related to model building
1 parent 69757b1 commit ed72bc2

48 files changed

Lines changed: 1465 additions & 193 deletions

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

icenet_mp/cli/main.py

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,6 @@
44
import warnings
55

66
import typer
7-
from hydra.core.utils import simple_stdout_log_config
87

98
from icenet_mp.plugins import register_plugins
109

@@ -13,7 +12,11 @@
1312
from .train import training_cli
1413

1514
# Configure logging
16-
simple_stdout_log_config()
15+
logging.basicConfig(
16+
format="😈 [%(asctime)s] %(message)s",
17+
datefmt=r"%Y-%m-%d %H:%M:%S",
18+
level=logging.INFO,
19+
)
1720
logger = logging.getLogger(__name__)
1821

1922
# Register all plugins
Lines changed: 57 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,57 @@
1+
full-weathernorth-era5-25p0km-1979-2025-24h-v3:
2+
name: full-weathernorth-era5-25p0km-1979-2025-24h-v3
3+
group_as: era5
4+
description: Weather from ERA5 for partial northern hemisphere
5+
attribution: ECMWF/C3S
6+
licence: CC-BY-4.0
7+
8+
dates:
9+
start: '1979-01-01T12:00:00'
10+
end: '2025-12-31T23:00:00'
11+
frequency: 24h
12+
13+
input:
14+
pipe:
15+
- join:
16+
- mars:
17+
use_cdsapi_dataset: "reanalysis-era5-complete"
18+
class: ea
19+
expver: "0001"
20+
grid: "0.25 / 0.25"
21+
area: "90/-180/20/180" # northern hemisphere
22+
levtype: sfc
23+
param:
24+
- 2t # 2-metre air temperature
25+
- sp # Surface pressure
26+
- 10u # 10m x-direction wind speed
27+
- 10v # 10m y-direction wind speed
28+
- msl # Sea level pressure
29+
- mars:
30+
use_cdsapi_dataset: "reanalysis-era5-complete"
31+
class: ea
32+
expver: "0001"
33+
grid: "0.25 / 0.25"
34+
area: "90/-180/20/180" # northern hemisphere
35+
level:
36+
- 10 # hPa
37+
- 250 # hPa
38+
- 500 # hPa
39+
- 1000 # hPa
40+
levtype: pl # pressure level
41+
param:
42+
- z # geopotential height
43+
- t # air temperature
44+
- q # specific humidity
45+
- u # x-direction wind speed
46+
- v # y-direction wind speed
47+
- forcings:
48+
param:
49+
- cos_julian_day
50+
- insolation
51+
- sin_julian_day
52+
template: \${input.pipe.0.join.0.mars}
53+
# Reproject to EASE2 25km grid
54+
- reproject:
55+
crs: "EPSG:6931"
56+
resolution: "25p0km"
57+
shape: [432, 432]
Lines changed: 57 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,57 @@
1+
full-weathersouth-era5-25p0km-1979-2025-24h-v3:
2+
name: full-weathersouth-era5-25p0km-1979-2025-24h-v3
3+
group_as: era5
4+
description: Weather from ERA5 for partial southern hemisphere
5+
attribution: ECMWF/C3S
6+
licence: CC-BY-4.0
7+
8+
dates:
9+
start: '1979-01-01T12:00:00'
10+
end: '2025-12-31T23:00:00'
11+
frequency: 24h
12+
13+
input:
14+
pipe:
15+
- join:
16+
- mars:
17+
use_cdsapi_dataset: "reanalysis-era5-complete"
18+
class: ea
19+
expver: "0001"
20+
grid: "0.25 / 0.25"
21+
area: "-20/-180/-90/180" # southern hemisphere
22+
levtype: sfc
23+
param:
24+
- 2t # 2-metre air temperature
25+
- sp # Surface pressure
26+
- 10u # 10m x-direction wind speed
27+
- 10v # 10m y-direction wind speed
28+
- msl # Sea level pressure
29+
- mars:
30+
use_cdsapi_dataset: "reanalysis-era5-complete"
31+
class: ea
32+
expver: "0001"
33+
grid: "0.25 / 0.25"
34+
area: "-20/-180/-90/180" # southern hemisphere
35+
level:
36+
- 10 # hPa
37+
- 250 # hPa
38+
- 500 # hPa
39+
- 1000 # hPa
40+
levtype: pl # pressure level
41+
param:
42+
- z # geopotential height
43+
- t # air temperature
44+
- q # specific humidity
45+
- u # x-direction wind speed
46+
- v # y-direction wind speed
47+
- forcings:
48+
param:
49+
- cos_julian_day
50+
- insolation
51+
- sin_julian_day
52+
template: \${input.pipe.0.join.0.mars}
53+
# Reproject to EASE2 25km grid
54+
- reproject:
55+
crs: "EPSG:6932"
56+
resolution: "25p0km"
57+
shape: [432, 432]
Lines changed: 57 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,57 @@
1+
samp-weathernorth-era5-25p0km-2017-2019-24h-v3:
2+
name: samp-weathernorth-era5-25p0km-2017-2019-24h-v3
3+
group_as: era5
4+
description: Weather from ERA5 for partial northern hemisphere
5+
attribution: ECMWF/C3S
6+
licence: CC-BY-4.0
7+
8+
dates:
9+
start: '2017-01-01T12:00:00'
10+
end: '2019-01-31T23:00:00'
11+
frequency: 24h
12+
13+
input:
14+
pipe:
15+
- join:
16+
- mars:
17+
use_cdsapi_dataset: "reanalysis-era5-complete"
18+
class: ea
19+
expver: "0001"
20+
grid: "0.25 / 0.25"
21+
area: "90/-180/20/180" # northern hemisphere
22+
levtype: sfc
23+
param:
24+
- 2t # 2-metre air temperature
25+
- sp # Surface pressure
26+
- 10u # 10m x-direction wind speed
27+
- 10v # 10m y-direction wind speed
28+
- msl # Sea level pressure
29+
- mars:
30+
use_cdsapi_dataset: "reanalysis-era5-complete"
31+
class: ea
32+
expver: "0001"
33+
grid: "0.25 / 0.25"
34+
area: "90/-180/20/180" # northern hemisphere
35+
level:
36+
- 10 # hPa
37+
- 250 # hPa
38+
- 500 # hPa
39+
- 1000 # hPa
40+
levtype: pl # pressure level
41+
param:
42+
- z # geopotential height
43+
- t # air temperature
44+
- q # specific humidity
45+
- u # x-direction wind speed
46+
- v # y-direction wind speed
47+
- forcings:
48+
param:
49+
- cos_julian_day
50+
- insolation
51+
- sin_julian_day
52+
template: \${input.pipe.0.join.0.mars}
53+
# Reproject to EASE2 25km grid
54+
- reproject:
55+
crs: "EPSG:6931"
56+
resolution: "25p0km"
57+
shape: [432, 432]
Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,53 @@
1+
samp-weathersouth-era5-25p0km-2017-2019-24h-v3:
2+
name: samp-weathersouth-era5-25p0km-2017-2019-24h-v3
3+
group_as: era5
4+
description: Weather from ERA5 for partial southern hemisphere
5+
attribution: ECMWF/C3S
6+
licence: CC-BY-4.0
7+
8+
dates:
9+
start: '2017-01-01T12:00:00'
10+
end: '2019-01-31T23:00:00'
11+
frequency: 24h
12+
13+
input:
14+
pipe:
15+
- join:
16+
- mars:
17+
use_cdsapi_dataset: "reanalysis-era5-complete"
18+
area: "-20/-180/-90/180" # southern hemisphere
19+
grid: "0.25 / 0.25"
20+
levtype: sfc
21+
param:
22+
- 2t # 2-metre air temperature
23+
- sp # Surface pressure
24+
- 10u # 10m x-direction wind speed
25+
- 10v # 10m y-direction wind speed
26+
- msl # Sea level pressure
27+
- mars:
28+
use_cdsapi_dataset: "reanalysis-era5-complete"
29+
area: "-20/-180/-90/180" # southern hemisphere
30+
grid: "0.25 / 0.25"
31+
level:
32+
- 10 # hPa
33+
- 250 # hPa
34+
- 500 # hPa
35+
- 1000 # hPa
36+
levtype: pl # pressure level
37+
param:
38+
- z # geopotential height
39+
- t # air temperature
40+
- q # specific humidity
41+
- u # x-direction wind speed
42+
- v # y-direction wind speed
43+
- forcings:
44+
param:
45+
- cos_julian_day
46+
- insolation
47+
- sin_julian_day
48+
template: \${input.pipe.0.join.0.mars}
49+
# Reproject to EASE2 25km grid
50+
- reproject:
51+
crs: "EPSG:6932"
52+
resolution: "25p0km"
53+
shape: [432, 432]

icenet_mp/config/model/cnn_null_cnn.yaml

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2,9 +2,12 @@ _target_: icenet_mp.models.EncodeProcessDecode
22

33
name: cnn-null-cnn
44

5-
encoder:
6-
_target_: icenet_mp.models.encoders.CNNEncoder
7-
latent_space: [144, 144] # Shape of the latent space
5+
encoders:
6+
latent_space: [144, 144]
7+
era5:
8+
_target_: icenet_mp.models.encoders.CNNEncoder
9+
sic-icenet:
10+
_target_: icenet_mp.models.encoders.CNNEncoder
811

912
processor:
1013
_target_: icenet_mp.models.processors.NullProcessor

icenet_mp/config/model/cnn_unet_cnn.yaml

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2,9 +2,12 @@ _target_: icenet_mp.models.EncodeProcessDecode
22

33
name: cnn-unet-cnn
44

5-
encoder:
6-
_target_: icenet_mp.models.encoders.CNNEncoder
7-
latent_space: [144, 144] # Shape of the latent space
5+
encoders:
6+
latent_space: [144, 144]
7+
era5:
8+
_target_: icenet_mp.models.encoders.CNNEncoder
9+
sic-icenet:
10+
_target_: icenet_mp.models.encoders.CNNEncoder
811

912
processor:
1013
_target_: icenet_mp.models.processors.UNetProcessor

icenet_mp/config/model/cnn_vit_cnn.yaml

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2,9 +2,12 @@ _target_: icenet_mp.models.EncodeProcessDecode
22

33
name: cnn-vit-cnn
44

5-
encoder:
6-
_target_: icenet_mp.models.encoders.CNNEncoder
7-
latent_space: [144, 144] # Shape of the latent space
5+
encoders:
6+
latent_space: [144, 144]
7+
era5:
8+
_target_: icenet_mp.models.encoders.CNNEncoder
9+
sic-icenet:
10+
_target_: icenet_mp.models.encoders.CNNEncoder
811

912
processor:
1013
_target_: icenet_mp.models.processors.VitProcessor

icenet_mp/config/model/naive_null_naive.yaml

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2,9 +2,12 @@ _target_: icenet_mp.models.EncodeProcessDecode
22

33
name: naive-null-naive
44

5-
encoder:
6-
_target_: icenet_mp.models.encoders.NaiveLinearEncoder
7-
latent_space: [432, 432] # Shape of the latent space
5+
encoders:
6+
latent_space: [432, 432]
7+
era5:
8+
_target_: icenet_mp.models.encoders.NaiveLinearEncoder
9+
sic-icenet:
10+
_target_: icenet_mp.models.encoders.NaiveLinearEncoder
811

912
processor:
1013
_target_: icenet_mp.models.processors.NullProcessor

icenet_mp/config/model/naive_unet_naive.yaml

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2,9 +2,12 @@ _target_: icenet_mp.models.EncodeProcessDecode
22

33
name: naive-unet-naive
44

5-
encoder:
6-
_target_: icenet_mp.models.encoders.NaiveLinearEncoder
7-
latent_space: [432, 432] # Shape of the latent space
5+
encoders:
6+
latent_space: [432, 432]
7+
era5:
8+
_target_: icenet_mp.models.encoders.NaiveLinearEncoder
9+
sic-icenet:
10+
_target_: icenet_mp.models.encoders.NaiveLinearEncoder
811

912
processor:
1013
_target_: icenet_mp.models.processors.UNetProcessor

0 commit comments

Comments
 (0)