Merge pull request #1 from openclimatefix/pre-commit-ci-update-config

[pre-commit.ci] pre-commit autoupdate

Author: JackKelly

.pre-commit-config.yaml

@@ -1,6 +1,6 @@
 repos:
   - repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v4.1.0
+    rev: v4.3.0
     hooks:
       # list of supported hooks: https://pre-commit.com/hooks.html
       - id: trailing-whitespace
@@ -12,7 +12,7 @@ repos:
 
   # python code formatting
   - repo: https://github.com/psf/black
-    rev: 22.1.0
+    rev: 22.6.0
     hooks:
       - id: black
         args: [--line-length, "99"]
@@ -33,19 +33,19 @@ repos:
 
   # yaml formatting
   - repo: https://github.com/pre-commit/mirrors-prettier
-    rev: v2.5.1
+    rev: v2.7.1
     hooks:
       - id: prettier
         types: [yaml]
 
   # python code analysis
   - repo: https://github.com/PyCQA/flake8
-    rev: 4.0.1
+    rev: 5.0.4
     hooks:
       - id: flake8
 
   # jupyter notebook cell output clearing
   - repo: https://github.com/kynan/nbstripout
-    rev: 0.5.0
+    rev: 0.6.0
     hooks:
       - id: nbstripout

configs/datamodule/xresunet.yaml

@@ -3,5 +3,18 @@ _target_: src.datamodules.xresunet_datamodule.SatDataModule
 data_dir: ${data_dir} # data_dir is specified in config.yaml
 batch_size: 32
 num_workers: 8
-channels: ["IR_016", "IR_039", "IR_087", "IR_097", "IR_108", "IR_120", "IR_134", "VIS006", "VIS008", "WV_062", "WV_073"]
+channels:
+  [
+    "IR_016",
+    "IR_039",
+    "IR_087",
+    "IR_097",
+    "IR_108",
+    "IR_120",
+    "IR_134",
+    "VIS006",
+    "VIS008",
+    "WV_062",
+    "WV_073",
+  ]
 pin_memory: True

configs/model/xresunet.yaml

@@ -4,7 +4,7 @@ weight_decay: 0.0005
 
 net:
   _target_: src.models.components.xresnet_unet.XResUNet
-  input_size: [64,64]
+  input_size: [64, 64]
   forecast_steps: 24
   history_steps: 9
   pretrained: False

src/datamodules/components/sat_dataloader.py

@@ -1,9 +1,9 @@
-import numpy as np
-import xarray as xr
-import fsspec
 import glob
 import os
 
