Add "datastores" to represent input data from zarr, npy, etc (#66)

Introduce new "datastores" concept where loading of input data from disk into `pytorch.Dataset` (`neural_lam.WeatherDataset`) is split into two layers: a) datastores that work with and return `xr.DataArray` objects for the whole time-series and b) `neural_lam.WeatherDataset` which consumes output from a datastore, takes care of time-sampling and produces `pytorch.Tensor`-based training samples. Currently, two kinds of datastores are implemented: 1) reading of zarr-based training datasets produced with `mllam-data-prep` and 2) reading of the npyfiles-based MEPS example dataset included with neural-lam `v0.1.0`. --------- Co-authored-by: SimonKamuk <[email protected]> Co-authored-by: joeloskarsson <[email protected]> Co-authored-by: Leif Denby <[email protected]> Co-authored-by: Joel Oskarsson <[email protected]> Co-authored-by: Simon Adamov <[email protected]> Co-authored-by: Simon Adamov <[email protected]> Co-authored-by: Kasper Hintz <[email protected]>
mllam · Nov 21, 2024 · c3c1722 · c3c1722
1 parent 2cc617e
commit c3c1722
Show file tree

Hide file tree

Showing 50 changed files with 5,765 additions and 1,179 deletions.
diff --git a/.github/workflows/ci-pdm-install-and-test-cpu.yml b/.github/workflows/ci-pdm-install-and-test-cpu.yml
@@ -39,17 +39,17 @@ jobs:
       - name: Load cache data
         uses: actions/cache/restore@v4
         with:
-          path: data
-          key: ${{ runner.os }}-meps-reduced-example-data-v0.1.0
+          path: tests/datastore_examples/npyfilesmeps/meps_example_reduced.zip
+          key: ${{ runner.os }}-meps-reduced-example-data-v0.2.0
           restore-keys: |
-            ${{ runner.os }}-meps-reduced-example-data-v0.1.0
+            ${{ runner.os }}-meps-reduced-example-data-v0.2.0
 
       - name: Run tests
         run: |
-          pdm run pytest
+          pdm run pytest -vv -s
 
       - name: Save cache data
         uses: actions/cache/save@v4
         with:
-          path: data
-          key: ${{ runner.os }}-meps-reduced-example-data-v0.1.0
+          path: tests/datastore_examples/npyfilesmeps/meps_example_reduced.zip
+          key: ${{ runner.os }}-meps-reduced-example-data-v0.2.0
diff --git a/.github/workflows/ci-pdm-install-and-test-gpu.yml b/.github/workflows/ci-pdm-install-and-test-gpu.yml
@@ -44,17 +44,17 @@ jobs:
       - name: Load cache data
         uses: actions/cache/restore@v4
         with:
-          path: data
-          key: ${{ runner.os }}-meps-reduced-example-data-v0.1.0
+          path: tests/datastore_examples/npyfilesmeps/meps_example_reduced.zip
+          key: ${{ runner.os }}-meps-reduced-example-data-v0.2.0
           restore-keys: |
-            ${{ runner.os }}-meps-reduced-example-data-v0.1.0
+            ${{ runner.os }}-meps-reduced-example-data-v0.2.0
 
       - name: Run tests
         run: |
-          pdm run pytest
+          pdm run pytest -vv -s
 
       - name: Save cache data
         uses: actions/cache/save@v4
         with:
-          path: data
-          key: ${{ runner.os }}-meps-reduced-example-data-v0.1.0
+          path: tests/datastore_examples/npyfilesmeps/meps_example_reduced.zip
+          key: ${{ runner.os }}-meps-reduced-example-data-v0.2.0
diff --git a/.github/workflows/ci-pip-install-and-test-cpu.yml b/.github/workflows/ci-pip-install-and-test-cpu.yml
@@ -29,17 +29,17 @@ jobs:
       - name: Load cache data
         uses: actions/cache/restore@v4
         with:
-          path: data
-          key: ${{ runner.os }}-meps-reduced-example-data-v0.1.0
+          path: tests/datastore_examples/npyfilesmeps/meps_example_reduced.zip
+          key: ${{ runner.os }}-meps-reduced-example-data-v0.2.0
           restore-keys: |
-            ${{ runner.os }}-meps-reduced-example-data-v0.1.0
+            ${{ runner.os }}-meps-reduced-example-data-v0.2.0
 
       - name: Run tests
         run: |
-          python -m pytest
+          python -m pytest -vv -s
 
       - name: Save cache data
         uses: actions/cache/save@v4
         with:
-          path: data
-          key: ${{ runner.os }}-meps-reduced-example-data-v0.1.0
+          path: tests/datastore_examples/npyfilesmeps/meps_example_reduced.zip
+          key: ${{ runner.os }}-meps-reduced-example-data-v0.2.0
diff --git a/.github/workflows/ci-pip-install-and-test-gpu.yml b/.github/workflows/ci-pip-install-and-test-gpu.yml
@@ -34,17 +34,17 @@ jobs:
       - name: Load cache data
         uses: actions/cache/restore@v4
         with:
-          path: data
-          key: ${{ runner.os }}-meps-reduced-example-data-v0.1.0
+          path: tests/datastore_examples/npyfilesmeps/meps_example_reduced.zip
+          key: ${{ runner.os }}-meps-reduced-example-data-v0.2.0
           restore-keys: |
-            ${{ runner.os }}-meps-reduced-example-data-v0.1.0
+            ${{ runner.os }}-meps-reduced-example-data-v0.2.0
 
       - name: Run tests
         run: |
-          python -m pytest
+          python -m pytest -vv -s
 
       - name: Save cache data
         uses: actions/cache/save@v4
         with:
-          path: data
-          key: ${{ runner.os }}-meps-reduced-example-data-v0.1.0
+          path: tests/datastore_examples/npyfilesmeps/meps_example_reduced.zip
+          key: ${{ runner.os }}-meps-reduced-example-data-v0.2.0
diff --git a/.github/workflows/pre-commit.yml b/.github/workflows/pre-commit.yml
@@ -13,7 +13,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: ["3.9", "3.10", "3.11", "3.12"]
+        python-version: ["3.9", "3.10", "3.11"]
     steps:
     - uses: actions/checkout@v2
     - name: Set up Python

diff --git a/.gitignore b/.gitignore
@@ -1,14 +1,14 @@
 ### Project Specific ###
 wandb
-slurm_log*
 saved_models
 lightning_logs
 data
 graphs
 *.sif
 sweeps
-test_*.sh
 .vscode
+*.html
+*.zarr
 *slurm*
 
 ### Python ###
@@ -75,8 +75,15 @@ tags
 
 # Coc configuration directory
 .vim
+.vscode
+
+# macos
+.DS_Store
+__MACOSX
 
 # pdm (https://pdm-project.org/en/stable/)
 .pdm-python
+.venv
+
 # exclude pdm.lock file so that both cpu and gpu versions of torch will be accepted by pdm
 pdm.lock
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -7,6 +7,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [unreleased](https://github.com/joeloskarsson/neural-lam/compare/v0.2.0...HEAD)
 
+### Added
+
+- Introduce Datastores to represent input data from different sources, including zarr and numpy.
+  [\#66](https://github.com/mllam/neural-lam/pull/66)
+ @leifdenby @sadamov
+
 ## [v0.2.0](https://github.com/joeloskarsson/neural-lam/releases/tag/v0.2.0)
 
 ### Added