Feature: add tests for meps dataset

.github/workflows/run_tests.yml

@@ -0,0 +1,32 @@
+name: Run Unit Test via Pytest
+
+on:
+  # trigger on pushes to any branch, but not main
+  push:
+    branches-ignore:
+      - main
+  # and also on PRs to main
+  pull_request:
+    branches:
+      - main
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version: ["3.9", "3.10", "3.11", "3.12"]
+
+    steps:
+      - uses: actions/checkout@v3
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v4
+        with:
+          python-version: ${{ matrix.python-version }}
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          if [ -f requirements.txt ]; then pip install -r requirements.txt; fi
+      - name: Test with pytest
+        run: |
+          pytest  -v -s

create_mesh.py

@@ -153,7 +153,7 @@ def prepend_node_index(graph, new_index):
     return networkx.relabel_nodes(graph, to_mapping, copy=True)
 
 
-def main():
+def main(input_args=None):
     parser = ArgumentParser(description="Graph generation arguments")
     parser.add_argument(
         "--data_config",
@@ -186,7 +186,7 @@ def main():
         default=0,
         help="Generate hierarchical mesh graph (default: 0, no)",
     )
-    args = parser.parse_args()
+    args = parser.parse_args(input_args)
 
     # Load grid positions
     config_loader = config.Config.from_file(args.data_config)

neural_lam/utils.py

@@ -1,5 +1,6 @@
 # Standard library
 import os
+import shutil
 
 # Third-party
 import numpy as np