+import fsspec
+import numpy as np
+import xarray as xr
 from torch.utils.data import ConcatDataset, DataLoader, Dataset, random_split
 
 SAT_MEAN = {
@@ -41,11 +41,12 @@
 
 
 def _set_sat_coords(dataset: xr.Dataset) -> xr.Dataset:
-    """Set variables as coordinates"""
+    """Set variables as coordinates."""
     return dataset.set_coords(
         ["time_utc", "channel_name", "y_osgb", "x_osgb", "y_geostationary", "x_geostationary"]
     )
 
+
 def load_netcdf(filename, engine="h5netcdf", *args, **kwargs) -> xr.Dataset:
     """Load a NetCDF dataset from local file system or cloud bucket."""
     with fsspec.open(filename, mode="rb") as file:
@@ -54,7 +55,23 @@ def load_netcdf(filename, engine="h5netcdf", *args, **kwargs) -> xr.Dataset:
 
 
 class Satellite(DataLoader):
-    def __init__(self, channels=["IR_016", "IR_039", "IR_087", "IR_097", "IR_108", "IR_120", "IR_134", "VIS006", "VIS008", "WV_062", "WV_073"], data_dir="./"):
+    def __init__(
+        self,
+        channels=[
+            "IR_016",
+            "IR_039",
+            "IR_087",
+            "IR_097",
+            "IR_108",
+            "IR_120",
+            "IR_134",
+            "VIS006",
+            "VIS008",
+            "WV_062",
+            "WV_073",
+        ],
+        data_dir="./",
+    ):
         self.channels = channels
         self.data_dir = data_dir
         if "HRV" in self.channels:
@@ -114,9 +131,14 @@ def __getitem__(self, item):
         std = np.expand_dims(std, axis=[1, 2, 3])
         hrvsatellite = hrvsatellite - mean
         hrvsatellite = hrvsatellite / std
-        input_data = hrvsatellite.values[:,:7]
-        target_data = hrvsatellite.values[:,7:]
-        merged_data = (
-            np.concatenate((input_data, np.expand_dims(dataset["y_osgb"].values, axis=[1,2]),np.expand_dims(dataset["x_osgb"].values, axis=[1,2])), 1)
-        ) # Now in Batch, Time+Coord, Channel, W, H orderk
+        input_data = hrvsatellite.values[:, :7]
+        target_data = hrvsatellite.values[:, 7:]
+        merged_data = np.concatenate(
+            (
+                input_data,
+                np.expand_dims(dataset["y_osgb"].values, axis=[1, 2]),
+                np.expand_dims(dataset["x_osgb"].values, axis=[1, 2]),
+            ),
+            1,
+        )  # Now in Batch, Time+Coord, Channel, W, H orderk
         return np.squeeze(merged_data), np.squeeze(target_data)

src/datamodules/xresunet_datamodule.py

@@ -3,6 +3,7 @@
 import torch
 from pytorch_lightning import LightningDataModule
 from torch.utils.data import ConcatDataset, DataLoader, Dataset, random_split
+
 from src.datamodules.components.sat_dataloader import Satellite
 
 
@@ -27,7 +28,19 @@ def __init__(
         self,
         data_dir: str = "data/",
         batch_size: int = 64,
-        channels: list = ["IR_016", "IR_039", "IR_087", "IR_097", "IR_108", "IR_120", "IR_134", "VIS006", "VIS008", "WV_062", "WV_073"],
+        channels: list = [
+            "IR_016",
+            "IR_039",
+            "IR_087",
+            "IR_097",
+            "IR_108",
+            "IR_120",
+            "IR_134",
+            "VIS006",
+            "VIS008",
+            "WV_062",
+            "WV_073",
+        ],
         num_workers: int = 0,
         pin_memory: bool = False,
     ):
@@ -51,9 +64,15 @@ def setup(self, stage: Optional[str] = None):
 
         # load datasets only if they're not loaded already
         if not self.data_train and not self.data_val and not self.data_test:
-            self.data_train = Satellite(channels=self.channels, data_dir=self.hparams.data_dir + "/train/")
-            self.data_val = Satellite(channels=self.channels, data_dir=self.hparams.data_dir + "/test/")
-            self.data_test = Satellite(channels=self.channels, data_dir=self.hparams.data_dir + "/test/")
+            self.data_train = Satellite(
+                channels=self.channels, data_dir=self.hparams.data_dir + "/train/"
+            )
+            self.data_val = Satellite(
+                channels=self.channels, data_dir=self.hparams.data_dir + "/test/"
+            )
+            self.data_test = Satellite(
+                channels=self.channels, data_dir=self.hparams.data_dir + "/test/"
+            )
 
     def train_dataloader(self):
         return DataLoader(

src/models/components/xresnet_unet.py

@@ -1,12 +1,9 @@
-from fastai.vision.all import *
-from fastai.callback.wandb import *
 from fastai.callback.tracker import *
+from fastai.callback.wandb import *
 from fastai.distributed import *
-from fastai.vision.models.xresnet import *
 from fastai.layers import Mish
-
-
-
+from fastai.vision.all import *
+from fastai.vision.models.xresnet import *
 from torch import nn
 
 

src/models/xresunet_module.py

@@ -1,9 +1,9 @@
 from typing import Any, List
 
 import torch
-from torch.optim.lr_scheduler import CosineAnnealingLR
-from torch.optim import AdamW
 from pytorch_lightning import LightningModule
+from torch.optim import AdamW
+from torch.optim.lr_scheduler import CosineAnnealingLR
 
 from src.models.components.xresnet_unet import XResUNet
 
@@ -32,7 +32,7 @@ def __init__(
 
         # this line allows to access init params with 'self.hparams' attribute
         # it also ensures init params will be stored in ckpt
-        #self.save_hyperparameters(logger=False)
+        # self.save_hyperparameters(logger=False)
 
         self.net = net
         self.lr = lr
@@ -85,7 +85,9 @@ def test_step(self, batch: Any, batch_idx: int):
         # log test metrics
         self.log("test/loss", loss, on_step=False, on_epoch=True)
 
-        return {"loss": loss, }
+        return {
+            "loss": loss,
+        }
 
     def test_epoch_end(self, outputs: List[Any]):
         pass