Add stub data, healpy

Author: jacobbieker

.github/workflows/workflows.yaml

@@ -11,13 +11,18 @@ jobs:
     strategy:
       fail-fast: true
       matrix:
-        os: [ubuntu-latest]
+        os: [ubuntu-latest, macos-latest]
         python-version: ["3.11", "3.12"]
         torch-version: [2.4.0]
         include:
           - torch-version: 2.4.0
             torchvision-version: 0.19.0
     steps:
+      - uses: prefix-dev/setup-pixi@v0.8.3
+          with:
+            pixi-version: v0.41.4
+            cache: true
+      - run: pixi run test
       - uses: actions/checkout@v2
       - name: Set up Python ${{ matrix.python-version }}
         uses: actions/setup-python@v2
@@ -30,7 +35,7 @@ jobs:
       - name: Install internal dependencies
         run: |
           pip install pyg_lib torch_scatter torch_sparse torch_cluster torch_spline_conv -f https://data.pyg.org/whl/torch-${{ matrix.torch-version}}+cpu.html
-          if [ ${{ matrix.torch-version}} == 2.3.0 ]; then pip install dgl -f https://data.dgl.ai/wheels/torch-2.3/repo.html; fi
+          if [ ${{ matrix.torch-version}} == 2.4.0 ]; then pip install dgl -f https://data.dgl.ai/wheels/torch-2.4/repo.html; fi
       - name: Install main package
         run: |
           pip install -e .
@@ -64,7 +69,7 @@ jobs:
       - name: Run pytest
         run: |
           # import dgl to initialize backend
-          if [ ${{ matrix.torch-version}} == 2.3.0 ]; then python3 -c "import dgl"; fi
+          if [ ${{ matrix.torch-version}} == 2.4.0 ]; then python3 -c "import dgl"; fi
           export PYTEST_COMMAND="pytest $PYTESTCOV $PYTESTXDIST -s"
           echo "Will be running this command: $PYTEST_COMMAND"
           eval $PYTEST_COMMAND

graph_weather/models/layers/grid_to_points.py

@@ -0,0 +1,4 @@
+"""
+This is designed to abstract away the lat/lon points and as a layer that can be put in front of models that more expecta regular lat/lon grid as input.
+"""
+

graph_weather/models/layers/points_to_grid.py

@@ -133,10 +133,10 @@ def __init__(
         self.timesteps = timesteps
 
     def forward(
-        self, x_2d: torch.Tensor, x_3d: torch.Tensor, forecast_steps: int
+        self, surface: torch.Tensor, pressure: torch.Tensor, forecast_steps: int
     ) -> WeatherMeshOutput:
         # Encode input
-        latent = self.encoder(x_2d, x_3d)
+        latent = self.encoder(surface, pressure)
 
         # Apply processors for each forecast step
         for _ in range(forecast_steps):

graph_weather/models/weathermesh/weathermesh2.py

@@ -77,6 +77,7 @@ pysolar = "*"
 rtree = "*"
 pixi-pycharm = ">=0.0.8,<0.0.9"
 uv = ">=0.6.2,<0.7"
+healpy = ">=1.18.1,<2"
 
 
 [tool.pixi.environments]