@@ -250,7 +251,8 @@ def fractional_plot_bundle(fraction):
     Get the tueplots bundle, but with figure width as a fraction of
     the page width.
     """
-    bundle = bundles.neurips2023(usetex=True, family="serif")
+    usetex = True if shutil.which("latex") else False
+    bundle = bundles.neurips2023(usetex=usetex, family="serif")
     bundle.update(figsizes.neurips2023())
     original_figsize = bundle["figure.figsize"]
     bundle["figure.figsize"] = (

neural_lam/vis.py

@@ -87,7 +87,7 @@ def plot_prediction(
         1,
         2,
         figsize=(13, 7),
-        subplot_kw={"projection": data_config.coords_projection()},
+        subplot_kw={"projection": data_config.coords_projection},
     )
 
     # Plot pred and target
@@ -136,7 +136,7 @@ def plot_spatial_error(error, obs_mask, data_config, title=None, vrange=None):
 
     fig, ax = plt.subplots(
         figsize=(5, 4.8),
-        subplot_kw={"projection": data_config.coords_projection()},
+        subplot_kw={"projection": data_config.coords_projection},
     )
 
     ax.coastlines()  # Add coastline outlines

requirements.txt

@@ -10,6 +10,18 @@ Cartopy>=0.22.0
 pyproj>=3.4.1
 tueplots>=0.0.8
 plotly>=5.15.0
+torch-geometric>=2.5.2
+loguru>=0.7.2
+xarray>=2024.3.0
+zarr>=2.17.2
+dask>=2024.4.2
 
 # for dev
 pre-commit>=2.15.0
+codespell>=2.0.0
+black>=21.9b0
+isort>=5.9.3
+flake8>=4.0.1
+pylint>=3.0.3
+pytest>=8.1.1
+pooch>=1.8.1

tests/test_mllam_dataset.py

@@ -0,0 +1,131 @@
+# Standard library
+import os
+
+# Third-party
+import pooch
+
+# First-party
+from create_mesh import main as create_mesh
+from neural_lam.config import Config
+from neural_lam.utils import load_static_data
+from neural_lam.weather_dataset import WeatherDataset
+from train_model import main as train_model
+
+os.environ["WANDB_DISABLED"] = "true"
+
+
+def test_retrieve_data_ewc():
+    # Initializing variables for the client
+    S3_BUCKET_NAME = "mllam-testdata"
+    S3_ENDPOINT_URL = "https://object-store.os-api.cci1.ecmwf.int"
+    S3_FILE_PATH = "neural-lam/npy/meps_example_reduced.v0.1.0.zip"
+    S3_FULL_PATH = "/".join([S3_ENDPOINT_URL, S3_BUCKET_NAME, S3_FILE_PATH])
+    known_hash = (
+        "98c7a2f442922de40c6891fe3e5d190346889d6e0e97550170a82a7ce58a72b7"
+    )
+
+    pooch.retrieve(
+        url=S3_FULL_PATH,
+        known_hash=known_hash,
+        processor=pooch.Unzip(extract_dir=""),
+        path="data",
+        fname="meps_example_reduced.zip",
+    )
+
+
+def test_load_reduced_meps_dataset():
+    data_config_file = "data/meps_example_reduced/data_config.yaml"
+    dataset_name = "meps_example_reduced"
+
+    dataset = WeatherDataset(dataset_name="meps_example_reduced")
+    config = Config.from_file(data_config_file)
+
+    var_names = config.values["dataset"]["var_names"]
+    var_units = config.values["dataset"]["var_units"]
+    var_longnames = config.values["dataset"]["var_longnames"]
+
+    assert len(var_names) == len(var_longnames)
+    assert len(var_names) == len(var_units)
+
+    # TODO: can these two variables be loaded from elsewhere?
+    n_grid_static_features = 4
+    n_input_steps = 2
+
+    n_forcing_features = config.values["dataset"]["num_forcing_features"]
+    n_state_features = len(var_names)
+    n_prediction_timesteps = dataset.sample_length - n_input_steps
+
+    nx, ny = config.values["grid_shape_state"]
+    n_grid = nx * ny
+
+    # check that the dataset is not empty
+    assert len(dataset) > 0
+
+    # get the first item
+    init_states, target_states, forcing = dataset[0]
+
+    # check that the shapes of the tensors are correct
+    assert init_states.shape == (n_input_steps, n_grid, n_state_features)
+    assert target_states.shape == (
+        n_prediction_timesteps,
+        n_grid,
+        n_state_features,
+    )
+    assert forcing.shape == (
+        n_prediction_timesteps,
+        n_grid,
+        n_forcing_features,
+    )
+
+    static_data = load_static_data(dataset_name=dataset_name)
+
+    required_props = {
+        "border_mask",
+        "grid_static_features",
+        "step_diff_mean",
+        "step_diff_std",
+        "data_mean",
+        "data_std",
+        "param_weights",
+    }
+
+    # check the sizes of the props
+    assert static_data["border_mask"].shape == (n_grid, 1)
+    assert static_data["grid_static_features"].shape == (
+        n_grid,
+        n_grid_static_features,
+    )
+    assert static_data["step_diff_mean"].shape == (n_state_features,)
+    assert static_data["step_diff_std"].shape == (n_state_features,)
+    assert static_data["data_mean"].shape == (n_state_features,)
+    assert static_data["data_std"].shape == (n_state_features,)
+    assert static_data["param_weights"].shape == (n_state_features,)
+
+    assert set(static_data.keys()) == required_props
+
+
+def test_create_graph_reduced_meps_dataset():
+    args = [
+        "--graph=hierarchical",
+        "--hierarchical=1",
+        "--data_config=data/meps_example_reduced/data_config.yaml",
+        "--levels=2",
+    ]
+    create_mesh(args)
+
+
+def test_train_model_reduced_meps_dataset():
+    args = [
+        "--model=hi_lam",
+        "--data_config=data/meps_example_reduced/data_config.yaml",
+        "--n_workers=4",
+        "--epochs=1",
+        "--graph=hierarchical",
+        "--hidden_dim=16",
+        "--hidden_layers=1",
+        "--processor_layers=1",
+        "--ar_steps=1",
+        "--eval=val",
+        "--n_example_pred=0",
+    ]
+    train_model(args)

train_model.py

@@ -22,7 +22,7 @@
 }
 
 
-def main():
+def main(input_args=None):
     """
     Main function for training and evaluating models
     """
@@ -206,7 +206,7 @@ def main():
         default={},
         help="Dict with variables and lead times to log watched metrics for",
     )
-    args = parser.parse_args()
+    args = parser.parse_args(input_args)
 
     config_loader = config.Config.from_file(args.data_config)