Repository: openclimatefix/satflow (language: python)

Issue:
Incorporate NWP data

Create output video visualizations 
Add creating video visualizations with the output from SatFlow models, possibly overlaying the cloud mask as well. The main idea is for the HuggingFace Spaces visualization (https://huggingface.co/spaces/openclimatefix/MetNet) so that we can have interactive widgets showing off the models we are training! 

@peterdudfield Would the visualizations you've been making work for making videos?


Discussion / hints:

Yea I'm sure we can do something

I was essentially just using - https://plotly.com/python/animations/

Let me know how I could help

For some not very clean code - https://github.com/openclimatefix/nowcasting_dataset/blob/main/notebooks/2021-08-26/video.py
Okay, great, thanks! I'll start looking through it
Love this idea!  And thanks loads for jumping in, @peterdudfield!

Fix patch:
diff --git a/.github/workflows/publish.yaml b/.github/workflows/publish.yaml
index 436afc9a..629df3f0 100644
--- a/.github/workflows/publish.yaml
+++ b/.github/workflows/publish.yaml
@@ -1,4 +1,3 @@
-
 # This workflow will upload a Python Package using Twine when a release is created
 # For more information see: https://help.github.com/en/actions/language-and-framework-guides/using-python-with-github-actions#publishing-to-package-registries
 
@@ -15,23 +14,22 @@ on:
 
 jobs:
   deploy:
-
     runs-on: ubuntu-latest
 
     steps:
-    - uses: actions/checkout@v2
-    - name: Set up Python
-      uses: actions/setup-python@v2
-      with:
-        python-version: '3.x'
-    - name: Install dependencies
-      run: |
-        python -m pip install --upgrade pip
-        pip install build
-    - name: Build package
-      run: python -m build
-    - name: Publish package
-      uses: pypa/gh-action-pypi-publish@27b31702a0e7fc50959f5ad993c78deac1bdfc29
-      with:
-        user: __token__
-        password: ${{ secrets.PYPI_API_TOKEN }}
+      - uses: actions/checkout@v2
+      - name: Set up Python
+        uses: actions/setup-python@v2
+        with:
+          python-version: "3.x"
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install build
+      - name: Build package
+        run: python -m build
+      - name: Publish package
+        uses: pypa/gh-action-pypi-publish@27b31702a0e7fc50959f5ad993c78deac1bdfc29
+        with:
+          user: __token__
+          password: ${{ secrets.PYPI_API_TOKEN }}
diff --git a/requirements.txt b/requirements.txt
index 793694e2..39cb9f0e 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -12,10 +12,11 @@ pytorch-msssim==0.2.1
 rich>=10.6.0
 torchvision>=0.10.0
 affine==2.3.0
-webdataset
 torch_optimizer
 huggingface_hub>=0.0.16
-einops==0.3.2
+einops>=0.3.2
 metnet>=0.0.3
 skillful_nowcasting>=0.0.2
 perceiver-model>=0.7.0
+nowcasting-dataset
+nowcasting-utils
diff --git a/satflow/configs/callbacks/default.yaml b/satflow/configs/callbacks/default.yaml
index 78ab1da5..655b10f3 100644
--- a/satflow/configs/callbacks/default.yaml
+++ b/satflow/configs/callbacks/default.yaml
@@ -17,4 +17,4 @@ early_stopping:
   min_delta: 0 # minimum change in the monitored metric needed to qualify as an improvement
 
 model_logging:
-  _target_: satflow.core.callbacks.NeptuneModelLogger
+  _target_: nowcasting_utils.training.callbacks.NeptuneModelLogger
diff --git a/satflow/configs/datamodule/cloudgan.yaml b/satflow/configs/datamodule/cloudgan.yaml
deleted file mode 100644
index b096883f..00000000
--- a/satflow/configs/datamodule/cloudgan.yaml
+++ /dev/null
@@ -1,43 +0,0 @@
-_target_: satflow.data.datamodules.SatFlowDataModule
-
-batch_size: 8
-data_dir: ${data_dir} # data_dir is specified in config.yaml
-shuffle: 0
-sources:
-  train: "satflow-train-{00000..00214}.tar"
-  val: "satflow-val-{00000..00060}.tar"
-  test: "satflow-test-{00000..00057}.tar"
-num_workers: 12
-pin_memory: True
-config:
-  visualize: False
-  num_timesteps: 6
-  skip_timesteps: 1
-  forecast_times: 24
-  output_shape: 64
-  target_type: "cloudmask"
-  num_crops: 10
-  use_topo: True
-  use_latlon: True
-  use_time: False
-  time_aux: False
-  use_mask: False
-  use_image: True
-  time_as_channels: True
-  add_pixel_coords: False
-  bands:
-    [
-      "HRV",
-      "IR016",
-      "IR039",
-      "IR087",
-      "IR097",
-      "IR108",
-      "IR120",
-      "IR134",
-      "VIS006",
-      "VIS008",
-      "WV062",
-      "WV073",
-    ]
-  transforms: {}
diff --git a/satflow/configs/datamodule/cloudgan_timestep.yaml b/satflow/configs/datamodule/cloudgan_timestep.yaml
deleted file mode 100644
index e700ea6b..00000000
--- a/satflow/configs/datamodule/cloudgan_timestep.yaml
+++ /dev/null
@@ -1,43 +0,0 @@
-_target_: satflow.data.datamodules.SatFlowDataModule
-
-batch_size: 8
-data_dir: ${data_dir} # data_dir is specified in config.yaml
-shuffle: 0
-sources:
-  train: "satflow-train-{00000..00214}.tar"
-  val: "satflow-val-{00000..00060}.tar"
-  test: "satflow-test-{00000..00057}.tar"
-num_workers: 12
-pin_memory: True
-config:
-  visualize: False
-  num_timesteps: 6
-  skip_timesteps: 1
-  forecast_times: 48
-  output_shape: 128
-  target_type: "cloudmask"
-  num_crops: 10
-  use_topo: True
-  use_latlon: True
-  use_time: False
-  time_aux: False
-  use_mask: False
-  use_image: True
-  time_as_channels: False
-  add_pixel_coords: False
-  bands:
-    [
-      "HRV",
-      "IR016",
-      "IR039",
-      "IR087",
-      "IR097",
-      "IR108",
-      "IR120",
-      "IR134",
-      "VIS006",
-      "VIS008",
-      "WV062",
-      "WV073",
-    ]
-  transforms: {}
diff --git a/satflow/configs/datamodule/maskflow.yaml b/satflow/configs/datamodule/maskflow.yaml
deleted file mode 100644
index 61eda84a..00000000
--- a/satflow/configs/datamodule/maskflow.yaml
+++ /dev/null
@@ -1,29 +0,0 @@
-_target_: satflow.data.datamodules.MaskFlowDataModule
-
-batch_size: 6
-data_dir: ${data_dir} # data_dir is specified in config.yaml
-shuffle: 0
-sources:
-  train: "satflow-train-{00000..00214}.tar"
-  val: "satflow-val-{00000..00060}.tar"
-  test: "satflow-test-{00000..00057}.tar"
-num_workers: 16
-pin_memory: True
-config:
-  visualize: False
-  num_timesteps: 5
-  skip_timesteps: 1
-  forecast_times: 5
-  output_shape: 128
-  target_type: "cloudmask"
-  num_crops: 10
-  use_topo: True
-  use_latlon: False
-  use_time: False
-  time_aux: False
-  use_mask: True
-  use_image: False
-  add_pixel_coords: True
-  time_as_channels: False
-  bands: ["IR016"]
-  transforms: {}
diff --git a/satflow/configs/datamodule/metnet.yaml b/satflow/configs/datamodule/metnet.yaml
index 69424b4e..3c0f6192 100644
--- a/satflow/configs/datamodule/metnet.yaml
+++ b/satflow/configs/datamodule/metnet.yaml
@@ -1,46 +1,17 @@
 _target_: satflow.data.datamodules.SatFlowDataModule
 
-batch_size: 2
-data_dir: ${data_dir} # data_dir is specified in config.yaml
-shuffle: 0
-sources:
-  train: "satflow-flow-144-tiled-{00001..00105}.tar"
-  val: "satflow-flow-144-tiled-{00106..00149}.tar"
-  test: "satflow-flow-144-tiled-{00130..00149}.tar"
-num_workers: 12
+temp_path: "."
+n_train_data: 24900
+n_val_data: 1000
+cloud: "aws"
+num_workers: 8
 pin_memory: True
-prefetch_factor: 4
-config:
-  visualize: False
-  num_timesteps: 6
-  skip_timesteps: 1
-  forecast_times: 24
-  output_shape: 256
-  output_target: 16
-  target_type: "cloudmask"
-  num_crops: 10
-  use_topo: True
-  use_latlon: True
-  use_time: False
-  time_aux: False
-  use_mask: False
-  use_image: False
-  add_pixel_coords: False
-  time_as_channels: False
-  metnet_normalization: True
-  bands:
-    [
-      "HRV",
-      "IR016",
-      "IR039",
-      "IR087",
-      "IR097",
-      "IR108",
-      "IR120",
-      "IR134",
-      "VIS006",
-      "VIS008",
-      "WV062",
-      "WV073",
-    ]
-  transforms: {}
+data_path: "prepared_ML_training_data/v4/"
+fake_data: False
+
+# SatFlow options
+num_timesteps: 6
+forecast_steps: 12
+topo: True
+latlon: True
+metnet_normalization: True
diff --git a/satflow/configs/datamodule/optical_flow.yaml b/satflow/configs/datamodule/optical_flow.yaml
deleted file mode 100644
index bc671b86..00000000
--- a/satflow/configs/datamodule/optical_flow.yaml
+++ /dev/null
@@ -1,43 +0,0 @@
-_target_: satflow.data.datamodules.MaskFlowDataModule
-batch_size: 1
-data_dir: ${data_dir} # data_dir is specified in config.yaml
-shuffle: 0
-sources:
-  train: "satflow-train-{00000..00214}.tar"
-  val: "satflow-val-{00000..00060}.tar"
-  test: "satflow-test-{00000..00057}.tar"
-num_workers: 1
-pin_memory: False
-config:
-  visualize: False
-  num_timesteps: 1
-  skip_timesteps: 1
-  forecast_times: 48
-  output_shape: 400
-  target_type: "cloudmask"
-  num_crops: 1
-  use_topo: False
-  use_latlon: False
-  use_time: False
-  time_aux: False
-  use_mask: False
-  use_image: True
-  add_pixel_coords: False
-  time_as_channels: False
-  # NIR1.6, VIS0.8 and VIS0.6 RGB for near normal view
-  bands:
-    [
-      "HRV",
-      "IR016",
-      "IR039",
-      "IR087",
-      "IR097",
-      "IR108",
-      "IR120",
-      "IR134",
-      "VIS006",
-      "VIS008",
-      "WV062",
-      "WV073",
-    ]
-  transforms: {}
diff --git a/satflow/configs/datamodule/past15hr.yaml b/satflow/configs/datamodule/past15hr.yaml
deleted file mode 100644
index 6d6e4157..00000000
--- a/satflow/configs/datamodule/past15hr.yaml
+++ /dev/null
@@ -1,44 +0,0 @@
-_target_: satflow.data.datamodules.SatFlowDataModule
-
-batch_size: 6
-data_dir: ${data_dir} # data_dir is specified in config.yaml
-shuffle: 0
-sources:
-  train: "satflow-train-{00000..00214}.tar"
-  val: "satflow-val-{00000..00060}.tar"
-  test: "satflow-test-{00000..00057}.tar"
-num_workers: 12
-pin_memory: True
-config:
-  visualize: False
-  num_timesteps: 5
-  skip_timesteps: 3
-  forecast_times: 24
-  output_shape: 64
-  output_target: 64
-  target_type: "cloudmask"
-  num_crops: 10
-  use_topo: True
-  use_latlon: False
-  use_time: False
-  time_aux: False
-  use_mask: False
-  use_image: True
-  add_pixel_coords: False
-  time_as_channels: False
-  bands:
-    [
-      "HRV",
-      "IR016",
-      "IR039",
-      "IR087",
-      "IR097",
-      "IR108",
-      "IR120",
-      "IR134",
-      "VIS006",
-      "VIS008",
-      "WV062",
-      "WV073",
-    ]
-  transforms: {}
diff --git a/satflow/configs/datamodule/past30min.yaml b/satflow/configs/datamodule/past30min.yaml
deleted file mode 100644
index af78cca2..00000000
--- a/satflow/configs/datamodule/past30min.yaml
+++ /dev/null
@@ -1,44 +0,0 @@
-_target_: satflow.data.datamodules.SatFlowDataModule
-
-batch_size: 6
-data_dir: ${data_dir} # data_dir is specified in config.yaml
-shuffle: 0
-sources:
-  train: "satflow-flow-144-tiled-{00001..00105}.tar"
-  val: "satflow-flow-144-tiled-{00106..00129}.tar"
-  test: "satflow-flow-144-tiled-{00130..00149}.tar"
-num_workers: 12
-pin_memory: True
-config:
-  visualize: False
-  num_timesteps: 5
-  skip_timesteps: 1
-  forecast_times: 24
-  output_shape: 64
-  output_target: 64
-  target_type: "cloudmask"
-  num_crops: 10
-  use_topo: True
-  use_latlon: False
-  use_time: False
-  time_aux: False
-  use_mask: False
-  use_image: True
-  add_pixel_coords: False
-  time_as_channels: False
-  bands:
-    [
-      "HRV",
-      "IR016",
-      "IR039",
-      "IR087",
-      "IR097",
-      "IR108",
-      "IR120",
-      "IR134",
-      "VIS006",
-      "VIS008",
-      "WV062",
-      "WV073",
-    ]
-  transforms: {}
diff --git a/satflow/configs/datamodule/pix2pix.yaml b/satflow/configs/datamodule/pix2pix.yaml
deleted file mode 100644
index d0235592..00000000
--- a/satflow/configs/datamodule/pix2pix.yaml
+++ /dev/null
@@ -1,43 +0,0 @@
-_target_: satflow.data.datamodules.SatFlowDataModule
-
-batch_size: 16
-data_dir: ${data_dir} # data_dir is specified in config.yaml
-shuffle: 0
-sources:
-  train: "satflow-train-{00000..00214}.tar"
-  val: "satflow-val-{00000..00060}.tar"
-  test: "satflow-test-{00000..00057}.tar"
-num_workers: 12
-pin_memory: True
-config:
-  visualize: False
-  num_timesteps: 3
-  skip_timesteps: 1
-  forecast_times: 6
-  output_shape: 128
-  target_type: "cloudmask"
-  num_crops: 10
-  use_topo: True
-  use_latlon: True
-  use_time: False
-  time_aux: False
-  use_mask: False
-  use_image: True
-  time_as_channels: True
-  add_pixel_coords: False
-  bands:
-    [
-      "HRV",
-      "IR016",
-      "IR039",
-      "IR087",
-      "IR097",
-      "IR108",
-      "IR120",
-      "IR134",
-      "VIS006",
-      "VIS008",
-      "WV062",
-      "WV073",
-    ]
-  transforms: {}
diff --git a/satflow/configs/datamodule/preprocess.yaml b/satflow/configs/datamodule/preprocess.yaml
deleted file mode 100644
index 50fed744..00000000
--- a/satflow/configs/datamodule/preprocess.yaml
+++ /dev/null
@@ -1,43 +0,0 @@
-_target_: satflow.data.datamodules.FileDataModule
-
-batch_size: 4
-data_dir: ${data_dir} # data_dir is specified in config.yaml
-shuffle: 0
-sources:
-  train: "preprocessed_satflow_prev6"
-  val: "preprocessed_satflow_prev6"
-  test: "preprocessed_satflow_prev6"
-num_workers: 16
-pin_memory: True
-config:
-  visualize: False
-  num_timesteps: 6
-  skip_timesteps: 1
-  forecast_times: 48
-  output_shape: 128
-  target_type: "cloudmask"
-  num_crops: 10
-  use_topo: True
-  use_latlon: False
-  use_time: True
-  time_aux: False
-  use_mask: False
-  use_image: True
-  time_as_channels: False
-  add_pixel_coords: False
-  bands:
-    [
-      "HRV",
-      "IR016",
-      "IR039",
-      "IR087",
-      "IR097",
-      "IR108",
-      "IR120",
-      "IR134",
-      "VIS006",
-      "VIS008",
-      "WV062",
-      "WV073",
-    ]
-  transforms: {}
diff --git a/satflow/configs/datamodule/satflow.yaml b/satflow/configs/datamodule/satflow.yaml
deleted file mode 100644
index d14df5ed..00000000
--- a/satflow/configs/datamodule/satflow.yaml
+++ /dev/null
@@ -1,31 +0,0 @@
-_target_: satflow.data.datamodules.SatFlowDataModule
-
-batch_size: 12
-data_dir: ${data_dir} # data_dir is specified in config.yaml
-shuffle: 0
-sources:
-  train: "satflow-train-{00000..00214}.tar"
-  val: "satflow-val-{00000..00060}.tar"
-  test: "satflow-test-{00000..00057}.tar"
-num_workers: 12
-pin_memory: True
-config:
-  visualize: False
-  num_timesteps: 6
-  skip_timesteps: 1
-  forecast_times: 24
-  output_shape: 128
-  target_type: "cloudmask"
-  num_crops: 10
-  num_times: 15
-  use_topo: False
-  use_latlon: False
-  use_time: False
-  time_aux: False
-  use_mask: False
-  use_image: False
-  time_as_channels: True
-  add_pixel_coords: False
-  add_polar_coords: False
-  bands: ["HRV", "IR016", "VIS006", "VIS008"]
-  transforms: {}
diff --git a/satflow/configs/datamodule/segmentation.yaml b/satflow/configs/datamodule/segmentation.yaml
deleted file mode 100644
index 9eb592eb..00000000
--- a/satflow/configs/datamodule/segmentation.yaml
+++ /dev/null
@@ -1,45 +0,0 @@
-_target_: satflow.data.datamodules.SatFlowDataModule
-
-batch_size: 1
-data_dir: ${data_dir} # data_dir is specified in config.yaml
-shuffle: 0
-sources:
-  train: "satflow-train-{00000..00214}.tar"
-  val: "satflow-val-{00000..00060}.tar"
-  test: "satflow-test-{00000..00057}.tar"
-num_workers: 12
-pin_memory: True
-config:
-  visualize: False
-  num_timesteps: 6
-  skip_timesteps: 1
-  forecast_times: 6
-  output_shape: 128
-  target_type: "cloudmask"
-  num_crops: 10
-  num_times: 10
-  use_topo: True
-  use_latlon: True
-  use_time: False
-  time_aux: False
-  use_mask: False
-  use_image: False
-  time_as_channels: True
-  add_pixel_coords: True
-  add_polar_coords: False
-  bands:
-    [
-      "HRV",
-      "IR016",
-      "IR039",
-      "IR087",
-      "IR097",
-      "IR108",
-      "IR120",
-      "IR134",
-      "VIS006",
-      "VIS008",
-      "WV062",
-      "WV073",
-    ]
-  transforms: {}
diff --git a/satflow/configs/datamodule/unet.yaml b/satflow/configs/datamodule/unet.yaml
deleted file mode 100644
index 85af201d..00000000
--- a/satflow/configs/datamodule/unet.yaml
+++ /dev/null
@@ -1,42 +0,0 @@
-_target_: satflow.data.datamodules.SatFlowDataModule
-
-batch_size: 8
-data_dir: ${data_dir} # data_dir is specified in config.yaml
-shuffle: 0
-sources:
-  train: "satflow-train-{00000..00214}.tar"
-  val: "satflow-val-{00000..00060}.tar"
-  test: "satflow-test-{00000..00057}.tar"
-num_workers: 12
-pin_memory: True
-config:
-  visualize: False
-  num_timesteps: 6
-  skip_timesteps: 1
-  forecast_times: 48
-  output_shape: 128
-  target_type: "cloudmask"
-  num_crops: 10
-  use_topo: True
-  use_latlon: True
-  use_time: False
-  time_aux: False
-  use_mask: False
-  use_image: False
-  time_as_channels: True
-  bands:
-    [
-      "HRV",
-      "IR016",
-      "IR039",
-      "IR087",
-      "IR097",
-      "IR108",
-      "IR120",
-      "IR134",
-      "VIS006",
-      "VIS008",
-      "WV062",
-      "WV073",
-    ]
-  transforms: {}
diff --git a/satflow/core/callbacks.py b/satflow/core/callbacks.py
deleted file mode 100644
index 676ae5d6..00000000
--- a/satflow/core/callbacks.py
+++ /dev/null
@@ -1,38 +0,0 @@
-from pytorch_lightning import Callback, LightningModule, Trainer
-from pytorch_lightning.callbacks import ModelCheckpoint
-import os
-
-
-class NeptuneModelLogger(Callback):
-    """
-    Saves out the last and best models after each validation epoch. If the files don't exists, does nothing.
-
-    Example::
-        from pl_bolts.callbacks import NeptuneModelLogger
-        trainer = Trainer(callbacks=[NeptuneModelLogger()])
-    """
-
-    def __init__(self) -> None:
-        super().__init__()
-
-    def on_validation_epoch_end(self, trainer: Trainer, pl_module: LightningModule) -> None:
-        try:
-            trainer.logger.experiment[0]["model_checkpoints/last.ckpt"].upload(
-                os.path.join(trainer.default_root_dir, "checkpoints", "last.ckpt")
-            )
-        except:
-            print(
-                f"No file to upload at {os.path.join(trainer.default_root_dir, 'checkpoints', 'last.ckpt')}"
-            )
-            pass
-
-    def on_fit_end(self, trainer: Trainer, pl_module: LightningModule) -> None:
-        try:
-            trainer.logger.experiment[0]["model_checkpoints/best.ckpt"].upload(
-                os.path.join(trainer.default_root_dir, "checkpoints", "best.ckpt"),
-            )
-        except:
-            print(
-                f"No file to upload at {os.path.join(trainer.default_root_dir, 'checkpoints', 'best.ckpt')}"
-            )
-            pass
diff --git a/satflow/core/schema.py b/satflow/core/schema.py
deleted file mode 100644
index e69de29b..00000000
diff --git a/satflow/core/training_utils.py b/satflow/core/training_utils.py
deleted file mode 100644
index 9a6656f9..00000000
--- a/satflow/core/training_utils.py
+++ /dev/null
@@ -1,104 +0,0 @@
-import argparse
-
-import torch
-import webdataset as wds
-from torch.utils.data import DataLoader
-
-import satflow.data.datasets
-from satflow.core.utils import load_config
-from satflow.data.datasets import get_dataset
-
-
-def get_loaders(config):
-    """
-    Get Dataloaders for train/test
-    Args:
-        config: Dict, configuration dictionary for the dataset
-
-    Returns:
-        Dict[Dataloader] containing the train and test dataloaders
-    """
-    print(config)
-    train_dset = wds.WebDataset(config["sources"]["train"])
-    val_dset = wds.WebDataset(config["sources"]["val"])
-    test_dset = wds.WebDataset(config["sources"]["test"])
-    train_dataset = get_dataset(config["name"])([train_dset], config=config, train=True)
-    val_dataset = get_dataset(config["name"])([val_dset], config=config, train=False)
-    test_dataset = get_dataset(config["name"])([test_dset], config=config, train=False)
-
-    train_dataloader = DataLoader(
-        train_dataset,
-        num_workers=config["num_workers"],
-        batch_size=config["batch_size"],
-        pin_memory=True,
-    )
-    val_dataloader = DataLoader(
-        val_dataset,
-        num_workers=config["num_workers"],
-        batch_size=config["batch_size"],
-        pin_memory=True,
-    )
-    test_dataloader = DataLoader(
-        test_dataset,
-        num_workers=config["num_workers"],
-        batch_size=config["batch_size"],
-        pin_memory=True,
-    )
-
-    return {"train": train_dataloader, "val": val_dataloader, "test": test_dataloader}
-
-
-def setup_experiment(args):
-    """
-    Sets up the basic logging, etc. common things for running experiments
-
-    Args:
-        args: Commandline arguments
-
-    Returns:
-
-    """
-
-    config = load_config(args.config)
-
-    config["dataset"]["num_workers"] = args.num_workers
-    return config
-
-
-def get_args():
-
-    parser = argparse.ArgumentParser(description="SatFlow")
-
-    # cuda
-    parser.add_argument(
-        "--with_cpu",
-        default=False,
-        action="store_true",
-        help="use CPU in case there's no GPU support",
-    )
-
-    # train
-    parser.add_argument(
-        "-c", "--config", default="./config.yaml", type=str, help="Path to Config File"
-    )
-    parser.add_argument(
-        "--local_rank",
-        type=int,
-        default=-1,
-        help="local rank passed from distributed launcher",
-    )
-
-    parser.add_argument(
-        "-nw",
-        "--num_workers",
-        type=int,
-        default=1,
-        help="Number of dataloader workers",
-    )
-
-    # Include DeepSpeed configuration arguments
-    # parser = deepspeed.add_config_arguments(parser)
-
-    args = parser.parse_args()
-
-    return args
diff --git a/satflow/core/utils.py b/satflow/core/utils.py
index e601551f..00ebb0d3 100644
--- a/satflow/core/utils.py
+++ b/satflow/core/utils.py
@@ -81,13 +81,9 @@ def extras(config: DictConfig) -> None:
         else channels
     )
     log.info(f"Channels: (Use Time) {channels}")
-    channels = (
-        channels + 2 if config.datamodule.config.get("add_pixel_coords", False) else channels
-    )
+    channels = channels + 2 if config.datamodule.config.get("add_pixel_coords", False) else channels
     log.info(f"Channels: (Add Pixel Coordinates) {channels}")
-    channels = (
-        channels + 1 if config.datamodule.config.get("add_polar_coords", False) else channels
-    )
+    channels = channels + 1 if config.datamodule.config.get("add_polar_coords", False) else channels
     log.info(f"Channels: (Add Polar Coordinates) {channels}")
     config.model.input_channels = channels
 
diff --git a/satflow/data/datamodules.py b/satflow/data/datamodules.py
index a4be914c..eb34b56e 100644
--- a/satflow/data/datamodules.py
+++ b/satflow/data/datamodules.py
@@ -1,340 +1,208 @@
-import pytorch_lightning as pl
-from torch.utils.data import DataLoader
-from typing import Optional
-import webdataset as wds
-from satflow.data.datasets import SatFlowDataset, CloudFlowDataset, FileDataset, PerceiverDataset
 import os
-from glob import glob
-
-
-def is_streaming(pattern):
+from nowcasting_dataset.dataset.datasets import worker_init_fn
+from satflow.data.datasets import SatFlowDataset
+from typing import Union, List, Tuple, Optional
+from nowcasting_dataset.consts import (
+    SATELLITE_DATA,
+    SATELLITE_X_COORDS,
+    SATELLITE_Y_COORDS,
+    SATELLITE_DATETIME_INDEX,
+    NWP_DATA,
+    NWP_Y_COORDS,
+    NWP_X_COORDS,
+    DATETIME_FEATURE_NAMES,
+)
+import logging
+import torch
+from pytorch_lightning import LightningDataModule
+
+
+_LOG = logging.getLogger(__name__)
+_LOG.setLevel(logging.DEBUG)
+
+
+class SatFlowDataModule(LightningDataModule):
     """
-    Determine whether Webdataset is being streamed in or not
-
-    Very simple for now and kinda hacky
-    Args:
-        pattern:
-
-    Returns:
-
+    Example of LightningDataModule for NETCDF dataset.
+    A DataModule implements 5 key methods:
+        - prepare_data (things to do on 1 GPU/TPU, not on every GPU/TPU in distributed mode)
+        - setup (things to do on every accelerator in distributed mode)
+        - train_dataloader (the training dataloader)
+        - val_dataloader (the validation dataloader(s))
+        - test_dataloader (the test dataloader(s))
+    This allows you to share a full dataset without explaining how to download,
+    split, transform and process the data.
+    Read the docs:
+        https://pytorch-lightning.readthedocs.io/en/latest/extensions/datamodules.html
     """
-    if "pipe" in pattern:
-        return True
-    else:
-        return False
-
 
-class SatFlowDataModule(pl.LightningDataModule):
     def __init__(
         self,
-        config: dict,
-        sources: dict,
-        batch_size: int = 2,
-        shuffle: int = 0,
-        data_dir: str = "./",
-        num_workers: int = 1,
+        temp_path: str = ".",
+        n_train_data: int = 24900,
+        n_val_data: int = 1000,
+        cloud: str = "aws",
+        num_workers: int = 8,
         pin_memory: bool = True,
-        prefetch_factor: int = 2,
+        data_path="prepared_ML_training_data/v4/",
+        fake_data: bool = False,
+        required_keys: Union[Tuple[str], List[str]] = [
+            NWP_DATA,
+            NWP_X_COORDS,
+            NWP_Y_COORDS,
+            SATELLITE_DATA,
+            SATELLITE_X_COORDS,
+            SATELLITE_Y_COORDS,
+            SATELLITE_DATETIME_INDEX,
+        ]
+        + list(DATETIME_FEATURE_NAMES),
+        history_minutes: Optional[int] = None,
+        forecast_minutes: Optional[int] = None,
     ):
+        """
+        fake_data: random data is created and used instead. This is useful for testing
+        """
         super().__init__()
-        self.data_dir = data_dir
-        self.config = config
-        self.batch_size = batch_size
-        self.shuffle = shuffle
-        self.sources = sources
+
+        self.temp_path = temp_path
+        self.data_path = data_path
+        self.cloud = cloud
+        self.n_train_data = n_train_data
+        self.n_val_data = n_val_data
         self.num_workers = num_workers
         self.pin_memory = pin_memory
-        self.prefetch_factor = prefetch_factor
-
-        self.training_dataloader_ref = None
+        self.fake_data = fake_data
+        self.required_keys = required_keys
+        self.forecast_minutes = forecast_minutes
+        self.history_minutes = history_minutes
 
-    def prepare_data(self):
-        # download
-        pass
+        self.dataloader_config = dict(
+            pin_memory=self.pin_memory,
+            num_workers=self.num_workers,
+            prefetch_factor=8,
+            worker_init_fn=worker_init_fn,
+            persistent_workers=True,
+            # Disable automatic batching because dataset
+            # returns complete batches.
+            batch_size=None,
+        )
 
-    def setup(self, stage: Optional[str] = None):
-        # Assign train/val datasets for use in dataloaders
-        if stage == "fit" or stage is None:
-            train_dset = wds.WebDataset(
-                self.sources["train"]
-                if is_streaming(self.sources["train"])
-                else os.path.join(self.data_dir, self.sources["train"])
-            )
-            val_dset = wds.WebDataset(
-                self.sources["val"]
-                if is_streaming(self.sources["val"])
-                else os.path.join(self.data_dir, self.sources["val"])
+    def train_dataloader(self):
+        if self.fake_data:
+            train_dataset = FakeDataset(
+                history_minutes=self.history_minutes, forecast_minutes=self.forecast_minutes
             )
-            if self.shuffle > 0:
-                # Add shuffling, each sample is still quite large, so too many examples ends up running out of ram
-                train_dset = train_dset.shuffle(self.shuffle)
-            self.train_dataset = SatFlowDataset([train_dset], config=self.config, train=True)
-            self.val_dataset = SatFlowDataset([val_dset], config=self.config, train=False)
-            # This seems necessary for the reload_dataloader to not reload the training_dataloader
-            if self.training_dataloader_ref is None:
-                training_dataloader = DataLoader(
-                    self.train_dataset,
-                    batch_size=self.batch_size,
-                    pin_memory=self.pin_memory,
-                    num_workers=self.num_workers,
-                    prefetch_factor=self.prefetch_factor,
-                )
-                self.training_dataloader_ref = training_dataloader
-
-        # Assign test dataset for use in dataloader(s)
-        if stage == "test" or stage is None:
-            test_dset = wds.WebDataset(
-                self.sources["test"]
-                if is_streaming(self.sources["test"])
-                else os.path.join(self.data_dir, self.sources["test"])
+        else:
+            train_dataset = SatFlowDataset(
+                self.n_train_data,
+                os.path.join(self.data_path, "train"),
+                os.path.join(self.temp_path, "train"),
+                cloud=self.cloud,
+                required_keys=self.required_keys,
+                history_minutes=self.history_minutes,
+                forecast_minutes=self.forecast_minutes,
             )
-            self.test_dataset = SatFlowDataset([test_dset], config=self.config, train=False)
 
-    def train_dataloader(self):
-        return self.training_dataloader_ref
+        return torch.utils.data.DataLoader(train_dataset, **self.dataloader_config)
 
     def val_dataloader(self):
-        return DataLoader(
-            self.val_dataset,
-            batch_size=self.batch_size,
-            pin_memory=self.pin_memory,
-            num_workers=self.num_workers // 2,
-            prefetch_factor=self.prefetch_factor,
-        )
-
-    def test_dataloader(self):
-        return DataLoader(
-            self.test_dataset,
-            batch_size=self.batch_size,
-            pin_memory=self.pin_memory,
-            num_workers=self.num_workers // 2,
-            prefetch_factor=self.prefetch_factor,
-        )
-
-
-class FileDataModule(pl.LightningDataModule):
-    def __init__(
-        self,
-        config: dict,
-        sources: dict,
-        batch_size: int = 2,
-        shuffle: int = 0,
-        data_dir: str = "./",
-        num_workers: int = 1,
-        pin_memory: bool = True,
-    ):
-        super().__init__()
-        self.data_dir = data_dir
-        self.config = config
-        self.batch_size = batch_size
-        self.shuffle = shuffle
-        self.sources = sources
-        self.num_workers = num_workers
-        self.pin_memory = pin_memory
-
-        self.training_dataloader_ref = None
-
-    def prepare_data(self):
-        # download
-        pass
-
-    def setup(self, stage: Optional[str] = None):
-        # Assign train/val datasets for use in dataloaders
-        if stage == "fit" or stage is None:
-            self.train_dataset = FileDataset(
-                os.path.join(self.data_dir, self.sources["train"]),
-                use_image=self.config.get("use_image", True),
-                train=True,
+        if self.fake_data:
+            val_dataset = FakeDataset(
+                history_minutes=self.history_minutes, forecast_minutes=self.forecast_minutes
             )
-            self.val_dataset = FileDataset(
-                os.path.join(self.data_dir, self.sources["val"]),
-                use_image=self.config.get("use_image", True),
-                train=False,
+        else:
+            val_dataset = SatFlowDataset(
+                self.n_val_data,
+                os.path.join(self.data_path, "validation"),
+                os.path.join(self.temp_path, "validation"),
+                cloud=self.cloud,
+                required_keys=self.required_keys,
+                history_minutes=self.history_minutes,
+                forecast_minutes=self.forecast_minutes,
             )
 
-        # Assign test dataset for use in dataloader(s)
-        if stage == "test" or stage is None:
-            self.test_dataset = FileDataset(
-                os.path.join(self.data_dir, self.sources["test"]),
-                use_image=self.config.get("use_image", True),
-                train=False,
-            )
-
-    def train_dataloader(self):
-        if self.training_dataloader_ref:
-            return self.training_dataloader_ref
-
-        training_dataloader = DataLoader(
-            self.train_dataset,
-            batch_size=self.batch_size,
-            pin_memory=self.pin_memory,
-            num_workers=self.num_workers,
-            shuffle=True,
-        )
-        self.training_dataloader_ref = training_dataloader
-
-        return self.training_dataloader_ref
-
-    def val_dataloader(self):
-        return DataLoader(
-            self.val_dataset,
-            batch_size=self.batch_size,
-            pin_memory=self.pin_memory,
-            num_workers=self.num_workers,
-        )
+        return torch.utils.data.DataLoader(val_dataset, **self.dataloader_config)
 
     def test_dataloader(self):
-        return DataLoader(
-            self.test_dataset,
-            batch_size=self.batch_size,
-            pin_memory=self.pin_memory,
-            num_workers=self.num_workers,
-        )
-
-
-class MaskFlowDataModule(pl.LightningDataModule):
-    def __init__(
-        self,
-        config: dict,
-        sources: dict,
-        batch_size: int = 2,
-        shuffle: int = 0,
-        data_dir: str = "./",
-        num_workers: int = 1,
-        pin_memory: bool = True,
-    ):
-        super().__init__()
-        self.data_dir = data_dir
-        self.config = config
-        self.batch_size = batch_size
-        self.shuffle = shuffle
-        self.sources = sources
-        self.num_workers = num_workers
-        self.pin_memory = pin_memory
-
-        self.training_dataloader_ref = None
-
-    def prepare_data(self):
-        # download
-        pass
-
-    def setup(self, stage: Optional[str] = None):
-        # Assign train/val datasets for use in dataloaders
-        if stage == "fit" or stage is None:
-            train_dset = wds.WebDataset(os.path.join(self.data_dir, self.sources["train"]))
-            val_dset = wds.WebDataset(os.path.join(self.data_dir, self.sources["val"]))
-            if self.shuffle > 0:
-                # Add shuffling, each sample is still quite large, so too many examples ends up running out of ram
-                train_dset = train_dset.shuffle(self.shuffle)
-            self.train_dataset = CloudFlowDataset([train_dset], config=self.config, train=True)
-            self.val_dataset = CloudFlowDataset([val_dset], config=self.config, train=False)
-
-        # Assign test dataset for use in dataloader(s)
-        if stage == "test" or stage is None:
-            test_dset = wds.WebDataset(os.path.join(self.data_dir, self.sources["test"]))
-            self.test_dataset = CloudFlowDataset([test_dset], config=self.config, train=False)
+        if self.fake_data:
+            test_dataset = FakeDataset(
+                history_minutes=self.history_minutes, forecast_minutes=self.forecast_minutes
+            )
+        else:
+            # TODO need to change this to a test folder
+            test_dataset = SatFlowDataset(
+                self.n_val_data,
+                os.path.join(self.data_path, "validation"),
+                os.path.join(self.temp_path, "validation"),
+                cloud=self.cloud,
+                required_keys=self.required_keys,
+                history_minutes=self.history_minutes,
+                forecast_minutes=self.forecast_minutes,
+            )
 
-    def train_dataloader(self):
-        # Stores reference and returns it for the reload_dataloaders_every_n_epochs so that the training dataloader keeps
-        # iterating through all the data, but the validation dataloader is reset and helps with keeping the same examples
-        if self.training_dataloader_ref:
-            return self.training_dataloader_ref
+        return torch.utils.data.DataLoader(test_dataset, **self.dataloader_config)
 
-        training_dataloader = DataLoader(
-            self.train_dataset,
-            batch_size=self.batch_size,
-            pin_memory=self.pin_memory,
-            num_workers=self.num_workers,
-        )
-        self.training_dataloader_ref = training_dataloader
 
-        return self.training_dataloader_ref
+class FakeDataset(torch.utils.data.Dataset):
+    """Fake dataset."""
 
-    def val_dataloader(self):
-        return DataLoader(
-            self.val_dataset,
-            batch_size=self.batch_size,
-            pin_memory=self.pin_memory,
-            num_workers=self.num_workers,
-        )
-
-    def test_dataloader(self):
-        return DataLoader(
-            self.test_dataset,
-            batch_size=self.batch_size,
-            pin_memory=self.pin_memory,
-            num_workers=self.num_workers,
-        )
-
-
-class PerceiverDataModule(pl.LightningDataModule):
     def __init__(
         self,
-        config: dict,
-        sources: dict,
-        batch_size: int = 2,
-        shuffle: int = 0,
-        data_dir: str = "./",
-        num_workers: int = 1,
-        pin_memory: bool = True,
+        batch_size=32,
+        width=16,
+        height=16,
+        number_sat_channels=12,
+        length=10,
+        history_minutes=30,
+        forecast_minutes=30,
     ):
-        super().__init__()
-        self.data_dir = data_dir
-        self.config = config
         self.batch_size = batch_size
-        self.shuffle = shuffle
-        self.sources = sources
-        self.num_workers = num_workers
-        self.pin_memory = pin_memory
-
-        self.training_dataloader_ref = None
-
-    def prepare_data(self):
-        # download
+        if history_minutes is None or forecast_minutes is None:
+            history_minutes = 30  # Half an hour
+            forecast_minutes = 240  # 4 hours
+        self.history_steps = history_minutes // 5
+        self.forecast_steps = forecast_minutes // 5
+        self.seq_length = self.history_steps + 1
+        self.width = width
+        self.height = height
+        self.number_sat_channels = number_sat_channels
+        self.length = length
+
+    def __len__(self):
+        return self.length
+
+    def per_worker_init(self, worker_id: int):
         pass
 
-    def setup(self, stage: Optional[str] = None):
-        # Assign train/val datasets for use in dataloaders
-        if stage == "fit" or stage is None:
-            train_dset = wds.WebDataset(os.path.join(self.data_dir, self.sources["train"]))
-            val_dset = wds.WebDataset(os.path.join(self.data_dir, self.sources["val"]))
-            if self.shuffle > 0:
-                # Add shuffling, each sample is still quite large, so too many examples ends up running out of ram
-                train_dset = train_dset.shuffle(self.shuffle)
-            self.train_dataset = PerceiverDataset([train_dset], config=self.config, train=True)
-            self.val_dataset = PerceiverDataset([val_dset], config=self.config, train=False)
-
-        # Assign test dataset for use in dataloader(s)
-        if stage == "test" or stage is None:
-            test_dset = wds.WebDataset(os.path.join(self.data_dir, self.sources["test"]))
-            self.test_dataset = PerceiverDataset([test_dset], config=self.config, train=False)
-
-    def train_dataloader(self):
-        if self.training_dataloader_ref:
-            return self.training_dataloader_ref
-
-        training_dataloader = DataLoader(
-            self.train_dataset,
-            batch_size=self.batch_size,
-            pin_memory=self.pin_memory,
-            num_workers=self.num_workers,
+    def __getitem__(self, idx):
+
+        x = {
+            SATELLITE_DATA: torch.randn(
+                self.batch_size, self.seq_length, self.width, self.height, self.number_sat_channels
+            ),
+            NWP_DATA: torch.randn(self.batch_size, 10, self.seq_length, 2, 2),
+            "hour_of_day_sin": torch.randn(self.batch_size, self.seq_length),
+            "hour_of_day_cos": torch.randn(self.batch_size, self.seq_length),
+            "day_of_year_sin": torch.randn(self.batch_size, self.seq_length),
+            "day_of_year_cos": torch.randn(self.batch_size, self.seq_length),
+        }
+
+        # add fake x and y coords, and make sure they are sorted
+        x[SATELLITE_X_COORDS], _ = torch.sort(torch.randn(self.batch_size, self.seq_length))
+        x[SATELLITE_Y_COORDS], _ = torch.sort(
+            torch.randn(self.batch_size, self.seq_length), descending=True
         )
-        self.training_dataloader_ref = training_dataloader
-
-        return self.training_dataloader_ref
 
-    def val_dataloader(self):
-        return DataLoader(
-            self.val_dataset,
-            batch_size=self.batch_size,
-            pin_memory=self.pin_memory,
-            num_workers=self.num_workers,
-        )
-
-    def test_dataloader(self):
-        return DataLoader(
-            self.test_dataset,
-            batch_size=self.batch_size,
-            pin_memory=self.pin_memory,
-            num_workers=self.num_workers,
-        )
+        # add sorted (fake) time series
+        x[SATELLITE_DATETIME_INDEX], _ = torch.sort(torch.randn(self.batch_size, self.seq_length))
+
+        y = {
+            SATELLITE_DATA: torch.randn(
+                self.batch_size,
+                self.forecast_steps,
+                self.width,
+                self.height,
+                self.number_sat_channels,
+            ),
+        }
+        return x, y
diff --git a/satflow/data/datasets.py b/satflow/data/datasets.py
index 26482169..c53bcfc1 100644
--- a/satflow/data/datasets.py
+++ b/satflow/data/datasets.py
@@ -1,700 +1,94 @@
-import glob
-from typing import Iterator, List, Type, Union, Tuple
-
-import albumentations as A
-import numpy as np
-import torch.utils.data as thd
-import webdataset as wds
-from torch.utils.data.dataset import T_co
-import logging
-import pickle
-import random
-import os
-from satflow.data.utils.normalization import metnet_normalization, standard_normalization
-from satflow.data.utils.utils import (
-    create_time_layer,
-    load_np,
-    binarize_mask,
-    create_pixel_coord_layers,
-    check_channels,
-    crop_center,
-    load_config,
+from typing import Tuple, Union, List
+from nowcasting_dataset.dataset.datasets import NetCDFDataset
+from nowcasting_dataset.consts import (
+    SATELLITE_DATA,
+    SATELLITE_X_COORDS,
+    SATELLITE_Y_COORDS,
+    SATELLITE_DATETIME_INDEX,
+    NWP_DATA,
+    NWP_Y_COORDS,
+    NWP_X_COORDS,
+    NWP_TARGET_TIME,
+    DATETIME_FEATURE_NAMES,
 )
 
-logger = logging.getLogger("satflow.dataset")
-logger.setLevel(logging.WARN)
-
-REGISTERED_DATASET_CLASSES = {}
-
-
-def register_dataset(cls: Type[thd.IterableDataset]):
-    global REGISTERED_DATASET_CLASSES
-    name = cls.__name__
-    assert name not in REGISTERED_DATASET_CLASSES, f"exists class: {REGISTERED_DATASET_CLASSES}"
-    REGISTERED_DATASET_CLASSES[name] = cls
-    return cls
-
-
-def get_dataset(name: str) -> Type[thd.IterableDataset]:
-    global REGISTERED_DATASET_CLASSES
-    assert name in REGISTERED_DATASET_CLASSES, f"available class: {REGISTERED_DATASET_CLASSES}"
-    return REGISTERED_DATASET_CLASSES[name]
-
-
-# Taken from training set
-MSG_MEAN = {
-    "HRV": 14.04300588,
-    "IR016": 12.08545261,
-    "IR039": 277.3749233,
-    "IR087": 269.09229239,
-    "IR097": 246.08281192,
-    "IR108": 271.22961027,
-    "IR120": 269.87252372,
-    "IR134": 251.17556403,
-    "VIS006": 12.23808318,
-    "VIS008": 14.80151262,
-    "WV062": 232.57341978,
-    "WV073": 248.14469363,
-}
-
-MSG_STD = {
-    "HRV": 7.6144786,
-    "IR016": 6.70064364,
-    "IR039": 10.4374892,
-    "IR087": 13.27530427,
-    "IR097": 6.9411872,
-    "IR108": 14.14880209,
-    "IR120": 14.15595176,
-    "IR134": 8.41474376,
-    "VIS006": 7.16105213,
-    "VIS008": 8.04250388,
-    "WV062": 4.20345723,
-    "WV073": 6.93812301,
-}
-
-TOPO_MEAN = 224.3065682349895
-TOPO_STD = 441.7514422990341
-
-
-@register_dataset
-class SatFlowDataset(thd.IterableDataset, wds.Shorthands, wds.Composable):
-    def __init__(self, datasets: List[wds.WebDataset], config: dict, train: bool = True):
-        super().__init__()
-        self.config = config
-        self.datasets = datasets
-        self.train = train
-        self.num_timesteps = config["num_timesteps"]
-        self.skip_timesteps = config.get(
-            "skip_timesteps", 1
-        )  # Every nth historical timestep to take
-        self.forecast_times = config.get(
-            "forecast_times", 48
-        )  # Max timesteps to predict ahead (minutes / 5) default 4 hours
-
-        # Defined output sizes, etc.
-        self.output_shape = config["output_shape"]
-        self.output_target = config.get("output_target", config["output_shape"])
-        self.target_type = config.get("target", "cloudmask")
-        # Should load the common data here
-        self.bands = config.get(
-            "bands",
-            (
-                "HRV",  # mask[mask == 255] = 0  # hardcoded incase missing any of the cloud mask background problem
-                # mask[mask < 2] = 0  # 2 is cloud, others are clear over land (1) and clear over water (0)
-                # mask[mask > 0] = 1
-                "IR016",
-                "IR039",
-                "IR087",
-                "IR097",
-                "IR108",
-                "IR120",
-                "IR134",
-                "VIS006",
-                "VIS008",
-                "WV062",
-                "WV073",
-            ),
-        )
-
-        self.mean = np.array([MSG_MEAN[b] for b in self.bands])
-        self.std = np.array([MSG_STD[b] for b in self.bands])
-        # Make it match the dimensions of the output so that it can be broadcasted
-        self.mean = np.expand_dims(self.mean, (0, 2, 3))
-        self.std = np.expand_dims(self.std, (0, 2, 3))
-
-        self.use_topo = config.get("use_topo", False)
-        self.use_latlon = config.get("use_latlon", False)
-        self.use_time = config.get("use_time", True)
-        self.time_aux = config.get("time_aux", False)  # Time as an auxiliary input as 1D array
-        self.use_mask = config.get("use_mask", True)
-        self.use_image = config.get("use_image", False)
-        self.return_target_stack = config.get("stack_targets", False)
-        self.time_as_channels = config.get("time_as_channels", False)
-        self.add_pixel_coords = config.get("add_pixel_coords", False)
-        self.metnet_norm = config.get("metnet_normalization", False)
-
-        # Number of channels in final one
-        self.num_channels = check_channels(config)
-        self.num_bands = len(self.bands)
-        self.total_per_timestep_channels = (
-            self.num_bands + 3 if self.use_time and not self.time_aux else self.num_bands
-        )
-        self.total_per_timestep_channels = (
-            self.total_per_timestep_channels + 1
-            if self.use_mask
-            else self.total_per_timestep_channels
-        )
-        self.input_cube = np.empty(
-            (self.num_timesteps + 1, self.num_channels, self.output_shape, self.output_shape)
-        )
-        self.input_mask_cube = np.empty(
-            (self.num_timesteps + 1, 1, self.output_shape, self.output_shape)
-        )
-        self.target_cube = np.empty((self.forecast_times, 1, self.output_shape, self.output_shape))
-        self.target_image_cube = np.empty(
-            (self.forecast_times, self.num_bands, self.output_shape, self.output_shape)
-        )
-        self.image_input = True
-        self.pixel_coords = create_pixel_coord_layers(
-            self.output_shape, self.output_shape, with_r=config.get("add_polar_coords", False)
-        )
-        self.pixel_coords = np.squeeze(self.pixel_coords)
-
-        self.topo = None
-        self.location = None
-
-        self.num_crops = config.get("num_crops", 5)
-        self.num_times = config.get("num_times", 10)
-
-        self.vis = config.get("visualize", False)
-
-        transforms = []
-        if self.train and False:
-            # TODO Change if we want to actually flip things
-            # Pointed out that flips might mess up learning dominant winds, etc. physical phenomena, disable for now
-            transforms = [
-                A.HorizontalFlip(p=0.5),
-                A.VerticalFlip(p=0.5),
-            ]
-        if self.train:
-            transforms.append(A.RandomCrop(self.output_shape, self.output_shape))
-        else:
-            transforms.append(
-                A.RandomCrop(self.output_shape, self.output_shape)
-            )  # TODO Make sure is reproducible
-        self.aug = A.ReplayCompose(
-            transforms,
-        )
-        self.replay = None
-
-    def visualize(self, image: np.ndarray, target_image: np.ndarray, mask: np.ndarray):
-        import matplotlib.pyplot as plt
-
-        fig, axs = plt.subplots(
-            len(self.bands) + 1, self.num_timesteps + 1 + self.forecast_times, figsize=(15, 15)
-        )
-        for t_step, img in enumerate(image):
-            for channel, channel_img in enumerate(img):
-                if channel >= len(self.bands):
-                    break
-                # Now should be 2D array
-                axs[channel, t_step].imshow(channel_img)
-                # axs[channel,t_step].set_title(f"{self.bands[channel]} T{'+' if t_step-self.num_timesteps >= 0 else '-'}{t_step - self.num_timesteps}")
-        for t_step, img in enumerate(target_image):
-            for channel, channel_img in enumerate(img):
-                if channel >= len(self.bands):
-                    break
-                # Now should be 2D array
-                axs[channel, t_step + self.num_timesteps + 1].imshow(channel_img)
-                # axs[channel,t_step+self.num_timesteps+1].set_title(f"{self.bands[channel]} T{'+' if t_step+1 >= 0 else '-'}{t_step+1}")
-        for t_step, img in enumerate(mask):
-            # Now should be 2D array
-            axs[-1, t_step + self.num_timesteps + 1].imshow(img[0])
-            # axs[-1,t_step+self.num_timesteps+1].set_title(f"Mask T{'+' if t_step+1 >= 0 else '-'}{t_step+1}")
-        for ax in axs.flat:
-            ax.label_outer()
-        plt.show()
-        plt.close()
-
-    def create_target_time_cube(self, target_timestep: int) -> np.ndarray:
-        """Create target time layer"""
-        time_cube = np.zeros(
-            (1, self.forecast_times, self.output_shape, self.output_shape), dtype=np.int8
-        )
-        time_cube[:, target_timestep, :, :] = 1
-        return time_cube
 
-    def create_target_time_layer(self, target_timestep: int) -> np.ndarray:
+class SatFlowDataset(NetCDFDataset):
+    """Loads data saved by the `prepare_ml_training_data.py` script.
+    Adapted from predict_pv_yield
+    """
+
+    def __init__(
+        self,
+        n_batches: int,
+        src_path: str,
+        tmp_path: str,
+        cloud: str = "gcp",
+        required_keys: Union[Tuple[str], List[str]] = [
+            NWP_DATA,
+            NWP_X_COORDS,
+            NWP_Y_COORDS,
+            SATELLITE_DATA,
+            SATELLITE_X_COORDS,
+            SATELLITE_Y_COORDS,
+            SATELLITE_DATETIME_INDEX,
+        ]
+        + list(DATETIME_FEATURE_NAMES),
+        history_minutes: int = 30,
+        forecast_minutes: int = 60,
+        current_timestep_index: int = 7,
+    ):
         """
-        Creates a one-hot encoded layer, a lot more space efficient than timecube, but needs to be an aux layer
         Args:
-            target_timestep: The target timestep for predicting
-
-        Returns:
-            1-D numpy array where there is a 1 for the target timestep
-        """
-
-        time_layer = np.zeros((self.output_shape, self.forecast_times), dtype=np.int8)
-        time_layer[target_timestep] = 1
-        return time_layer
-
-    def get_timestep(
-        self, sample: dict, idx: int, return_target=False, return_image=True
-    ) -> Union[Tuple[np.ndarray, np.ndarray], Tuple[None, np.ndarray]]:
+          n_batches: Number of batches available on disk.
+          src_path: The full path (including 'gs://') to the data on
+            Google Cloud storage.
+          tmp_path: The full path to the local temporary directory
+            (on a local filesystem).
         """
-        Gets the image stack of the given timestep, if return_target is true, only returns the mask and satellite channels
-        as the model does not need to predict the time, topogrpahic data, etc.
-        :param sample:
-        :param idx:
-        :param return_target:
-        :return:
-        """
-        target = load_np(sample[f"{self.target_type}.{idx:03d}.npy"])
-        if "mask" in self.target_type:
-            target = binarize_mask(target)  # Not actual target, but for now, should be good
-
-        if return_target and not return_image:
-            return None, target
-        img_cube = np.empty(
-            shape=(target.shape[0], target.shape[1], self.total_per_timestep_channels)
-        )
-        bands = np.array([load_np(sample[f"{b.lower()}.{idx:03d}.npy"]) for b in self.bands])
-        img_cube[:, :, : self.num_bands] = bands.transpose((1, 2, 0))
-
-        if self.use_time and not self.time_aux:
-            t = create_time_layer(
-                pickle.loads(sample["time.pyd"])[idx],
-                shape=(img_cube.shape[0], img_cube.shape[1]),
-            )
-            img_cube[:, :, self.num_bands : self.num_bands + 3] = t
-        if self.use_mask:
-            mask = binarize_mask(load_np(sample[f"cloudmask.{idx:03d}.npy"]))
-            # Resample between -1 and 1 like the rest
-            mask = (mask * 2) - 1.0
-            img_cube[:, :, -1] = mask
-        return img_cube, target
-
-    def __iter__(self):
-        # Need to make sure same time step for all of them.
-        # As its all from rapid scan, should be fairly easy.
-        # Main missing one is the regional and rapid weather ones, which are every 15 minutes,
-        # but could be interpolated between the previous step and next one by weighting by time difference
-        # Topographic is same of course, just need to resize to 1km x 1km?
-        # grid by taking the mean value of the interior ones
-        while True:
-            sources = [iter(ds) for ds in self.datasets]
-            if not self.train:  # Same for validation each time for each source
-                np.random.seed(42)
-                # Have to set Python random seed for Albumentations
-                random.seed(a=42)
-                # Reduce number of repeat examples from each source, gives more variety in the inputs
-                self.num_crops = 5
-                self.num_times = 4
-            for source in sources:
-                try:
-                    sample = next(source)
-                except StopIteration:
-                    continue
-                timesteps = pickle.loads(sample["time.pyd"])
-                logger.debug(f"Timesteps: {timesteps}")
-                available_steps = len(timesteps)  # number of available timesteps
-                logger.debug(f"Available Timesteps: {available_steps}")
-                # Check to make sure all timesteps exist
-                sample_keys = [key for key in sample.keys() if self.bands[0].lower() in key]
-                key_checker = [
-                    f"{self.bands[0].lower()}.{idx:03d}.npy" for idx in range(1, available_steps)
-                ]
-                if (
-                    not all(e in sample_keys for e in key_checker)
-                    or self.num_timesteps * self.skip_timesteps + 1
-                    >= available_steps - self.forecast_times
-                ):
-                    logger.debug(
-                        f"Issue with {self.num_timesteps * self.skip_timesteps + 1} >= {available_steps - self.forecast_times} with {available_steps} available timesteps"
-                    )
-                    continue  # Skip this sample as it is missing timesteps, or has none
-                # Times that have enough previous timesteps and post timesteps for training
-                # pick one at random
-                idxs = np.random.randint(
-                    self.num_timesteps * self.skip_timesteps + 1,
-                    available_steps - self.forecast_times,
-                    size=self.num_times,
-                )
-                self.get_topo_latlon(sample)
-                for idx in idxs:
-                    target_timesteps = np.full(self.num_crops, idx + self.forecast_times)
-                    for _ in range(self.num_crops):  # Do random crops as well for training
-                        for target_timestep in target_timesteps:
-                            input_idxs = list(
-                                range(
-                                    idx - (self.num_timesteps * self.skip_timesteps),
-                                    idx + self.skip_timesteps,
-                                    self.skip_timesteps,
-                                )
-                            )
-                            logger.debug(f"Input IDXs: {input_idxs}")
-                            if not self.create_stack(
-                                input_idxs, sample, is_input=self.image_input
-                            ):
-                                self.replay = None
-                                continue
-                            # Now in a Time x W x H x Channel order
-                            target_idxs = list(range(idx + 1, target_timestep + 1))
-                            logger.debug(f"Target IDXs: {target_idxs}")
-                            logger.debug(
-                                f"Timesteps: Current: {timesteps[input_idxs[-1]]} Prev: {timesteps[input_idxs[0]]} Next: {timesteps[target_idxs[0]]} Final: {timesteps[target_idxs[-1]]} "
-                                f"Timedelta: Next - Curr: {timesteps[target_idxs[0]] - timesteps[input_idxs[-1]] } End - Curr: {timesteps[target_idxs[-1]] - timesteps[input_idxs[-1]]}"
-                            )
-                            if not self.create_stack(target_idxs, sample, is_input=False):
-                                self.replay = None
-                                continue
-                            logger.debug(f"Target Masks Shape: {self.target_cube.shape}")
-                            logger.debug(
-                                f"After Time Changes Image/Masks Shape: {self.input_cube.shape} {self.target_cube.shape}"
-                            )
-                            self.input_cube[:, : self.num_bands, :, :] = (
-                                metnet_normalization(self.input_cube[:, : self.num_bands, :, :])
-                                if self.metnet_norm
-                                else standard_normalization(
-                                    self.input_cube[:, : self.num_bands, :, :],
-                                    std=self.std,
-                                    mean=self.mean,
-                                )
-                            )
-                            if self.use_image:
-                                self.target_image_cube[:, : self.num_bands, :, :] = (
-                                    metnet_normalization(
-                                        self.target_image_cube[:, : self.num_bands, :, :]
-                                    )
-                                    if self.metnet_norm
-                                    else standard_normalization(
-                                        self.target_image_cube[:, : self.num_bands, :, :],
-                                        std=self.std,
-                                        mean=self.mean,
-                                    )
-                                )
-                            start_channel = self.num_bands + 3 if self.use_time else self.num_bands
-                            start_channel = start_channel + 1 if self.use_mask else start_channel
-                            logger.info(f"Start Channel: {start_channel}")
-                            self.add_aux_layers(start_channel=start_channel)
-                            logger.debug(
-                                f"After Aux Layers Image/Masks Shape: {self.input_cube.shape} {self.target_cube.shape}"
-                            )
-                            # Reset the replay
-                            self.replay = None
-                            if self.output_target != self.output_shape:
-                                if self.use_image:
-                                    target_image = crop_center(
-                                        self.target_image_cube,
-                                        self.output_target,
-                                        self.output_target,
-                                    )
-                                target_mask = crop_center(
-                                    self.target_cube, self.output_target, self.output_target
-                                )
-                            else:
-                                if self.use_image:
-                                    target_image = self.target_image_cube
-                                target_mask = self.target_cube
-                            # Now convert to channels if time_as_channels
-                            if self.time_as_channels:
-                                image, target_mask, target_image = self.time_changes(
-                                    self.input_cube if self.image_input else self.input_mask_cube,
-                                    self.target_cube,
-                                    self.target_image_cube,
-                                )
-                            else:
-                                image = (
-                                    self.input_cube if self.image_input else self.input_mask_cube
-                                )
-                            # Need to make sure no NaNs and change dtype
-                            image = np.nan_to_num(
-                                image, copy=False, neginf=0.0, posinf=0.0
-                            ).astype(np.float32)
-                            target_mask = np.nan_to_num(
-                                target_mask, copy=False, neginf=0.0, posinf=0.0
-                            ).astype(np.float32)
-                            if self.use_image:
-                                target_image = np.nan_to_num(
-                                    target_image, copy=False, neginf=0.0, posinf=0.0
-                                ).astype(np.float32)
-                            if self.vis:
-                                self.visualize(image, target_image=target_image, mask=target_mask)
-                            if self.use_time and self.time_aux:
-                                time_layer = create_time_layer(
-                                    target_timestep - idx - 1, self.output_shape
-                                )
-                                yield image, time_layer, target_image, target_mask
-                            if not self.use_image:
-                                yield image, target_mask
-                            else:
-                                yield image, target_image
-
-    def get_topo_latlon(self, sample: dict) -> None:
-        if self.use_topo:
-            topo = load_np(sample["topo.npy"])
-            topo[topo < 100] = 0  # Elevation shouldn't really be below 0 here (ocean mostly)
-            self.topo = (topo - TOPO_MEAN) / TOPO_STD
-            self.topo = np.expand_dims(self.topo, axis=0)
-        if self.use_latlon:
-            self.location = load_np(sample["location.npy"])
-            self.location = np.moveaxis(self.location, [2], [0])
-
-    def create_stack(self, idxs: list, sample: dict, is_input: bool = False) -> bool:
-        # Now create stack here
-        time_idx = 0
-        for i in idxs:
-            t_image, t_mask = self.get_timestep(
-                sample,
-                i,
-                return_target=True,
-                return_image=self.use_image or is_input,
-            )
-            if self.replay is None:
-                data = self.aug(image=t_mask)
-                self.replay = data["replay"]
-                logger.debug(self.replay)
-            t_mask = self.aug.replay(self.replay, image=t_mask)["image"]
-            if t_image is not None:
-                t_image = self.aug.replay(self.replay, image=t_image)["image"]
-            if is_input:
-                self.input_mask_cube[time_idx, :, :, :] = t_mask
-            else:
-                self.target_cube[time_idx, :, :, :] = t_mask
-            if is_input:
-                self.input_cube[time_idx, : t_image.shape[2], :, :] = t_image.transpose((2, 1, 0))
-            elif t_image is not None:
-                remove_last_channels = 3 if self.use_time else 0
-                remove_last_channels = (
-                    remove_last_channels + 1 if self.use_mask else remove_last_channels
-                )
-                t_image = (
-                    t_image[:, :, :-remove_last_channels] if remove_last_channels > 0 else t_image
-                )
-                self.target_image_cube[time_idx, : t_image.shape[2], :, :] = t_image.transpose(
-                    (2, 1, 0)
-                )
-
-            time_idx += 1
-        # Convert to Time x Channel x W x H
-        # target_mask = np.expand_dims(target_mask, axis=1)
-        # One timestep as well
-        # Ensure last target mask is also different than previous ones -> only want ones where things change
-        if np.allclose(self.target_cube[0], self.target_cube[-1]) and not is_input:
-            return False
-        return True
-
-    def add_aux_layers(self, start_channel: int) -> None:
-        if self.use_topo:
-            self.apply_aug_to_time(self.topo, start_channel=start_channel)
-            start_channel += 1  # Topo is single channel
-        if self.use_latlon:
-            self.apply_aug_to_time(self.location, start_channel=start_channel)
-            start_channel += 3  # Location is triple channel
-        if self.add_pixel_coords:
-            self.input_cube[
-                :, start_channel : start_channel + self.pixel_coords.shape[0], :, :
-            ] = self.pixel_coords
-
-    def apply_aug_to_time(self, data: np.ndarray, start_channel: int = 1) -> None:
-        data = np.moveaxis(
-            data,
-            [0],
-            [2],
+        super().__init__(
+            n_batches,
+            src_path,
+            tmp_path,
+            cloud,
+            required_keys,
+            history_minutes,
+            forecast_minutes,
+            current_timestep_index,
         )
-        data = self.aug.replay(self.replay, image=data)["image"]
-        data = np.moveaxis(
-            data,
-            [2],
-            [0],
-        )
-        self.input_cube[:, start_channel : start_channel + data.shape[0], :, :] = data
-
-    def time_changes(
-        self, inputs: np.ndarray, target: np.ndarray, target_image: np.ndarray
-    ) -> Union[Tuple[np.ndarray, np.ndarray, np.ndarray], Tuple[np.ndarray, np.ndarray, None]]:
-        # Time changes has to be different make new array and copy into that one
-        if self.image_input and self.num_channels - self.total_per_timestep_channels != 0:
-            time_flattened_input = np.empty(
-                (
-                    self.total_per_timestep_channels * inputs.shape[0]
-                    + (self.num_channels - self.total_per_timestep_channels),
-                    inputs.shape[2],
-                    inputs.shape[3],
-                )
-            )
-            time_flattened_input[
-                : self.total_per_timestep_channels * inputs.shape[0], :, :
-            ] = inputs[:, : self.total_per_timestep_channels, :, :].reshape(
-                -1, inputs.shape[2], inputs.shape[3]
-            )
-            time_flattened_input[
-                -(self.num_channels - self.total_per_timestep_channels) :, :, :
-            ] = inputs[0, self.total_per_timestep_channels :, :, :]
-            inputs = time_flattened_input
-        else:
-            inputs = inputs.reshape((-1, inputs.shape[2], inputs.shape[3]))
-        targets = target.reshape((-1, target.shape[2], target.shape[3]))
-        target_image = target_image.reshape((-1, target_image.shape[2], target_image.shape[3]))
-        return inputs, targets, target_image
-
-
-class CloudFlowDataset(SatFlowDataset):
-    def __init__(self, datasets: List[wds.WebDataset], config: dict, train: bool = True):
-        super(CloudFlowDataset, self).__init__(datasets, config, train)
-        self.image_input = False  # Only want Mask -> Mask training
-
-
-class PerceiverDataset(SatFlowDataset):
-    def __init__(self, datasets: List[wds.WebDataset], config: dict, train: bool = True):
-        super(PerceiverDataset, self).__init__(datasets, config, train)
-        self.image_input = True
-
-    def add_position_encoding(self):
-        pass
-
-    def encode_images(self):
-        pass
-
-
-class OpticalFlowDataset(SatFlowDataset):
-    def __iter__(self) -> Iterator[T_co]:
-        # Need to make sure same time step for all of them.
-        # As its all from rapid scan, should be fairly easy.
-        # Main missing one is the regional and rapid weather ones, which are every 15 minutes,
-        # but could be interpolated between the previous step and next one by weighting by time difference
-        # Topographic is same of course, just need to resize to 1km x 1km?
-        # grid by taking the mean value of the interior ones
-        sources = [iter(ds) for ds in self.datasets]
-        while True:
-            for source in sources:
-                sample = next(source)
-                timesteps = pickle.loads(sample["time.pyd"])
-                available_steps = len(timesteps)  # number of available timesteps
-                # Check to make sure all timesteps exist
-                sample_keys = [key for key in sample.keys() if self.bands[0].lower() in key]
-                key_checker = [
-                    f"{self.bands[0].lower()}.{idx:03d}.npy" for idx in range(1, available_steps)
-                ]
-                if (
-                    not all(e in sample_keys for e in key_checker)
-                    or len(sample_keys)
-                    <= self.num_timesteps * self.skip_timesteps + self.forecast_times
-                ):
-                    continue  # Skip this sample as it is missing timesteps, or has none
-                idxs = list(range(2, available_steps - self.forecast_times))
-                for idx in idxs:
-                    for _ in range(self.num_crops):  # Do random crops as well for training
-                        logger.debug(f"IDX: {idx}")
-                        logger.debug(
-                            f"Timesteps: Current: {timesteps[idx]} Prev: {timesteps[idx - 1]} Next: {timesteps[idx + 1]} Final: {timesteps[idx + self.forecast_times - 1]} "
-                            f"Timedelta: Next - Curr: {timesteps[idx + 1] - timesteps[idx] } End - Curr: {timesteps[idx + self.forecast_times - 1] - timesteps[idx]}"
-                        )
-                        image, mask = self.get_timestep(
-                            sample,
-                            idx,
-                            return_target=True,
-                            return_image=True,
-                        )  # First timestep considered
-                        data = self.aug(image=mask)
-                        replay = data["replay"]
-                        mask = data["image"]
-                        image = self.aug.replay(replay, image=image)["image"]
-                        prev_image, prev_mask = self.get_timestep(
-                            sample,
-                            idx - 1,
-                            return_target=True,
-                            return_image=True,
-                        )  # First timestep considered
-                        prev_mask = self.aug.replay(replay, image=prev_mask)["image"]
-                        prev_image = self.aug.replay(replay, image=prev_image)["image"]
-                        # Now in a Time x W x H x Channel order
-                        _, target_mask = self.get_timestep(
-                            sample,
-                            idx + 1,
-                            return_target=True,
-                            return_image=False,
-                        )
-                        target_mask = self.aug.replay(replay, image=target_mask)["image"]
-                        target_mask = np.expand_dims(target_mask, axis=0)
-
-                        if np.isclose(np.min(target_mask), np.max(target_mask)):
-                            continue  # Ignore if target timestep has no clouds, or only clouds
-                        # Now create stack here
-                        for i in range(idx + 2, idx + self.forecast_times + 1):
-                            _, t_mask = self.get_timestep(
-                                sample,
-                                i,
-                                return_target=True,
-                                return_image=False,
-                            )
-                            t_mask = self.aug.replay(replay, image=t_mask)["image"]
-                            target_mask = np.concatenate(
-                                [target_mask, np.expand_dims(t_mask, axis=0)]
-                            )
-                        target_mask = np.round(target_mask).astype(np.int8)
-                        # Convert to float/half-precision
-                        mask = np.round(mask).astype(np.int8)
-                        prev_mask = np.round(prev_mask).astype(np.int8)
-                        # Move channel to Time x Channel x W x H
-                        mask = np.expand_dims(mask, axis=0)
-                        prev_mask = np.expand_dims(prev_mask, axis=0)
-                        mask = np.nan_to_num(mask, posinf=0.0, neginf=0.0)
-                        prev_mask = np.nan_to_num(prev_mask, posinf=0.0, neginf=0.0)
-                        target_mask = np.nan_to_num(target_mask, posinf=0, neginf=0)
-                        logger.debug(f"Mask: {mask.shape} Target: {target_mask.shape}")
-                        yield prev_mask, mask, target_mask, image, prev_image
-
-
-class FileDataset(thd.Dataset):
-    def __getitem__(self, index) -> T_co:
-        try:
-            if index in self.bad_indexes:
-                index = np.random.randint(0, len(self.files))
-                while index in self.bad_indexes:
-                    index = np.random.randint(0, len(self.files))
-            arrays = np.load(self.files[index])
-            return arrays["images"], arrays["future_images" if self.use_image else "masks"][:24]
-        except:
-            self.bad_indexes.append(index)
-            replacement_index = np.random.randint(0, len(self.files))
-            while replacement_index in self.bad_indexes:
-                replacement_index = np.random.randint(0, len(self.files))
-            arrays = np.load(self.files[replacement_index])
-            return arrays["images"], arrays["future_images" if self.use_image else "masks"][:24]
-
-    def __init__(self, directory: str = "./", train: bool = True, use_image: bool = True):
-        super().__init__()
-        search_space = os.path.join(directory, f"{'train' if train else 'val'}_*.npz")
-        self.files = glob.glob(search_space)
-        self.use_image = use_image
-        self.bad_indexes = []
-
-    def __len__(self):
-        return len(self.files)
-
-
-if __name__ == "__main__":
-    # NIR1.6, VIS0.8 and VIS0.6 RGB for near normal view
-    import numpy as np
-    import webdataset as wds
-
-    dataset = wds.WebDataset("../../datasets/satflow-test.tar")
-    config = load_config("../tests/configs/satflow.yaml")
-    cloudflow = SatFlowDataset([dataset], config)
-    datas = iter(cloudflow)
-    i = 0
-    print("Starting Data")
-    for data in datas:
-        print(f"Data {i}")
-        i += 1
-        if i + 1 > 50:
-            break
-    config = load_config("../tests/configs/satflow_all.yaml")
-    cloudflow = SatFlowDataset([dataset], config)
-    datas = iter(cloudflow)
-    i = 0
-    for data in datas:
-        i += 1
-        if i + 1 > 50:
-            exit()
+        # SatFlow specific changes, i.e. which timestep to split on
+        self.history_steps = history_minutes // 5
+        self.forecast_steps = forecast_minutes // 5
+        self.current_index = (
+            self.history_steps + 1
+        )  # +2 as indexing does not include this index, so need to go one beyond
+        self.required_keys = list(required_keys)
+
+    def __getitem__(self, batch_idx: int):
+        batch = super().__getitem__(batch_idx)
+
+        # Need to partition out past and future sat images here, along with the rest of the data
+        past_satellite_data = batch[SATELLITE_DATA][:, : self.current_index]
+        future_sat_data = batch[SATELLITE_DATA][:, self.current_index :]
+        x = {
+            SATELLITE_DATA: past_satellite_data,
+            SATELLITE_X_COORDS: batch.get(SATELLITE_X_COORDS, None),
+            SATELLITE_Y_COORDS: batch.get(SATELLITE_Y_COORDS, None),
+            SATELLITE_DATETIME_INDEX: batch[SATELLITE_DATETIME_INDEX][:, : self.current_index],
+        }
+        y = {
+            SATELLITE_DATA: future_sat_data,
+            SATELLITE_DATETIME_INDEX: batch[SATELLITE_DATETIME_INDEX][:, self.current_index :],
+        }
+
+        for k in list(DATETIME_FEATURE_NAMES):
+            if k in self.required_keys:
+                x[k] = batch[k][:, : self.current_index]
+
+        if NWP_DATA in self.required_keys:
+            past_nwp_data = batch[NWP_DATA][:, :, : self.current_index]
+            x[NWP_DATA] = past_nwp_data
+            x[NWP_X_COORDS] = batch.get(NWP_X_COORDS, None)
+            x[NWP_Y_COORDS] = batch.get(NWP_Y_COORDS, None)
+
+        return x, y
diff --git a/satflow/data/download.py b/satflow/data/download.py
deleted file mode 100644
index b0c85030..00000000
--- a/satflow/data/download.py
+++ /dev/null
@@ -1,251 +0,0 @@
-import datetime
-import os
-
-from satip import eumetsat
-from satip.eumetsat import compress_downloaded_files
-
-from satflow.data.utils.utils import eumetsat_name_to_datetime, eumetsat_filename_to_datetime
-import logging
-import os
-from datetime import datetime, timedelta
-
-import numpy as np
-import rasterio
-import satpy
-import webdataset as wds
-from pyresample import load_area
-
-logger = logging.getLogger("satflow")
-logger.setLevel(logging.DEBUG)
-
-#
-areas = load_area("/home/bieker/Development/satflow/satflow/resources/areas.yaml")
-topo_data = np.load("../resources/cutdown_europe_dem.npy")
-topo_data[topo_data < -100] = 0  # Ocean
-location_array = np.load("../resources/location_array.npy")
-box_1 = (0, 446, 0, 978)
-box_2 = (446, -1, 978, -1)
-
-data_dir = "/run/media/bieker/Round1/EUMETSAT/"
-debug_fp = "../logs/EUMETSAT_download.txt"
-metadata_db_fp = "../data/EUMETSAT_metadata.db"
-from glob import glob
-
-grb_files = glob(data_dir + "*.grb")
-
-dm = eumetsat.DownloadManager(user_key, user_secret, data_dir, metadata_db_fp, debug_fp)
-
-start_date = "2020-01-01 08:00"
-end_date = "2020-01-01 20:00"
-from satflow.data.utils.utils import map_satellite_to_mercator
-
-
-def make_day(data, flow=True, batch=146, tile=True):
-    root_dir, sat_date, shard_num = data
-    shard_num += 150  # Start after the current ones
-    # reset optical flow samples
-    flow_sample = {}
-    # Initialize new flow sample
-    overall_datetime = sat_date
-    shard_num = overall_datetime.strftime("%Y%m%d")
-    batch_num = 0
-    if batch < 1:
-        flow_sample["__key__"] = overall_datetime.strftime("%Y/%m/%d")
-    else:
-        flow_sample["__key__"] = overall_datetime.strftime("%Y/%m/%d") + f"{batch_num}"
-    # Split data into 8 chunks? 446x489 squareish areas
-    samples = [
-        {"__key__": overall_datetime.strftime("%Y/%m/%d") + f"_{i}", "time.pyd": []}
-        for i in range(8)
-    ]
-    flow_sample["time.pyd"] = []
-    if flow:
-        flow_sample["topo.npy"] = topo_data
-        flow_sample["location.npy"] = location_array
-    print(f"On shard: {shard_num} Date: {overall_datetime.strftime('%Y/%m/%d')}")
-    interday_frame = 0
-    if os.path.exists(
-        f"/run/media/bieker/data/EUMETSAT/satflow{'-' if not flow else '-flow'}{'-' if not flow and batch > 0 else f'-{batch}-'}{'tiled-' if tile else ''}{shard_num}.tar"
-    ):
-        return
-    sink_flow = wds.TarWriter(
-        f"/run/media/bieker/data/EUMETSAT/satflow{'-' if not flow else '-flow'}{'-' if not flow and batch > 0 else f'-{batch}-'}{'tiled-' if tile else ''}{shard_num}.tar",
-        compress=True,
-    )
-    for root, dirs, files in os.walk(root_dir):
-        dirs.sort(
-            key=int
-        )  # Ensure this is done in sorted order, so its always from earliest to latest
-        # As its stored in year, month, day, hour, minute, format, once we are at the day level, create a new shard
-        # Also cut down everything to only the square of all value data
-        if files:
-            # At lowest one, so have everything in it
-            if len(files) < 11:  # Not everything worked
-                continue
-
-            date_data = root.split("/")
-            date_str = os.path.join(
-                date_data[-5], date_data[-4], date_data[-3], date_data[-2], date_data[-1]
-            )
-            datetime_object = datetime.strptime(date_str, "%Y/%m/%d/%H/%M")
-            if not flow:
-                flow_sample = {}
-                flow_sample["__key__"] = datetime_object.strftime("%Y/%m/%d/%H/%M")
-
-            # try:
-            # Want this at the top, as if this extraction fails, we still then know if there is a missing frame
-            interday_frame += 1
-            for f in files:
-                if ".tif" in f:
-                    # Only need geotiff ones
-                    src = rasterio.open(os.path.join(root, f))
-                    data = src.read(1)
-                    cropped_data = data[
-                        515:-641, 603:
-                    ]  # Done by handish to exclude all NODATA and invalid masks for clouds and images (clouds have a smaller extant)
-                    channel = f.split("_")
-                    cropped_data = cropped_data.astype(np.float16)
-                    if channel[0] in ["IR", "WV", "cloud"]:
-                        channel = channel[0] + channel[1]  # These are split across multiple
-                    else:
-                        channel = channel[0]
-                    if tile:
-                        height = (0, 446, -1)
-                        width = (0, 490, 980, 1470, -1)
-                        s_num = 0
-                        for j in range(4):
-                            for i in range(2):
-                                tile_cropped_data = cropped_data[
-                                    height[i] : height[i + 1], width[j] : width[j + 1]
-                                ]
-                                tiled_topo = topo_data[
-                                    height[i] : height[i + 1], width[j] : width[j + 1]
-                                ]
-                                tiled_loc = location_array[
-                                    height[i] : height[i + 1], width[j] : width[j + 1]
-                                ]
-                                samples[s_num][
-                                    channel + f".{interday_frame:03d}.npy"
-                                ] = tile_cropped_data.astype(np.float16)
-                                samples[s_num]["location.npy"] = tiled_loc
-                                samples[s_num]["topo.npy"] = tiled_topo
-                                s_num += 1
-
-                    elif not flow:
-                        flow_sample[channel + ".npy"] = cropped_data
-                    else:
-                        flow_sample[channel + f".{interday_frame:03d}.npy"] = cropped_data.astype(
-                            np.float16
-                        )
-            # Now all channels added to thing, write to shard
-            if flow:
-                if not tile:
-                    flow_sample["time.pyd"].append(datetime_object)
-                for i, s in enumerate(samples):
-                    samples[i]["time.pyd"].append(datetime_object)
-                if batch > 0:
-                    print(f"In Batch: {len(samples[0]['time.pyd'])} == {batch} Shard: {shard_num}")
-                    if len(samples[0]["time.pyd"]) == batch:
-                        if tile:
-                            for s in samples:
-                                sink_flow.write(s)
-                        else:
-                            sink_flow.write(flow_sample)
-                            flow_sample["__key__"] = datetime_object.strftime("%Y/%m/%d/%H/%M")
-                            flow_sample["topo.npy"] = topo_data
-                            flow_sample["location.npy"] = location_array
-                            flow_sample["time.pyd"] = []
-                        interday_frame = 0
-                        flow_sample = {}
-                        samples = [
-                            {
-                                "__key__": overall_datetime.strftime("%Y/%m/%d") + f"_{i}",
-                                "time.pyd": [],
-                            }
-                            for i in range(8)
-                        ]
-                        batch_num += 1
-            else:
-                flow_sample["time.pyd"] = datetime_object
-                sink_flow.write(flow_sample)
-    if tile:
-        for s in samples:
-            sink_flow.write(s)
-    sink_flow.close()
-
-
-import shutil
-
-for day in range(15, 28):
-    for month in [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]:
-        for year in [2020]:
-            curr_date = f"{year}/{month}/{day}"
-            curr_date = datetime.strptime(curr_date, "%Y/%m/%d")
-            # make_day((os.path.join(data_dir, curr_date.strftime("%Y/%m/%d/")), curr_date, 0))
-            # continue
-            dm.download_date_range(
-                f"{year}-{month}-{day} 07:59",
-                f"{year}-{month}-{day} 20:05",
-                product_id="EO:EUM:DAT:MSG:RSS-CLM",
-            )
-            dm.download_date_range(
-                f"{year}-{month}-{day} 07:59",
-                f"{year}-{month}-{day} 20:05",
-                product_id="EO:EUM:DAT:MSG:RII",
-            )
-            dm.download_date_range(
-                f"{year}-{month}-{day} 07:59",
-                f"{year}-{month}-{day} 20:05",
-                product_id="EO:EUM:DAT:MSG:RSS-AMV",
-            )
-            dm.download_date_range(
-                f"{year}-{month}-{day} 07:59",
-                f"{year}-{month}-{day} 20:05",
-                product_id="EO:EUM:DAT:MSG:RSS-MPE-GRIB",
-            )
-            dm.download_date_range(f"{year}-{month}-{day} 07:59", f"{year}-{month}-{day} 20:05")
-            # Now go through and create the webdataset for the day
-            nat_files = sorted(glob(data_dir + "*.nat"))
-
-            for n in nat_files:
-                base_native_filename = os.path.basename(n)
-                dt = eumetsat_filename_to_datetime(base_native_filename)
-                if dt.date() != curr_date.date():
-                    # New date, so theoretically all done for the day
-                    continue
-                grb_files = sorted(glob(data_dir + "*.grb"))
-                for grb in grb_files:
-                    base_grb_filename = os.path.basename(grb)
-                    grb_dt = eumetsat_name_to_datetime(base_grb_filename) - timedelta(minutes=1)
-                    if (
-                        grb_dt.date() == dt.date()
-                        and grb_dt.date() == curr_date.date()
-                        and grb_dt.hour == dt.hour
-                        and grb_dt.minute == dt.minute
-                    ):
-                        # Have the data to create webdataset
-                        map_satellite_to_mercator(
-                            native_satellite=n,
-                            grib_files=grb,
-                            save_loc=os.path.join(data_dir, dt.strftime("%Y/%m/%d/%H/%M")),
-                            bands=(
-                                "HRV",
-                                "IR_016",
-                                "IR_039",
-                                "IR_087",
-                                "IR_097",
-                                "IR_108",
-                                "IR_120",
-                                "IR_134",
-                                "VIS006",
-                                "VIS008",
-                                "WV_062",
-                                "WV_073",
-                                "cloud_mask",
-                            ),
-                        )
-                        os.remove(grb)
-                        os.remove(n)
-            make_day((os.path.join(data_dir, curr_date.strftime("%Y/%m/%d/")), curr_date, 0))
-            # Clean up directory that isn't needed anymore
-            # shutil.rmtree(os.path.join(data_dir, dt.strftime("%Y/%m/%d/")))
diff --git a/satflow/data/preprocess.py b/satflow/data/preprocess.py
deleted file mode 100644
index 6ef75cd7..00000000
--- a/satflow/data/preprocess.py
+++ /dev/null
@@ -1,51 +0,0 @@
-import torch
-
-from satflow.data.datasets import SatFlowDataset
-import numpy as np
-from torch.utils.data import DataLoader
-from satflow.data.datasets import SatFlowDataset, CloudFlowDataset
-import webdataset as wds
-import yaml
-
-
-def load_config(config_file):
-    with open(config_file, "r") as cfg:
-        return yaml.load(cfg, Loader=yaml.FullLoader)["config"]
-
-
-config = load_config("preprocess.yaml")
-
-train = wds.WebDataset("/run/media/jacob/data/satflow-flow-144-tiled-{00001..00105}.tar")
-val = wds.WebDataset("/run/media/jacob/data/satflow-flow-144-tiled-{00106..00149}.tar")
-
-t_flow = SatFlowDataset([train], config, train=True)
-v_flow = SatFlowDataset([val], config, train=False)
-
-training_dataloader = DataLoader(
-    t_flow,
-    batch_size=1,
-    pin_memory=True,
-    num_workers=16,
-)
-val_dataloader = DataLoader(
-    v_flow,
-    batch_size=1,
-    pin_memory=True,
-    num_workers=8,
-)
-
-output_dir = "/run/media/jacob/data/satflow_prev6_skip3_notime/"
-
-for i, data in enumerate(training_dataloader):
-    if i > 40000:
-        break
-for i, data in enumerate(val_dataloader):
-    image, target_image, target_mask = data
-    image = image[0].numpy()
-    target_image = target_image[0].numpy()
-    target_mask = target_mask[0].numpy()
-    np.savez_compressed(
-        output_dir + f"val_{i}.npz", images=image, future_images=target_image, masks=target_mask
-    )
-    if i > 10000:
-        break
diff --git a/satflow/data/preprocess.yaml b/satflow/data/preprocess.yaml
deleted file mode 100644
index b9a09410..00000000
--- a/satflow/data/preprocess.yaml
+++ /dev/null
@@ -1,34 +0,0 @@
-config:
-  visualize: False
-  num_timesteps: 6
-  skip_timesteps: 3
-  forecast_times: 48
-  output_shape: 128
-  target_type: "cloudmask"
-  num_crops: 6
-  use_topo: True
-  use_latlon: False
-  use_time: True
-  time_aux: False
-  use_mask: False
-  use_image: True
-  num_times: 6
-  add_pixel_coords: False
-  add_polar_coords: False
-  time_as_channels: False
-  bands:
-    [
-      "HRV",
-      "IR016",
-      "IR039",
-      "IR087",
-      "IR097",
-      "IR108",
-      "IR120",
-      "IR134",
-      "VIS006",
-      "VIS008",
-      "WV062",
-      "WV073",
-    ]
-  transforms: {}
diff --git a/satflow/data/utils/nimrod.py b/satflow/data/utils/nimrod.py
deleted file mode 100644
index de55eced..00000000
--- a/satflow/data/utils/nimrod.py
+++ /dev/null
@@ -1,438 +0,0 @@
-#!/usr/bin/python
-"""
-Extract data from UK Met Office Rain Radar NIMROD image files.
-
-Parse NIMROD format image files, display header data and allow extraction of
-raster image to an ESRI ASCII (.asc) format file. A bounding box may be
-specified to clip the image to the area of interest. Can be imported as a
-Python module or run directly as a command line script.
-
-Author: Richard Thomas
-Version: 1.0 (13 April 2015)
-Public Repository: https://github.com/richard-thomas/MetOffice_NIMROD
-
-Command line usage:
-  python nimrod.py [-h] [-q] [-x] [-bbox XMIN XMAX YMIN YMAX] [infile] [outfile]
-
-positional arguments:
-  infile                (Uncompressed) NIMROD input filename
-  outfile               Output raster filename (*.asc)
-
-optional arguments:
-  -h, --help            show this help message and exit
-  -q, --query           Display metadata
-  -x, --extract         Extract raster file in ASC format
-  -bbox XMIN XMAX YMIN YMAX
-                        Bounding box to clip raster data to
-
-Note that any bounding box must be specified in the same units and projection
-as the input file. The bounding box does not need to be contained by the input
-raster but must intersect it.
-
-Example command line usage:
-  python nimrod.py -bbox 279906 285444 283130 290440
-    -xq 200802252000_nimrod_ng_radar_rainrate_composite_1km_merged_UK_zip
-    plynlimon_catchments_rainfall.asc
-
-Example Python module usage:
-    import nimrod
-    a = nimrod.Nimrod(open(
-        '200802252000_nimrod_ng_radar_rainrate_composite_1km_merged_UK_zip'))
-    a.query()
-    a.extract_asc(open('full_raster.asc', 'w'))
-    a.apply_bbox(279906, 285444, 283130, 290440)
-    a.query()
-    a.extract_asc(open('clipped_raster.asc', 'w'))
-
-Notes:
-  1. Valid for v1.7 and v2.6-4 of NIMROD file specification
-  2. Assumes image origin is top left (i.e. that header[24] = 0)
-  3. Tested on UK composite 1km and 5km data, under Linux and Windows XP
-  4. Further details of NIMROD data and software at the NERC BADC website:
-      http://badc.nerc.ac.uk/browse/badc/ukmo-nimrod/
-
-Copyright (c) 2015 Richard Thomas
-(Nimrod.__init__() method based on read_nimrod.py by Charles Kilburn Aug 2008)
-
-This program is free software: you can redistribute it and/or modify
-it under the terms of the Artistic License 2.0 as published by the
-Open Source Initiative (http://opensource.org/licenses/Artistic-2.0)
-
-This program is distributed in the hope that it will be useful,
-but WITHOUT ANY WARRANTY; without even the implied warranty of
-MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
-"""
-
-import array
-import struct
-import sys
-
-
-class Nimrod:
-    """Reading, querying and processing of NIMROD format rainfall data files."""
-
-    class RecordLenError(Exception):
-        """
-        Exception Type: NIMROD record length read from file not as expected.
-        """
-
-        def __init__(self, actual, expected, location):
-            self.message = "Incorrect record length %d bytes (expected %d) at %s." % (
-                actual,
-                expected,
-                location,
-            )
-
-    class HeaderReadError(Exception):
-        """Exception Type: Read error whilst parsing NIMROD header elements."""
-
-        pass
-
-    class PayloadReadError(Exception):
-        """Exception Type: Read error whilst parsing NIMROD raster data."""
-
-        pass
-
-    class BboxRangeError(Exception):
-        """
-        Exception Type: Bounding box specified out of range of raster image.
-        """
-
-        pass
-
-    def __init__(self, infile):
-        """
-        Parse all header and data info from a NIMROD data file into this object.
-        (This method based on read_nimrod.py by Charles Kilburn Aug 2008)
-
-        Args:
-            infile: NIMROD file object opened for binary reading
-        Raises:
-            RecordLenError: NIMROD record length read from file not as expected
-            HeaderReadError: Read error whilst parsing NIMROD header elements
-            PayloadReadError: Read error whilst parsing NIMROD raster data
-        """
-
-        def check_record_len(infile, expected, location):
-            """
-            Check record length in C struct is as expected.
-
-            Args:
-                infile: file to read from
-                expected: expected value of record length read
-                location: description of position in file (for reporting)
-            Raises:
-                HeaderReadError: Read error whilst reading record length
-                RecordLenError: Unexpected NIMROD record length read from file
-            """
-
-            # Unpack length from C struct (Big Endian, 4-byte long)
-            try:
-                (record_length,) = struct.unpack(">l", infile.read(4))
-            except Exception:
-                raise Nimrod.HeaderReadError
-            if record_length != expected:
-                raise Nimrod.RecordLenError(record_length, expected, location)
-
-        # Header should always be a fixed length record
-        check_record_len(infile, 512, "header start")
-
-        try:
-            # Read first 31 2-byte integers (header fields 1-31)
-            gen_ints = array.array("h")
-            gen_ints.fromfile(infile, 31)
-            gen_ints.byteswap()
-
-            # Read next 28 4-byte floats (header fields 32-59)
-            gen_reals = array.array("f")
-            gen_reals.fromfile(infile, 28)
-            gen_reals.byteswap()
-
-            # Read next 45 4-byte floats (header fields 60-104)
-            spec_reals = array.array("f")
-            spec_reals.fromfile(infile, 45)
-            spec_reals.byteswap()
-
-            # Read next 56 characters (header fields 105-107)
-            characters = array.array("b")
-            characters.fromfile(infile, 56)
-
-            # Read next 51 2-byte integers (header fields 108-)
-            spec_ints = array.array("h")
-            spec_ints.fromfile(infile, 51)
-            spec_ints.byteswap()
-        except Exception:
-            infile.close()
-            raise Nimrod.HeaderReadError
-
-        check_record_len(infile, 512, "header end")
-
-        # Extract strings and make duplicate entries to give meaningful names
-        chars = characters.tobytes()
-        self.units = chars[0:8]
-        self.data_source = chars[8:32]
-        self.title = chars[32:55]
-
-        # Store header values in a list so they can be indexed by "element
-        # number" shown in NIMROD specification (starts at 1)
-        self.hdr_element = [None]  # Dummy value at element 0
-        self.hdr_element.extend(gen_ints)
-        self.hdr_element.extend(gen_reals)
-        self.hdr_element.extend(spec_reals)
-        self.hdr_element.extend([self.units])
-        self.hdr_element.extend([self.data_source])
-        self.hdr_element.extend([self.title])
-        self.hdr_element.extend(spec_ints)
-
-        # Duplicate some of values to give more meaningful names
-        self.nrows = self.hdr_element[16]
-        self.ncols = self.hdr_element[17]
-        self.n_data_specific_reals = self.hdr_element[22]
-        self.n_data_specific_ints = self.hdr_element[23] + 1
-        # Note "+ 1" because header value is count from element 109nimrod_to_kmz
-        self.y_top = self.hdr_element[34]
-        self.y_pixel_size = self.hdr_element[35]
-        self.x_left = self.hdr_element[36]
-        self.x_pixel_size = self.hdr_element[37]
-
-        # Calculate other image bounds (note these are pixel centres)
-        self.x_right = self.x_left + self.x_pixel_size * (self.ncols - 1)
-        self.y_bottom = self.y_top - self.y_pixel_size * (self.nrows - 1)
-
-        # Read payload (actual raster data)
-        array_size = self.ncols * self.nrows
-        check_record_len(infile, array_size * 2, "data start")
-
-        self.data = array.array("h")
-        try:
-            self.data.fromfile(infile, array_size)
-            self.data.byteswap()
-        except Exception:
-            infile.close()
-            raise Nimrod.PayloadReadError
-
-        check_record_len(infile, array_size * 2, "data end")
-        infile.close()
-
-    def query(self):
-        """Print complete NIMROD file header information."""
-
-        print("NIMROD file raw header fields listed by element number:")
-        print("General (Integer) header entries:")
-        for i in range(1, 32):
-            print(" ", i, "\t", self.hdr_element[i])
-        print("General (Real) header entries:")
-        for i in range(32, 60):
-            print(" ", i, "\t", self.hdr_element[i])
-        print("Data Specific (Real) header entries (%d):" % self.n_data_specific_reals)
-        for i in range(60, 60 + self.n_data_specific_reals):
-            print(" ", i, "\t", self.hdr_element[i])
-        print("Data Specific (Integer) header entries (%d):" % self.n_data_specific_ints)
-        for i in range(108, 108 + self.n_data_specific_ints):
-            print(" ", i, "\t", self.hdr_element[i])
-        print("Character header entries:")
-        print("  105 Units:           ", self.units)
-        print("  106 Data source:     ", self.data_source)
-        print("  107 Title of field:  ", self.title)
-
-        # Print out info & header fields
-        # Note that ranges are given to the edge of each pixel
-        print(
-            "\nValidity Time:  %2.2d:%2.2d on %2.2d/%2.2d/%4.4d"
-            % (
-                self.hdr_element[4],
-                self.hdr_element[5],
-                self.hdr_element[3],
-                self.hdr_element[2],
-                self.hdr_element[1],
-            )
-        )
-        print(
-            "Easting range:  %.1f - %.1f (at pixel steps of %.1f)"
-            % (
-                self.x_left - self.x_pixel_size / 2,
-                self.x_right + self.x_pixel_size / 2,
-                self.x_pixel_size,
-            )
-        )
-        print(
-            "Northing range: %.1f - %.1f (at pixel steps of %.1f)"
-            % (
-                self.y_bottom - self.y_pixel_size / 2,
-                self.y_top + self.y_pixel_size / 2,
-                self.y_pixel_size,
-            )
-        )
-        print("Image size: %d rows x %d cols" % (self.nrows, self.ncols))
-
-    def apply_bbox(self, xmin, xmax, ymin, ymax):
-        """
-        Clip raster data to all pixels that intersect specified bounding box.
-
-        Note that existing object data is modified and all header values
-        affected are appropriately adjusted. Because pixels are specified by
-        their centre points, a bounding box that comes within half a pixel
-        width of the raster edge will intersect with the pixel.
-
-        Args:
-            xmin: Most negative easting or longitude of bounding box
-            xmax: Most positive easting or longitude of bounding box
-            ymin: Most negative northing or latitude of bounding box
-            ymax: Most positive northing or latitude of bounding box
-        Raises:
-            BboxRangeError: Bounding box specified out of range of raster image
-        """
-
-        # Check if there is no overlap of bounding box with raster
-        if (
-            xmin > self.x_right + self.x_pixel_size / 2
-            or xmax < self.x_left - self.x_pixel_size / 2
-            or ymin > self.y_top + self.y_pixel_size / 2
-            or ymax < self.y_bottom - self.x_pixel_size / 2
-        ):
-            raise Nimrod.BboxRangeError
-
-        # Limit bounds to within raster image
-        xmin = max(xmin, self.x_left)
-        xmax = min(xmax, self.x_right)
-        ymin = max(ymin, self.y_bottom)
-        ymax = min(ymax, self.y_top)
-
-        # Calculate min and max pixel index in each row and column to use
-        # Note addition of 0.5 as x_left location is centre of pixel
-        # ('int' truncates floats towards zero)
-        xMinPixelId = int((xmin - self.x_left) / self.x_pixel_size + 0.5)
-        xMaxPixelId = int((xmax - self.x_left) / self.x_pixel_size + 0.5)
-
-        # For y (northings), note the first data row stored is most north
-        yMinPixelId = int((self.y_top - ymax) / self.y_pixel_size + 0.5)
-        yMaxPixelId = int((self.y_top - ymin) / self.y_pixel_size + 0.5)
-
-        bbox_data = []
-        for i in range(yMinPixelId, yMaxPixelId + 1):
-            bbox_data.extend(
-                self.data[i * self.ncols + xMinPixelId : i * self.ncols + xMaxPixelId + 1]
-            )
-
-        # Update object where necessary
-        self.data = bbox_data
-        self.x_right = self.x_left + xMaxPixelId * self.x_pixel_size
-        self.x_left += xMinPixelId * self.x_pixel_size
-        self.ncols = xMaxPixelId - xMinPixelId + 1
-        self.y_bottom = self.y_top - yMaxPixelId * self.y_pixel_size
-        self.y_top -= yMinPixelId * self.y_pixel_size
-        self.nrows = yMaxPixelId - yMinPixelId + 1
-        self.hdr_element[16] = self.nrows
-        self.hdr_element[17] = self.ncols
-        self.hdr_element[34] = self.y_top
-        self.hdr_element[36] = self.x_left
-
-    def extract_asc(self, outfile):
-        """
-        Write raster data to an ESRI ASCII (.asc) format file.
-
-        Args:
-            outfile: file object opened for writing text
-        """
-
-        # As ESRI ASCII format only supports square pixels, warn if not so
-        if self.x_pixel_size != self.y_pixel_size:
-            print(
-                "Warning: x_pixel_size(%d) != y_pixel_size(%d)"
-                % (self.x_pixel_size, self.y_pixel_size)
-            )
-
-        # Write header to output file. Note that data is valid at the centre
-        # of each pixel so "xllcenter" rather than "xllcorner" must be used
-        outfile.write("ncols          %d\n" % self.ncols)
-        outfile.write("nrows          %d\n" % self.nrows)
-        outfile.write("xllcenter     %d\n" % self.x_left)
-        outfile.write("yllcenter     %d\n" % self.y_bottom)
-        outfile.write("cellsize       %.1f\n" % self.y_pixel_size)
-        outfile.write("nodata_value  %.1f\n" % self.hdr_element[38])
-
-        # Write raster data to output file
-        for i in range(self.nrows):
-            for j in range(self.ncols - 1):
-                outfile.write("%d " % self.data[i * self.ncols + j])
-            outfile.write("%d\n" % self.data[i * self.ncols + self.ncols - 1])
-        outfile.close()
-
-
-# -------------------------------------------------------------------------------
-# Handle if called as a command line script
-# (And as an example of how to invoke class methods from an importing module)
-# -------------------------------------------------------------------------------
-import argparse
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser(
-        description="Extract information and data from a NIMROD format file",
-        epilog="""Note that any bounding box must be specified in the same
-                  units and projection as the input file. The bounding box
-                  does not need to be contained by the input raster but
-                  must intersect it.""",
-    )
-    parser.add_argument("-q", "--query", action="store_true", help="Display metadata")
-    parser.add_argument(
-        "-x", "--extract", action="store_true", help="Extract raster file in ASC format"
-    )
-    parser.add_argument(
-        "infile",
-        nargs="?",
-        type=argparse.FileType("rb"),
-        default=sys.stdin,
-        help="(Uncompressed) NIMROD input filename",
-    )
-    parser.add_argument(
-        "outfile",
-        nargs="?",
-        type=argparse.FileType("w"),
-        default=sys.stdout,
-        help="Output raster filename (*.asc)",
-    )
-    parser.add_argument(
-        "-bbox",
-        type=float,
-        nargs=4,
-        metavar=("XMIN", "XMAX", "YMIN", "YMAX"),
-        help="Bounding box to clip raster data to",
-    )
-    args = parser.parse_args()
-
-    if not args.query and not args.extract:
-        parser.print_help()
-        sys.exit(1)
-
-    # Initialise data object by reading NIMROD file
-    # (Only trap record length exception as others self-explanatory)
-    try:
-        rainfall_data = Nimrod(args.infile)
-    except Nimrod.RecordLenError as error:
-        sys.stderr.write("ERROR: %s\n" % error.message)
-        sys.exit(1)
-
-    if args.bbox:
-        sys.stderr.write("Trimming NIMROD raster to bounding box...\n")
-        try:
-            rainfall_data.apply_bbox(*args.bbox)
-        except Nimrod.BboxRangeError:
-            sys.stderr.write("ERROR: bounding box not within raster image.\n")
-            sys.exit(1)
-
-    # Perform query after any bounding box trimming to allow sanity checking of
-    # size of resulting image
-    if args.query:
-        rainfall_data.query()
-
-    if args.extract:
-        sys.stderr.write("Extracting NIMROD raster to ASC file...\n")
-        sys.stderr.write(
-            "  Outputting data array (%d rows x %d cols = %d pixels)\n"
-            % (
-                rainfall_data.nrows,
-                rainfall_data.ncols,
-                rainfall_data.nrows * rainfall_data.ncols,
-            )
-        )
-        rainfall_data.extract_asc(args.outfile)
diff --git a/satflow/data/utils/normalization.py b/satflow/data/utils/normalization.py
deleted file mode 100644
index 098c5235..00000000
--- a/satflow/data/utils/normalization.py
+++ /dev/null
@@ -1,25 +0,0 @@
-import numpy as np
-
-
-def metnet_normalization(data: np.ndarray) -> np.ndarray:
-    """
-    Perform the normalization used in the MetNet paper on the GOES data
-    This involves subtracting by the median, dividing by the interquartile range,
-    then squashing to [-1,1] with the hyperbolic tangent
-    Args:
-        data:
-
-    Returns:
-
-    """
-    # Ensure no NaNs
-    data = np.nan_to_num(data)
-    # Get IQR
-    iqr = np.subtract(*np.percentile(data, [75, 25]))
-    data = (data - np.median(data)) / iqr
-    # Now hyperbolic tangent
-    return np.tanh(data)
-
-
-def standard_normalization(data: np.ndarray, std: np.ndarray, mean: np.ndarray) -> np.ndarray:
-    return (data - mean) / std
diff --git a/satflow/data/utils/processors.py b/satflow/data/utils/processors.py
deleted file mode 100644
index c4800e66..00000000
--- a/satflow/data/utils/processors.py
+++ /dev/null
@@ -1,227 +0,0 @@
-# Copyright 2021 DeepMind Technologies Limited
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     https://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Position encodings and utilities."""
-
-import abc
-import functools
-import numpy as np
-
-
-def generate_fourier_features(
-    pos, num_bands, max_resolution=(224, 224), concat_pos=True, sine_only=False
-):
-    """Generate a Fourier frequency position encoding with linear spacing.
-
-    Args:
-      pos: The position of n points in d dimensional space.
-        A np array of shape [n, d].
-      num_bands: The number of bands (K) to use.
-      max_resolution: The maximum resolution (i.e. the number of pixels per dim).
-        A tuple representing resolution for each dimension
-      concat_pos: Concatenate the input position encoding to the Fourier features?
-      sine_only: Whether to use a single phase (sin) or two (sin/cos) for each
-        frequency band.
-    Returns:
-      embedding: A 1D np array of shape [n, n_channels]. If concat_pos is True
-        and sine_only is False, output dimensions are ordered as:
-          [dim_1, dim_2, ..., dim_d,
-           sin(pi*f_1*dim_1), ..., sin(pi*f_K*dim_1), ...,
-           sin(pi*f_1*dim_d), ..., sin(pi*f_K*dim_d),
-           cos(pi*f_1*dim_1), ..., cos(pi*f_K*dim_1), ...,
-           cos(pi*f_1*dim_d), ..., cos(pi*f_K*dim_d)],
-         where dim_i is pos[:, i] and f_k is the kth frequency band.
-    """
-    min_freq = 1.0
-    # Nyquist frequency at the target resolution:
-
-    freq_bands = np.stack(
-        [np.linspace(min_freq, res / 2, num=num_bands, endpoint=True) for res in max_resolution],
-        axis=0,
-    )
-
-    # Get frequency bands for each spatial dimension.
-    # Output is size [n, d * num_bands]
-    per_pos_features = pos[:, :, None] * freq_bands[None, :, :]
-    per_pos_features = np.reshape(per_pos_features, [-1, np.prod(per_pos_features.shape[1:])])
-
-    if sine_only:
-        # Output is size [n, d * num_bands]
-        per_pos_features = np.sin(np.pi * (per_pos_features))
-    else:
-        # Output is size [n, 2 * d * num_bands]
-        per_pos_features = np.concatenate(
-            [np.sin(np.pi * per_pos_features), np.cos(np.pi * per_pos_features)], axis=-1
-        )
-    # Concatenate the raw input positions.
-    if concat_pos:
-        # Adds d bands to the encoding.
-        per_pos_features = np.concatenate([pos, per_pos_features], axis=-1)
-    return per_pos_features
-
-
-def build_linear_positions(index_dims, output_range=(-1.0, 1.0)):
-    """Generate an array of position indices for an N-D input array.
-
-    Args:
-      index_dims: The shape of the index dimensions of the input array.
-      output_range: The min and max values taken by each input index dimension.
-    Returns:
-      A np array of shape [index_dims[0], index_dims[1], .., index_dims[-1], N].
-    """
-
-    def _linspace(n_xels_per_dim):
-        return np.linspace(
-            output_range[0], output_range[1], num=n_xels_per_dim, endpoint=True, dtype=np.float32
-        )
-
-    dim_ranges = [_linspace(n_xels_per_dim) for n_xels_per_dim in index_dims]
-    array_index_grid = np.meshgrid(*dim_ranges, indexing="ij")
-
-    return np.stack(array_index_grid, axis=-1)
-
-
-class AbstractPositionEncoding(hk.Module, metaclass=abc.ABCMeta):
-    """Abstract Perceiver decoder."""
-
-    @abc.abstractmethod
-    def __call__(self, batch_size, pos):
-        raise NotImplementedError
-
-
-class TrainablePositionEncoding(AbstractPositionEncoding):
-    """Trainable position encoding."""
-
-    def __init__(self, index_dim, num_channels=128, init_scale=0.02, name=None):
-        super(TrainablePositionEncoding, self).__init__(name=name)
-        self._index_dim = index_dim
-        self._num_channels = num_channels
-        self._init_scale = init_scale
-
-    def __call__(self, batch_size, pos=None):
-        del pos  # Unused.
-        pos_embs = hk.get_parameter(
-            "pos_embs",
-            [self._index_dim, self._num_channels],
-            init=hk.initializers.TruncatedNormal(stddev=self._init_scale),
-        )
-
-        if batch_size is not None:
-            pos_embs = np.broadcast_to(pos_embs[None, :, :], (batch_size,) + pos_embs.shape)
-        return pos_embs
-
-
-def _check_or_build_spatial_positions(pos, index_dims, batch_size):
-    """Checks or builds spatial position features (x, y, ...).
-
-    Args:
-      pos: None, or an array of position features. If None, position features
-        are built. Otherwise, their size is checked.
-      index_dims: An iterable giving the spatial/index size of the data to be
-        featurized.
-      batch_size: The batch size of the data to be featurized.
-    Returns:
-      An array of position features, of shape [batch_size, prod(index_dims)].
-    """
-    if pos is None:
-        pos = build_linear_positions(index_dims)
-        pos = np.broadcast_to(pos[None], (batch_size,) + pos.shape)
-        pos = np.reshape(pos, [batch_size, np.prod(index_dims), -1])
-    else:
-        # Just a warning label: you probably don't want your spatial features to
-        # have a different spatial layout than your pos coordinate system.
-        # But feel free to override if you think it'll work!
-        assert pos.shape[-1] == len(index_dims)
-
-    return pos
-
-
-class FourierPositionEncoding(AbstractPositionEncoding):
-    """Fourier (Sinusoidal) position encoding."""
-
-    def __init__(
-        self,
-        index_dims,
-        num_bands,
-        concat_pos=True,
-        max_resolution=None,
-        sine_only=False,
-        name=None,
-    ):
-        super(FourierPositionEncoding, self).__init__(name=name)
-        self._num_bands = num_bands
-        self._concat_pos = concat_pos
-        self._sine_only = sine_only
-        self._index_dims = index_dims
-        # Use the index dims as the maximum resolution if it's not provided.
-        self._max_resolution = max_resolution or index_dims
-
-    def __call__(self, batch_size, pos=None):
-        pos = _check_or_build_spatial_positions(pos, self._index_dims, batch_size)
-        build_ff_fn = functools.partial(
-            generate_fourier_features,
-            num_bands=self._num_bands,
-            max_resolution=self._max_resolution,
-            concat_pos=self._concat_pos,
-            sine_only=self._sine_only,
-        )
-        return jax.vmap(build_ff_fn, 0, 0)(pos)
-
-
-class PositionEncodingProjector(AbstractPositionEncoding):
-    """Projects a position encoding to a target size."""
-
-    def __init__(self, output_size, base_position_encoding, name=None):
-        super(PositionEncodingProjector, self).__init__(name=name)
-        self._output_size = output_size
-        self._base_position_encoding = base_position_encoding
-
-    def __call__(self, batch_size, pos=None):
-        base_pos = self._base_position_encoding(batch_size, pos)
-        projected_pos = hk.Linear(output_size=self._output_size)(base_pos)
-        return projected_pos
-
-
-def build_position_encoding(
-    position_encoding_type,
-    index_dims,
-    project_pos_dim=-1,
-    trainable_position_encoding_kwargs=None,
-    fourier_position_encoding_kwargs=None,
-    name=None,
-):
-    """Builds the position encoding."""
-
-    if position_encoding_type == "trainable":
-        assert trainable_position_encoding_kwargs is not None
-        output_pos_enc = TrainablePositionEncoding(
-            # Construct 1D features:
-            index_dim=np.prod(index_dims),
-            name=name,
-            **trainable_position_encoding_kwargs,
-        )
-    elif position_encoding_type == "fourier":
-        assert fourier_position_encoding_kwargs is not None
-        output_pos_enc = FourierPositionEncoding(
-            index_dims=index_dims, name=name, **fourier_position_encoding_kwargs
-        )
-    else:
-        raise ValueError(f"Unknown position encoding: {position_encoding_type}.")
-
-    if project_pos_dim > 0:
-        # Project the position encoding to a target dimension:
-        output_pos_enc = PositionEncodingProjector(
-            output_size=project_pos_dim, base_position_encoding=output_pos_enc
-        )
-
-    return output_pos_enc
diff --git a/satflow/data/utils/sampling.py b/satflow/data/utils/sampling.py
deleted file mode 100644
index 29ce993c..00000000
--- a/satflow/data/utils/sampling.py
+++ /dev/null
@@ -1,11 +0,0 @@
-import numpy as np
-
-OCEAN = 0
-MOUNTAIN = 2000
-HILLY = 100
-
-
-def calc_geo(topo_map: np.ndarray):
-    ocean_frac = len(topo_map[topo_map <= OCEAN])
-    hill_frac = len(topo_map[(topo_map > OCEAN) & (topo_map <= HILLY)])
-    mt_frac = len(topo_map[(topo_map > HILLY) & (topo_map <= MOUNTAIN)])
diff --git a/satflow/examples/collate_data.py b/satflow/examples/collate_data.py
deleted file mode 100644
index 19cb3461..00000000
--- a/satflow/examples/collate_data.py
+++ /dev/null
@@ -1,68 +0,0 @@
-import os
-import subprocess
-
-from satflow.data.utils.utils import map_satellite_to_mercator
-
-eumetsat_dir = "/run/media/bieker/Round1/EUMETSAT/2020/"
-
-
-def decompress(full_bzip_filename: str) -> str:
-    base_bzip_filename = os.path.basename(full_bzip_filename)
-    base_nat_filename = os.path.splitext(base_bzip_filename)[0]
-    full_nat_filename = os.path.join("./", base_nat_filename)
-    if os.path.exists(full_nat_filename):
-        os.remove(full_nat_filename)
-    with open(full_nat_filename, "wb") as nat_file_handler:
-        process = subprocess.run(
-            ["pbzip2", "--decompress", "--keep", "--stdout", full_bzip_filename],
-            stdout=nat_file_handler,
-        )
-    process.check_returncode()
-    return full_nat_filename
-
-
-for root, dirs, files in os.walk(eumetsat_dir):
-    # For each of these, load cloud mask, etc. and save to GeoTIFF
-    sat_file = ""
-    cloud_mask = ""
-    try:
-        if len(files) > 10:
-            print(f"skipping: {root}")
-            continue
-        for f in files:
-            if ".bz2" in f:
-                sat_file = os.path.join(root, f)
-            elif ".nat" in f:
-                sat_file = os.path.join(root, f)
-            if ".grb" in f:
-                cloud_mask = os.path.join(root, f)
-            if sat_file != "" and cloud_mask != "":
-                sat_file = decompress(sat_file)
-                map_satellite_to_mercator(
-                    native_satellite=sat_file,
-                    grib_files=cloud_mask,
-                    save_loc=root,
-                    bands=(
-                        "HRV",
-                        "IR_016",
-                        "IR_039",
-                        "IR_087",
-                        "IR_097",
-                        "IR_108",
-                        "IR_120",
-                        "IR_134",
-                        "VIS006",
-                        "VIS008",
-                        "WV_062",
-                        "WV_073",
-                        "cloud_mask",
-                    ),
-                )
-                os.remove(sat_file)
-    except Exception as e:
-        print(e)
-        try:
-            os.remove(sat_file)
-        except:
-            continue
-        continue
diff --git a/satflow/examples/create_dem.py b/satflow/examples/create_dem.py
deleted file mode 100644
index 715e4b4b..00000000
--- a/satflow/examples/create_dem.py
+++ /dev/null
@@ -1,64 +0,0 @@
-import glob
-
-import numpy as np
-import rasterio
-from pyresample import load_area
-from rasterio.merge import merge
-from rasterio.plot import show
-from satpy import Scene
-
-
-dem_path = "/home/bieker/bda/Bulk Order 20210611_041625/SRTM 1 Arc-Second Global/*.tif"
-dem_files = glob.glob(dem_path)
-
-src_files_to_mosaic = []
-for fp in dem_files:
-    src = rasterio.open(fp)
-    # src.astype(np.int8)
-    src_files_to_mosaic.append(src)
-
-areas = load_area("/home/bieker/Development/satflow/satflow/examples/areas.yaml")
-for f in dem_files:
-    scene = Scene(filenames={"generic_image": [f]})
-    scene.load(["image"])
-    scene = scene.resample(areas[0])
-    fname = f.split("/")[-1]
-    scene.save_datasets(
-        writer="geotiff",
-        base_dir="/run/media/bieker/T7/ResampledElevation/",
-        filename=fname,
-        enhance=False,
-    )
-
-dem_path = "/run/media/bieker/T7/ResampledElevation/*.tif"
-dem_files = glob.glob(dem_path)
-
-src_files_to_mosaic = []
-for fp in dem_files:
-    src = rasterio.open(fp)
-    # src.astype(np.int8)
-    src_files_to_mosaic.append(src)
-mosaic, out_trans = merge(src_files_to_mosaic, nodata=0, method="max")
-print(mosaic.shape)
-show(mosaic[0], cmap="terrain")
-show(mosaic[0][515:-641, 603:], cmap="terrain")
-out_meta = src.meta.copy()
-
-out_meta.update(
-    {
-        "driver": "GTiff",
-        "height": mosaic.shape[1],
-        "width": mosaic.shape[2],
-        "transform": out_trans,
-        "crs": "+proj=utm +zone=35 +ellps=WGS84 +units=m +no_defs ",
-    }
-)
-out_fp = "europe_all_dem_3km_satproj.tif"
-with rasterio.open(out_fp, "w", **out_meta) as dest:
-    dest.write(mosaic)
-
-out = rasterio.open(out_fp)
-show(out, cmap="terrain")
-# Save the numpy array, we don't care about other geotspatial transforms here
-np.save("../resources/cutdown_europe_dem.npy", mosaic[0][515:-641, 603:])
-np.save("europe_dem.npy", mosaic[0])
diff --git a/satflow/examples/create_webdataset.py b/satflow/examples/create_webdataset.py
deleted file mode 100644
index 0cbebe2b..00000000
--- a/satflow/examples/create_webdataset.py
+++ /dev/null
@@ -1,235 +0,0 @@
-import logging
-import os
-from datetime import datetime
-
-import numpy as np
-import rasterio
-import satpy
-import webdataset as wds
-from pyresample import load_area
-from rasterio.plot import show
-from rasterio.warp import Resampling, calculate_default_transform, reproject
-from rasterio.windows import Window, from_bounds
-from satpy import Scene
-
-logger = logging.getLogger("satflow")
-logger.setLevel(logging.DEBUG)
-
-#
-areas = load_area("/home/bieker/Development/satflow/satflow/resources/areas.yaml")
-# filenames = {
-#    "seviri_l1b_native": [
-#        "/run/media/bieker/Round1/EUMETSAT/2021/03/14/10/19/MSG3-SEVI-MSG15-0100-NA-20210314101915.098000000Z-NA.nat"
-#    ]
-# }
-# scene = Scene(filenames=filenames)
-# scene.load(("HRV",))
-# By default resamples to 3km, as thats the native resolution of all bands other than HRV
-# scene = scene.resample(areas[0])
-# hrv_long, hrv_lat = scene["HRV"].attrs["area"].get_lonlats()
-# hrv_long = hrv_long[515:-641, 603:]
-# hrv_lat = hrv_lat[515:-641, 603:]
-# loc_x = np.cos(hrv_lat) * np.cos(hrv_long)
-# loc_y = np.cos(hrv_lat) * np.sin(hrv_long)
-# loc_z = np.sin(hrv_lat)
-
-# location_array = np.stack([loc_x, loc_y, loc_z], axis=-1).astype(np.float32) # We aren't using float64 anyway, so some small overlap of locations should be fine?
-# print(location_array.dtype)
-# print(len(np.unique(location_array)))
-# print(len(np.unique(location_array.astype(np.float32))))
-# print(len(np.unique(location_array.astype(np.float16))))
-# np.save("../resources/location_array.npy", location_array)
-topo_data = np.load("../resources/cutdown_europe_dem.npy")
-# print(f"Channel: {channel}")
-topo_data[topo_data < -100] = 0  # Ocean
-location_array = np.load("../resources/location_array.npy")
-# print(topo_data.dtype)
-# print(topo_data.shape)
-# print(location_array.shape)
-box_1 = (0, 446, 0, 978)
-box_2 = (446, -1, 978, -1)
-# Create WebDataset with each shard being a single day and all the images for that day
-eumetsat_dir = "/run/media/bieker/Round1/EUMETSAT/"
-
-
-def make_day(data, flow=True, batch=146, tile=True):
-    root_dir, shard_num = data
-    shard_num += 150  # Start after the current ones
-    # reset optical flow samples
-    flow_sample = {}
-    # Initialize new flow sample
-    date_data = root_dir.split("/")
-    date_str = os.path.join(date_data[-3], date_data[-2], date_data[-1])
-    overall_datetime = datetime.strptime(date_str, "%Y/%m/%d")
-    batch_num = 0
-    if batch < 1:
-        flow_sample["__key__"] = overall_datetime.strftime("%Y/%m/%d")
-    else:
-        flow_sample["__key__"] = overall_datetime.strftime("%Y/%m/%d") + f"{batch_num}"
-    # Split data into 8 chunks? 446x489 squareish areas
-    samples = [
-        {"__key__": overall_datetime.strftime("%Y/%m/%d") + f"_{i}", "time.pyd": []}
-        for i in range(8)
-    ]
-    flow_sample["time.pyd"] = []
-    if flow:
-        flow_sample["topo.npy"] = topo_data
-        flow_sample["location.npy"] = location_array
-    print(f"On shard: {shard_num} Date: {overall_datetime.strftime('%Y/%m/%d')}")
-    shard_num += 1
-    interday_frame = 0
-    if os.path.exists(
-        f"/run/media/bieker/data/EUMETSAT/satflow{'-' if not flow else '-flow'}{'-' if not flow and batch > 0 else f'-{batch}-'}{'tiled-' if tile else ''}{shard_num:05d}.tar"
-    ):
-        return
-    sink_flow = wds.TarWriter(
-        f"/run/media/bieker/data/EUMETSAT/satflow{'-' if not flow else '-flow'}{'-' if not flow and batch > 0 else f'-{batch}-'}{'tiled-' if tile else ''}{shard_num:05d}.tar",
-        compress=True,
-    )
-    for root, dirs, files in os.walk(root_dir):
-        dirs.sort(
-            key=int
-        )  # Ensure this is done in sorted order, so its always from earliest to latest
-        # As its stored in year, month, day, hour, minute, format, once we are at the day level, create a new shard
-        # Also cut down everything to only the square of all value data
-        if files:
-            # At lowest one, so have everything in it
-            if len(files) < 11:  # Not everything worked
-                continue
-
-            date_data = root.split("/")
-            date_str = os.path.join(
-                date_data[-5], date_data[-4], date_data[-3], date_data[-2], date_data[-1]
-            )
-            datetime_object = datetime.strptime(date_str, "%Y/%m/%d/%H/%M")
-            if not flow:
-                flow_sample = {}
-                flow_sample["__key__"] = datetime_object.strftime("%Y/%m/%d/%H/%M")
-
-            # try:
-            # Want this at the top, as if this extraction fails, we still then know if there is a missing frame
-            interday_frame += 1
-            for f in files:
-                if ".tif" in f:
-                    # Only need geotiff ones
-                    src = rasterio.open(os.path.join(root, f))
-                    data = src.read(1)
-                    cropped_data = data[
-                        515:-641, 603:
-                    ]  # Done by handish to exclude all NODATA and invalid masks for clouds and images (clouds have a smaller extant)
-                    channel = f.split("_")
-                    from scipy import stats
-
-                    # print(f"Channel: {channel}")
-                    # print(stats.describe(cropped_data, axis=None))
-                    # print(stats.describe(cropped_data.astype(np.float16), axis=None))
-                    # TODO Ensure this is fine, from above it seems like it should be, values didn't shift too much
-                    cropped_data = cropped_data.astype(np.float16)
-                    if channel[0] in ["IR", "WV", "cloud"]:
-                        channel = channel[0] + channel[1]  # These are split across multiple
-                    else:
-                        channel = channel[0]
-                    if tile:
-                        height = (0, 446, -1)
-                        width = (0, 490, 980, 1470, -1)
-                        s_num = 0
-                        for j in range(4):
-                            for i in range(2):
-                                tile_cropped_data = cropped_data[
-                                    height[i] : height[i + 1], width[j] : width[j + 1]
-                                ]
-                                tiled_topo = topo_data[
-                                    height[i] : height[i + 1], width[j] : width[j + 1]
-                                ]
-                                tiled_loc = location_array[
-                                    height[i] : height[i + 1], width[j] : width[j + 1]
-                                ]
-                                samples[s_num][
-                                    channel + f".{interday_frame:03d}.npy"
-                                ] = tile_cropped_data.astype(np.float16)
-                                samples[s_num]["location.npy"] = tiled_loc
-                                samples[s_num]["topo.npy"] = tiled_topo
-                                s_num += 1
-
-                    elif not flow:
-                        flow_sample[channel + ".npy"] = cropped_data
-                    else:
-                        flow_sample[channel + f".{interday_frame:03d}.npy"] = cropped_data.astype(
-                            np.float16
-                        )
-            # Now all channels added to thing, write to shard
-            if flow:
-                if not tile:
-                    flow_sample["time.pyd"].append(datetime_object)
-                for i, s in enumerate(samples):
-                    samples[i]["time.pyd"].append(datetime_object)
-                if batch > 0:
-                    print(f"In Batch: {len(samples[0]['time.pyd'])} == {batch} Shard: {shard_num}")
-                    if len(samples[0]["time.pyd"]) == batch:
-                        if tile:
-                            for s in samples:
-                                sink_flow.write(s)
-                        else:
-                            sink_flow.write(flow_sample)
-                            flow_sample["__key__"] = datetime_object.strftime("%Y/%m/%d/%H/%M")
-                            flow_sample["topo.npy"] = topo_data
-                            flow_sample["location.npy"] = location_array
-                            flow_sample["time.pyd"] = []
-                        interday_frame = 0
-                        flow_sample = {}
-                        samples = [
-                            {
-                                "__key__": overall_datetime.strftime("%Y/%m/%d") + f"_{i}",
-                                "time.pyd": [],
-                            }
-                            for i in range(8)
-                        ]
-                        batch_num += 1
-            else:
-                flow_sample["time.pyd"] = datetime_object
-                sink_flow.write(flow_sample)
-            # except Exception as e:
-            #    print(e)
-    # print(f"Write Sample: {flow_sample.keys()}")
-    # if flow and batch <= 0:
-    #    sink_flow.write(flow_sample)
-    if tile:
-        for s in samples:
-            sink_flow.write(s)
-    # Close old shard
-    # sink.close()
-    sink_flow.close()
-
-
-import multiprocessing
-
-pool = multiprocessing.Pool(3)
-
-old = os.listdir(os.path.join(eumetsat_dir, "2020"))
-old.sort(key=int)
-old = [os.path.join(eumetsat_dir, "2020", d) for d in old]
-tmp_old = []
-for d in old:
-    days = os.listdir(os.path.join(d))
-    days.sort(key=int)
-    tmp_old += [os.path.join(d, day) for day in days]
-old = tmp_old
-new = os.listdir(os.path.join(eumetsat_dir, "2021"))
-new.sort(key=int)
-new = [os.path.join(eumetsat_dir, "2021", d) for d in new]
-tmp_new = []
-for d in new:
-    days = os.listdir(os.path.join(d))
-    days.sort(key=int)
-    tmp_new += [os.path.join(d, day) for day in days]
-new = tmp_new
-all_dates = old + new
-print(all_dates)
-all_dates = zip(all_dates, range(len(all_dates)))
-# pool.map(make_day, all_dates)
-# exit()
-for data in all_dates:
-    try:
-        make_day(data)
-    except:
-        continue
diff --git a/satflow/examples/metnet_example.py b/satflow/examples/metnet_example.py
index 7131343a..69647745 100644
--- a/satflow/examples/metnet_example.py
+++ b/satflow/examples/metnet_example.py
@@ -1,4 +1,4 @@
-from satflow.models import MetNet
+from satflow.models import LitMetNet
 import torch
 import urllib.request
 
@@ -13,7 +13,7 @@ def get_input_target(number: int):
 # Setup the model (need to add loading weights from HuggingFace :)
 # 12 satellite channels + 1 Topographic + 3 Lat/Lon + 1 Cloud Mask
 # Output Channels: 1 Cloud mask, 12 for Satellite image
-model = MetNet(input_channels=17, sat_channels=13, input_size=64, out_channels=1)
+model = LitMetNet(input_channels=17, sat_channels=13, input_size=64, out_channels=1)
 torch.set_grad_enabled(False)
 model.eval()
 # The inputs are Tensors of size (Batch, Curr+Prev Timesteps, Channel, Width, Height)
diff --git a/satflow/experiments/pl_train.py b/satflow/experiments/pl_train.py
deleted file mode 100644
index fa41b13a..00000000
--- a/satflow/experiments/pl_train.py
+++ /dev/null
@@ -1,98 +0,0 @@
-from typing import List, Optional
-
-import hydra
-from omegaconf import DictConfig
-from pytorch_lightning import (
-    Callback,
-    LightningModule,
-    LightningDataModule,
-    Trainer,
-    seed_everything,
-)
-from pytorch_lightning.loggers import LightningLoggerBase
-from pytorch_lightning.profiler import AdvancedProfiler, PyTorchProfiler
-from satflow.core import utils
-from satflow.core.callbacks import NeptuneModelLogger
-from pytorch_lightning.callbacks import LearningRateMonitor
-
-log = utils.get_logger(__name__)
-
-
-def train(config: DictConfig) -> Optional[float]:
-    """Contains training pipeline.
-    Instantiates all PyTorch Lightning objects from config.
-
-    Args:
-        config (DictConfig): Configuration composed by Hydra.
-
-    Returns:
-        Optional[float]: Metric score for hyperparameter optimization.
-    """
-
-    # Set seed for random number generators in pytorch, numpy and python.random
-    if "seed" in config:
-        seed_everything(config.seed, workers=True)
-
-    # If required
-    # Init Dataloaders
-    log.info(f"Instantiating datamodule <{config.datamodule._target_}>")
-    datamodule: LightningDataModule = hydra.utils.instantiate(
-        config.datamodule, _convert_="partial"
-    )
-
-    # Init Lightning model
-    log.info(f"Instantiating model <{config.model._target_}>")
-    model: LightningModule = hydra.utils.instantiate(config.model)
-
-    # Init Lightning callbacks
-    lr_monitor = LearningRateMonitor(logging_interval="step")
-    callbacks: List[Callback] = [lr_monitor, NeptuneModelLogger()]
-    if "callbacks" in config:
-        for _, cb_conf in config["callbacks"].items():
-            if "_target_" in cb_conf:
-                log.info(f"Instantiating callback <{cb_conf._target_}>")
-                callbacks.append(hydra.utils.instantiate(cb_conf))
-
-    # Init Lightning loggers
-    logger: List[LightningLoggerBase] = []
-    if "logger" in config:
-        for _, lg_conf in config["logger"].items():
-            if "_target_" in lg_conf:
-                log.info(f"Instantiating logger <{lg_conf._target_}>")
-                logger.append(hydra.utils.instantiate(lg_conf))
-
-    # Init Lightning trainer
-    log.info(f"Instantiating trainer <{config.trainer._target_}>")
-    trainer: Trainer = hydra.utils.instantiate(
-        config.trainer,
-        callbacks=callbacks,
-        logger=logger,
-    )
-
-    # Send some parameters from config to all lightning loggers
-    log.info("Logging hyperparameters!")
-    utils.log_hyperparameters(
-        config=config,
-        model=model,
-        trainer=trainer,
-    )
-
-    # Train the model
-    if config.trainer.auto_lr_find or config.trainer.auto_scale_batch_size:
-        log.info("Starting tuning!")
-        trainer.tune(model=model, datamodule=datamodule)
-    log.info("Starting training!")
-    trainer.fit(model=model, datamodule=datamodule)
-
-    # Evaluate model on test set after training
-    if not config.trainer.get("fast_dev_run", False):
-        log.info("Starting testing!")
-        trainer.test(model=model, datamodule=datamodule)
-
-    # Print path to best checkpoint
-    log.info(f"Best checkpoint path:\n{trainer.checkpoint_callback.best_model_path}")
-
-    # Return metric score for hyperparameter optimization
-    optimized_metric = config.get("optimized_metric")
-    if optimized_metric:
-        return trainer.callback_metrics[optimized_metric]
diff --git a/satflow/experiments/train.py b/satflow/experiments/train.py
index 388c52e7..fa41b13a 100644
--- a/satflow/experiments/train.py
+++ b/satflow/experiments/train.py
@@ -1,56 +1,98 @@
-import os
+from typing import List, Optional
 
-import torch
-import torch.nn.functional as F
-import webdataset as wds
-from pytorch_lightning import Trainer, seed_everything
-from pytorch_lightning.tuner.tuning import Tuner
-from torch.utils.data import DataLoader
+import hydra
+from omegaconf import DictConfig
+from pytorch_lightning import (
+    Callback,
+    LightningModule,
+    LightningDataModule,
+    Trainer,
+    seed_everything,
+)
+from pytorch_lightning.loggers import LightningLoggerBase
+from pytorch_lightning.profiler import AdvancedProfiler, PyTorchProfiler
+from satflow.core import utils
+from satflow.core.callbacks import NeptuneModelLogger
+from pytorch_lightning.callbacks import LearningRateMonitor
 
-import satflow.models
-from satflow.core.training_utils import get_args, get_loaders, setup_experiment
-from satflow.core.utils import load_config, make_logger
-from satflow.data.datasets import SatFlowDataset
-from satflow.models import get_model
+log = utils.get_logger(__name__)
 
-logger = make_logger("satflow.train")
 
+def train(config: DictConfig) -> Optional[float]:
+    """Contains training pipeline.
+    Instantiates all PyTorch Lightning objects from config.
 
-def run_experiment(args):
-    config = setup_experiment(args)
-    config["device"] = (
-        ("cuda" if torch.cuda.is_available() else "cpu") if not args.with_cpu else "cpu"
+    Args:
+        config (DictConfig): Configuration composed by Hydra.
+
+    Returns:
+        Optional[float]: Metric score for hyperparameter optimization.
+    """
+
+    # Set seed for random number generators in pytorch, numpy and python.random
+    if "seed" in config:
+        seed_everything(config.seed, workers=True)
+
+    # If required
+    # Init Dataloaders
+    log.info(f"Instantiating datamodule <{config.datamodule._target_}>")
+    datamodule: LightningDataModule = hydra.utils.instantiate(
+        config.datamodule, _convert_="partial"
+    )
+
+    # Init Lightning model
+    log.info(f"Instantiating model <{config.model._target_}>")
+    model: LightningModule = hydra.utils.instantiate(config.model)
+
+    # Init Lightning callbacks
+    lr_monitor = LearningRateMonitor(logging_interval="step")
+    callbacks: List[Callback] = [lr_monitor, NeptuneModelLogger()]
+    if "callbacks" in config:
+        for _, cb_conf in config["callbacks"].items():
+            if "_target_" in cb_conf:
+                log.info(f"Instantiating callback <{cb_conf._target_}>")
+                callbacks.append(hydra.utils.instantiate(cb_conf))
+
+    # Init Lightning loggers
+    logger: List[LightningLoggerBase] = []
+    if "logger" in config:
+        for _, lg_conf in config["logger"].items():
+            if "_target_" in lg_conf:
+                log.info(f"Instantiating logger <{lg_conf._target_}>")
+                logger.append(hydra.utils.instantiate(lg_conf))
+
+    # Init Lightning trainer
+    log.info(f"Instantiating trainer <{config.trainer._target_}>")
+    trainer: Trainer = hydra.utils.instantiate(
+        config.trainer,
+        callbacks=callbacks,
+        logger=logger,
     )
 
-    # Load Model
-    model = get_model(config["model"]["name"]).from_config(config["model"])
-    # Load Datasets
-    loaders = get_loaders(config["dataset"])
-
-    # Get batch size that fits in memory
-
-    # trainer = Trainer(auto_scale_batch_size='power')
-    # trainer.tune(model)
-    # tuner = Tuner(trainer)
-
-    # new_batch_size = tuner.scale_batch_size(model)
-    trainer = Trainer(
-        gpus=1,
-        auto_lr_find=True,
-        max_steps=config["training"]["iterations"],
-        min_epochs=0,
-        val_check_interval=config["training"]["eval_steps"],
-        limit_train_batches=2000,
-        limit_val_batches=config["training"]["eval_examples"],
-        accumulate_grad_batches=4,
-        profiler="simple",
+    # Send some parameters from config to all lightning loggers
+    log.info("Logging hyperparameters!")
+    utils.log_hyperparameters(
+        config=config,
+        model=model,
+        trainer=trainer,
     )
 
-    trainer.tune(model, loaders["train"], loaders["test"])
+    # Train the model
+    if config.trainer.auto_lr_find or config.trainer.auto_scale_batch_size:
+        log.info("Starting tuning!")
+        trainer.tune(model=model, datamodule=datamodule)
+    log.info("Starting training!")
+    trainer.fit(model=model, datamodule=datamodule)
 
-    trainer.fit(model, loaders["train"], loaders["test"])
+    # Evaluate model on test set after training
+    if not config.trainer.get("fast_dev_run", False):
+        log.info("Starting testing!")
+        trainer.test(model=model, datamodule=datamodule)
 
+    # Print path to best checkpoint
+    log.info(f"Best checkpoint path:\n{trainer.checkpoint_callback.best_model_path}")
 
-if __name__ == "__main__":
-    args = get_args()
-    run_experiment(args)
+    # Return metric score for hyperparameter optimization
+    optimized_metric = config.get("optimized_metric")
+    if optimized_metric:
+        return trainer.callback_metrics[optimized_metric]
diff --git a/satflow/models/__init__.py b/satflow/models/__init__.py
index 06b8b048..63a609a0 100644
--- a/satflow/models/__init__.py
+++ b/satflow/models/__init__.py
@@ -1,8 +1,6 @@
-from .base import get_model, create_model
+from nowcasting_utils.models.base import get_model, create_model
 from .conv_lstm import EncoderDecoderConvLSTM, ConvLSTM
 from .pl_metnet import LitMetNet
-
 from .runet import R2U_Net, RUnet
 from .attention_unet import R2AttU_Net, AttU_Net
-
 from .perceiver import Perceiver
diff --git a/satflow/models/attention_unet.py b/satflow/models/attention_unet.py
index d2b65eb2..47486b6e 100644
--- a/satflow/models/attention_unet.py
+++ b/satflow/models/attention_unet.py
@@ -3,8 +3,9 @@
 import pytorch_lightning as pl
 import torchvision
 import numpy as np
-from satflow.models.losses import FocalLoss, get_loss
-from satflow.models.base import register_model
+from nowcasting_utils.models.loss import get_loss
+from nowcasting_utils.models.losses.FocalLoss import FocalLoss
+from nowcasting_utils.models.base import register_model
 
 
 @register_model
diff --git a/satflow/models/base.py b/satflow/models/base.py
deleted file mode 100644
index ec7782d9..00000000
--- a/satflow/models/base.py
+++ /dev/null
@@ -1,164 +0,0 @@
-from typing import Any, Dict, Type
-import torch.nn
-import pytorch_lightning as pl
-import torchvision
-from neptune.new.types import File
-from satflow.models.hub import load_model_config_from_hf, load_pretrained, SatFlowModelHubMixin
-
-
-REGISTERED_MODELS = {}
-
-
-def register_model(cls: Type[pl.LightningModule]):
-    global REGISTERED_MODELS
-    name = cls.__name__
-    assert name not in REGISTERED_MODELS, f"exists class: {REGISTERED_MODELS}"
-    REGISTERED_MODELS[name] = cls
-    return cls
-
-
-def get_model(name: str) -> Type[pl.LightningModule]:
-    global REGISTERED_MODELS
-    assert name in REGISTERED_MODELS, f"available class: {REGISTERED_MODELS}"
-    return REGISTERED_MODELS[name]
-
-
-def list_models():
-    global REGISTERED_MODELS
-    return REGISTERED_MODELS.keys()
-
-
-def split_model_name(model_name):
-    model_split = model_name.split(":", 1)
-    if len(model_split) == 1:
-        return "", model_split[0]
-    else:
-        source_name, model_name = model_split
-        assert source_name in ("satflow", "hf_hub")
-        return source_name, model_name
-
-
-def safe_model_name(model_name, remove_source=True):
-    def make_safe(name):
-        return "".join(c if c.isalnum() else "_" for c in name).rstrip("_")
-
-    if remove_source:
-        model_name = split_model_name(model_name)[-1]
-    return make_safe(model_name)
-
-
-def create_model(model_name, pretrained=False, checkpoint_path=None, **kwargs):
-    """Create a model
-
-    Almost entirely taken from timm https://github.com/rwightman/pytorch-image-models
-
-    Args:
-        model_name (str): name of model to instantiate
-        pretrained (bool): load pretrained ImageNet-1k weights if true
-        checkpoint_path (str): path of checkpoint to load after model is initialized
-
-    Keyword Args:
-        drop_rate (float): dropout rate for training (default: 0.0)
-        global_pool (str): global pool type (default: 'avg')
-        input_channels (int): number of input channels (default: 12)
-        forecast_steps (int): number of steps to forecast (default: 48)
-        lr (float): learning rate (default: 0.001)
-        **: other kwargs are model specific
-    """
-    source_name, model_name = split_model_name(model_name)
-
-    # Parameters that aren't supported by all models or are intended to only override model defaults if set
-    # should default to None in command line args/cfg. Remove them if they are present and not set so that
-    # non-supporting models don't break and default args remain in effect.
-    kwargs = {k: v for k, v in kwargs.items() if v is not None}
-
-    if source_name == "hf_hub":
-        # For model names specified in the form `hf_hub:path/architecture_name#revision`,
-        # load model weights + default_cfg from Hugging Face hub.
-        hf_default_cfg, model_name = load_model_config_from_hf(model_name)
-        # Want to set the kwargs to correct values
-        # Should ignore this after we get it in the model_name
-        kwargs.update(hf_default_cfg)
-        kwargs.pop("hf_hub")
-        kwargs.pop("architecture")
-
-    if model_name in REGISTERED_MODELS:
-        model = get_model(model_name)
-    else:
-        raise RuntimeError("Unknown model (%s)" % model_name)
-
-    if checkpoint_path or pretrained:
-        kwargs["checkpoint_path"] = checkpoint_path
-        # Readd hf_hub here
-        if source_name == "hf_hub":
-            kwargs["hf_hub"] = hf_default_cfg.get("hf_hub")
-        model = load_pretrained(model, default_cfg=kwargs, in_chans=kwargs["input_channels"])
-    else:
-        # Initialize model here as LightingModules need a special way of loading checkpoints, this initializes randomly
-        model = model(**kwargs)
-    return model
-
-
-class BaseModel(pl.LightningModule, SatFlowModelHubMixin):
-    def __init__(
-        self,
-        pretrained: bool = False,
-        forecast_steps: int = 48,
-        input_channels: int = 12,
-        output_channels: int = 12,
-        lr: float = 0.001,
-        visualize: bool = False,
-    ):
-        super(BaseModel, self).__init__()
-        self.forecast_steps = forecast_steps
-        self.input_channels = input_channels
-        self.output_channels = output_channels
-        self.lr = lr
-        self.pretrained = pretrained
-        self.visualize = visualize
-
-    @classmethod
-    def from_config(cls, config):
-        raise NotImplementedError
-
-    def _train_or_validate_step(self, batch, batch_idx, is_training: bool = True):
-        pass
-
-    def training_step(self, batch, batch_idx):
-        return self._train_or_validate_step(batch, batch_idx, is_training=True)
-
-    def validation_step(self, batch, batch_idx):
-        return self._train_or_validate_step(batch, batch_idx, is_training=False)
-
-    def forward(self, x, **kwargs) -> Any:
-        return self.model.forward(x, **kwargs)
-
-    def visualize_step(
-        self, x: torch.Tensor, y: torch.Tensor, y_hat: torch.Tensor, batch_idx: int, step: str
-    ) -> None:
-        # the logger you used (in this case tensorboard)
-        tensorboard = self.logger.experiment[0]
-        # Timesteps per channel
-        images = x[0].cpu().detach().float()  # T, C, H, W
-        future_images = y[0].cpu().detach().float()
-        generated_images = y_hat[0].cpu().detach().float()
-        for i, t in enumerate(images):  # Now would be (C, H, W)
-            t = [torch.unsqueeze(img, dim=0) for img in t]
-            image_grid = torchvision.utils.make_grid(
-                t, nrow=self.input_channels // 2, normalize=True
-            )
-            # Neptune needs it to be H x W x C, so have to permute
-            image_grid = image_grid.permute(1, 2, 0)
-            tensorboard[f"{step}/Input_Frame_{i}"].log(File.as_image(image_grid.numpy()))
-            t = [torch.unsqueeze(img, dim=0) for img in future_images[i]]
-            image_grid = torchvision.utils.make_grid(
-                t, nrow=self.output_channels // 2, normalize=True
-            )
-            image_grid = image_grid.permute(1, 2, 0)
-            tensorboard[f"{step}/Target_Frame_{i}"].log(File.as_image(image_grid.numpy()))
-            t = [torch.unsqueeze(img, dim=0) for img in generated_images[i]]
-            image_grid = torchvision.utils.make_grid(
-                t, nrow=self.output_channels // 2, normalize=True
-            )
-            image_grid = image_grid.permute(1, 2, 0)
-            tensorboard[f"{step}/Predicted_Frame_{i}"].log(File.as_image(image_grid.numpy()))
diff --git a/satflow/models/cloudgan.py b/satflow/models/cloudgan.py
index 12a8ef53..46cab944 100644
--- a/satflow/models/cloudgan.py
+++ b/satflow/models/cloudgan.py
@@ -6,7 +6,7 @@
 from satflow.models import R2U_Net, ConvLSTM
 from satflow.models.gan import GANLoss, define_generator, define_discriminator
 from satflow.models.layers import ConditionTime
-from satflow.models.losses import get_loss
+from nowcasting_utils.models.loss import get_loss
 from pl_bolts.optimizers.lr_scheduler import LinearWarmupCosineAnnealingLR
 import numpy as np
 
diff --git a/satflow/models/conv_lstm.py b/satflow/models/conv_lstm.py
index 8e25c7e1..1c30629c 100644
--- a/satflow/models/conv_lstm.py
+++ b/satflow/models/conv_lstm.py
@@ -6,8 +6,8 @@
 
 import numpy as np
 
-from satflow.models.base import register_model
-from satflow.models.losses import get_loss
+from nowcasting_utils.models.base import register_model
+from nowcasting_utils.models.loss import get_loss
 from satflow.models.layers.ConvLSTM import ConvLSTMCell
 import torchvision
 
diff --git a/satflow/models/deeplabv3.py b/satflow/models/deeplabv3.py
index 7c0ee1f3..adefe2be 100644
--- a/satflow/models/deeplabv3.py
+++ b/satflow/models/deeplabv3.py
@@ -1,11 +1,11 @@
 import torch
 import torch.nn.functional as F
 import pytorch_lightning as pl
-from satflow.models.base import register_model
+from nowcasting_utils.models.base import register_model
 from torchvision.models.segmentation import deeplabv3_resnet50, deeplabv3_resnet101
 import numpy as np
 from typing import Union
-from satflow.models.losses import FocalLoss
+from nowcasting_utils.models.losses.FocalLoss import FocalLoss
 
 
 @register_model
diff --git a/satflow/models/fcn.py b/satflow/models/fcn.py
index ebbaec1b..d0ea4030 100644
--- a/satflow/models/fcn.py
+++ b/satflow/models/fcn.py
@@ -1,11 +1,11 @@
 import torch
 import torch.nn.functional as F
 import pytorch_lightning as pl
-from satflow.models.base import register_model
+from nowcasting_utils.models.base import register_model
 from torchvision.models.segmentation import fcn_resnet50, fcn_resnet101
 import numpy as np
 from typing import Union
-from satflow.models.losses import FocalLoss
+from nowcasting_utils.models.losses.FocalLoss import FocalLoss
 
 
 @register_model
diff --git a/satflow/models/gan/common.py b/satflow/models/gan/common.py
index a82fcba2..394392e0 100644
--- a/satflow/models/gan/common.py
+++ b/satflow/models/gan/common.py
@@ -1,14 +1,6 @@
 import functools
-from typing import Tuple
-
 import torch
 from torch.nn import init
-from torch.distributions import uniform
-from torch.nn.utils import spectral_norm
-from torch.nn.modules.pixelshuffle import PixelShuffle, PixelUnshuffle
-from torch.nn.functional import interpolate
-from satflow.models.utils import get_conv_layer
-from satflow.models.layers.Attention import SelfAttention2d
 
 
 def get_norm_layer(norm_type="instance"):
@@ -108,9 +100,7 @@ def cal_gradient_penalty(
     Returns the gradient penalty loss
     """
     if lambda_gp > 0.0:
-        if (
-            type == "real"
-        ):  # either use real images, fake images, or a linear interpolation of two.
+        if type == "real":  # either use real images, fake images, or a linear interpolation of two.
             interpolatesv = real_data
         elif type == "fake":
             interpolatesv = fake_data
diff --git a/satflow/models/gan/discriminators.py b/satflow/models/gan/discriminators.py
index 507128d4..07b5e16a 100644
--- a/satflow/models/gan/discriminators.py
+++ b/satflow/models/gan/discriminators.py
@@ -1,11 +1,8 @@
 import functools
 import torch
-from torch.nn.modules.pixelshuffle import PixelShuffle, PixelUnshuffle
-from torch.nn.utils import spectral_norm
 from torch import nn as nn
-from torchvision.transforms import RandomCrop
 from satflow.models.utils import get_conv_layer
-from satflow.models.gan.common import get_norm_layer, init_net, LBlock, GBlock, DBlock
+from satflow.models.gan.common import get_norm_layer, init_net
 import antialiased_cnns
 
 
diff --git a/satflow/models/gan/generators.py b/satflow/models/gan/generators.py
index 6a3db73a..794e13a4 100644
--- a/satflow/models/gan/generators.py
+++ b/satflow/models/gan/generators.py
@@ -1,13 +1,9 @@
 import functools
-
 import torch
 from torch import nn as nn
-from torch.nn.modules.pixelshuffle import PixelShuffle
-from torch.nn.utils import spectral_norm
-from typing import Union, Tuple, List
-from satflow.models.gan.common import get_norm_layer, init_net, GBlock
+from typing import Union
+from satflow.models.gan.common import get_norm_layer, init_net
 from satflow.models.utils import get_conv_layer
-from satflow.models.layers import ConvGRU
 import antialiased_cnns
 
 
@@ -386,15 +382,11 @@ def __init__(
         conv2d = get_conv_layer(conv_type)
         if conv_type == "antialiased":
             antialiased = True
-            downconv = conv2d(
-                input_nc, inner_nc, kernel_size=4, stride=1, padding=1, bias=use_bias
-            )
+            downconv = conv2d(input_nc, inner_nc, kernel_size=4, stride=1, padding=1, bias=use_bias)
             blurpool = antialiased_cnns.BlurPool(inner_nc, stride=2)
         else:
             antialiased = False
-            downconv = conv2d(
-                input_nc, inner_nc, kernel_size=4, stride=2, padding=1, bias=use_bias
-            )
+            downconv = conv2d(input_nc, inner_nc, kernel_size=4, stride=2, padding=1, bias=use_bias)
         downrelu = nn.LeakyReLU(0.2, True)
         downnorm = norm_layer(inner_nc)
         uprelu = nn.ReLU(True)
diff --git a/satflow/models/hub.py b/satflow/models/hub.py
deleted file mode 100644
index 48b56aa4..00000000
--- a/satflow/models/hub.py
+++ /dev/null
@@ -1,272 +0,0 @@
-"""
-Originally Taken from https://github.com/rwightman/pytorch-image-models/blob/acd6c687fd1c0507128f0ce091829b233c8560b9/timm/models/hub.py
-"""
-
-import json
-import logging
-import os
-from functools import partial
-from typing import Optional, Union
-
-import pytorch_lightning
-import torch
-
-try:
-    from torch.hub import get_dir
-except ImportError:
-    from torch.hub import _get_torch_home as get_dir
-
-from satflow import __version__
-
-try:
-    from huggingface_hub import cached_download, hf_hub_url
-
-    cached_download = partial(cached_download, library_name="satflow", library_version=__version__)
-except ImportError:
-    hf_hub_url = None
-    cached_download = None
-
-from huggingface_hub import CONFIG_NAME, PYTORCH_WEIGHTS_NAME, ModelHubMixin, hf_hub_download
-
-MODEL_CARD_MARKDOWN = """---
-license: mit
-tags:
-- satflow
-- forecasting
-- timeseries
-- remote-sensing
----
-
-# {model_name}
-
-## Model description
-
-[More information needed]
-
-## Intended uses & limitations
-
-[More information needed]
-
-## How to use
-
-[More information needed]
-
-## Limitations and bias
-
-[More information needed]
-
-## Training data
-
-[More information needed]
-
-## Training procedure
-
-[More information needed]
-
-## Evaluation results
-
-[More information needed]
-
-"""
-
-_logger = logging.getLogger(__name__)
-
-
-def get_cache_dir(child_dir=""):
-    """
-    Returns the location of the directory where models are cached (and creates it if necessary).
-    """
-    hub_dir = get_dir()
-    child_dir = () if not child_dir else (child_dir,)
-    model_dir = os.path.join(hub_dir, "checkpoints", *child_dir)
-    os.makedirs(model_dir, exist_ok=True)
-    return model_dir
-
-
-def has_hf_hub(necessary=False):
-    if hf_hub_url is None and necessary:
-        # if no HF Hub module installed and it is necessary to continue, raise error
-        raise RuntimeError(
-            "Hugging Face hub model specified but package not installed. Run `pip install huggingface_hub`."
-        )
-    return hf_hub_url is not None
-
-
-def hf_split(hf_id):
-    rev_split = hf_id.split("@")
-    assert (
-        0 < len(rev_split) <= 2
-    ), "hf_hub id should only contain one @ character to identify revision."
-    hf_model_id = rev_split[0]
-    hf_revision = rev_split[-1] if len(rev_split) > 1 else None
-    return hf_model_id, hf_revision
-
-
-def load_cfg_from_json(json_file: Union[str, os.PathLike]):
-    with open(json_file, "r", encoding="utf-8") as reader:
-        text = reader.read()
-    return json.loads(text)
-
-
-def _download_from_hf(model_id: str, filename: str):
-    hf_model_id, hf_revision = hf_split(model_id)
-    url = hf_hub_url(hf_model_id, filename, revision=hf_revision)
-    return cached_download(url, cache_dir=get_cache_dir("hf"))
-
-
-def load_model_config_from_hf(model_id: str):
-    assert has_hf_hub(True)
-    cached_file = _download_from_hf(model_id, "config.json")
-    default_cfg = load_cfg_from_json(cached_file)
-    default_cfg[
-        "hf_hub"
-    ] = model_id  # insert hf_hub id for pretrained weight load during model creation
-    model_name = default_cfg.get("architecture")
-    return default_cfg, model_name
-
-
-def load_state_dict_from_hf(model_id: str):
-    assert has_hf_hub(True)
-    cached_file = _download_from_hf(model_id, "pytorch_model.pth")
-    state_dict = torch.load(cached_file, map_location="cpu")
-    return state_dict
-
-
-def cache_file_from_hf(model_id: str):
-    assert has_hf_hub(True)
-    cached_file = _download_from_hf(model_id, "pytorch_model.pth")
-    return cached_file
-
-
-def load_pretrained(
-    model,
-    default_cfg: Optional[dict] = None,
-    in_chans: int = 12,
-    strict: bool = True,
-) -> Union[torch.nn.Module, pytorch_lightning.LightningModule]:
-    """Load pretrained checkpoint
-
-    Taken from https://github.com/rwightman/pytorch-image-models/blob/acd6c687fd1c0507128f0ce091829b233c8560b9/timm/models/helpers.py
-
-    Args:
-        model (nn.Module) : PyTorch model module, or LightningModule
-        default_cfg (Optional[Dict]): default configuration for pretrained weights / target dataset
-        in_chans (int): in_chans for model
-        strict (bool): strict load of checkpoint
-    """
-    is_lightning_module = issubclass(model, pytorch_lightning.LightningModule)
-    default_cfg = default_cfg or getattr(model, "default_cfg", None) or {}
-    pretrained_path = default_cfg.pop("checkpoint_path", None)
-    hf_hub_id = default_cfg.pop("hf_hub", None)
-    if in_chans != default_cfg.get("input_channels", None):
-        strict = False
-        _logger.warning(
-            f"Unable to convert pretrained weights because of mismatch in input channels, using random init for first layer."
-        )
-    if not is_lightning_module:
-        # The model is passed uninitialized, so if not having to do the PL thing, should initialize here
-        model = model(**default_cfg)
-    if not pretrained_path and not hf_hub_id:
-        _logger.warning("No pretrained weights exist for this model. Using random initialization.")
-        return model
-    if hf_hub_id and has_hf_hub(necessary=not pretrained_path):
-        _logger.info(f"Loading pretrained weights from Hugging Face hub ({hf_hub_id})")
-        if is_lightning_module:
-            checkpoint = cache_file_from_hf(hf_hub_id)
-            model.load_from_checkpoint(checkpoint, strict=strict, **default_cfg)
-            return model
-        state_dict = load_state_dict_from_hf(hf_hub_id)
-    else:
-        if is_lightning_module:
-            model.load_from_checkpoint(pretrained_path, strict=strict, **default_cfg)
-            return model
-        state_dict = torch.load(pretrained_path, map_location="cpu")
-
-    model.load_state_dict(state_dict, strict=strict)
-    return model
-
-
-class SatFlowModelHubMixin(ModelHubMixin):
-    def __init__(self, *args, **kwargs):
-        """
-        Mix this class with your pl.LightningModule class to easily push / download the model via the Hugging Face Hub
-
-        Example::
-
-            >>> from satflow.models.hub import SatFlowModelHubMixin
-
-            >>> class MyModel(nn.Module, SatFlowModelHubMixin):
-            ...    def __init__(self, **kwargs):
-            ...        super().__init__()
-            ...        self.layer = ...
-            ...    def forward(self, ...)
-            ...        return ...
-
-            >>> model = MyModel()
-            >>> model.push_to_hub("mymodel") # Pushing model-weights to hf-hub
-
-            >>> # Downloading weights from hf-hub & model will be initialized from those weights
-            >>> model = MyModel.from_pretrained("username/mymodel")
-        """
-
-    def _create_model_card(self, path):
-        model_card = MODEL_CARD_MARKDOWN.format(model_name=type(self).__name__)
-        with open(os.path.join(path, "README.md"), "w") as f:
-            f.write(model_card)
-
-    def _save_config(self, module, save_directory):
-        config = dict(module.hparams)
-        path = os.path.join(save_directory, CONFIG_NAME)
-        with open(path, "w") as f:
-            json.dump(config, f)
-
-    def _save_pretrained(self, save_directory: str, save_config: bool = True):
-        # Save model weights
-        path = os.path.join(save_directory, PYTORCH_WEIGHTS_NAME)
-        model_to_save = self.module if hasattr(self, "module") else self
-        torch.save(model_to_save.state_dict(), path)
-        # Save model config
-        if save_config and model_to_save.hparams:
-            self._save_config(model_to_save, save_directory)
-        # Save model card
-        self._create_model_card(save_directory)
-
-    @classmethod
-    def _from_pretrained(
-        cls,
-        model_id,
-        revision,
-        cache_dir,
-        force_download,
-        proxies,
-        resume_download,
-        local_files_only,
-        use_auth_token,
-        map_location="cpu",
-        strict=False,
-        **model_kwargs,
-    ):
-        map_location = torch.device(map_location)
-
-        if os.path.isdir(model_id):
-            print("Loading weights from local directory")
-            model_file = os.path.join(model_id, PYTORCH_WEIGHTS_NAME)
-        else:
-            model_file = hf_hub_download(
-                repo_id=model_id,
-                filename=PYTORCH_WEIGHTS_NAME,
-                revision=revision,
-                cache_dir=cache_dir,
-                force_download=force_download,
-                proxies=proxies,
-                resume_download=resume_download,
-                use_auth_token=use_auth_token,
-                local_files_only=local_files_only,
-            )
-        model = cls(**model_kwargs["config"])
-
-        state_dict = torch.load(model_file, map_location=map_location)
-        model.load_state_dict(state_dict, strict=strict)
-        model.eval()
-
-        return model
diff --git a/satflow/models/layers/Attention.py b/satflow/models/layers/Attention.py
index 500899b2..49b10aa8 100644
--- a/satflow/models/layers/Attention.py
+++ b/satflow/models/layers/Attention.py
@@ -10,15 +10,9 @@ def __init__(
     ):
         super().__init__()
         self.model = nn.Sequential(
-            SeparableAttnCell(
-                in_dim, "T", activation, pooling_factor, padding_mode, padding_value
-            ),
-            SeparableAttnCell(
-                in_dim, "W", activation, pooling_factor, padding_mode, padding_value
-            ),
-            SeparableAttnCell(
-                in_dim, "H", activation, pooling_factor, padding_mode, padding_value
-            ),
+            SeparableAttnCell(in_dim, "T", activation, pooling_factor, padding_mode, padding_value),
+            SeparableAttnCell(in_dim, "W", activation, pooling_factor, padding_mode, padding_value),
+            SeparableAttnCell(in_dim, "H", activation, pooling_factor, padding_mode, padding_value),
         )
 
     def forward(self, x):
diff --git a/satflow/models/layers/ConditionTime.py b/satflow/models/layers/ConditionTime.py
index 29a96cad..97695ae6 100644
--- a/satflow/models/layers/ConditionTime.py
+++ b/satflow/models/layers/ConditionTime.py
@@ -23,9 +23,7 @@ def forward(self, x, fstep=0):
         "x stack of images, fsteps"
         if self.num_dims == 5:
             bs, seq_len, ch, h, w = x.shape
-            ct = condition_time(x, fstep, (h, w), seq_len=self.horizon).repeat(
-                bs, seq_len, 1, 1, 1
-            )
+            ct = condition_time(x, fstep, (h, w), seq_len=self.horizon).repeat(bs, seq_len, 1, 1, 1)
         else:
             bs, h, w, ch = x.shape
             ct = condition_time(x, fstep, (h, w), seq_len=self.horizon).repeat(bs, 1, 1, 1)
diff --git a/satflow/models/layers/ConvLSTM.py b/satflow/models/layers/ConvLSTM.py
index 16a70ab0..39e204bf 100644
--- a/satflow/models/layers/ConvLSTM.py
+++ b/satflow/models/layers/ConvLSTM.py
@@ -58,10 +58,6 @@ def forward(self, input_tensor, cur_state):
     def init_hidden(self, batch_size, image_size):
         height, width = image_size
         return (
-            torch.zeros(
-                batch_size, self.hidden_dim, height, width, device=self.conv.weight.device
-            ),
-            torch.zeros(
-                batch_size, self.hidden_dim, height, width, device=self.conv.weight.device
-            ),
+            torch.zeros(batch_size, self.hidden_dim, height, width, device=self.conv.weight.device),
+            torch.zeros(batch_size, self.hidden_dim, height, width, device=self.conv.weight.device),
         )
diff --git a/satflow/models/layers/Generator.py b/satflow/models/layers/Generator.py
index 3a010294..510ce2a6 100644
--- a/satflow/models/layers/Generator.py
+++ b/satflow/models/layers/Generator.py
@@ -84,9 +84,7 @@ def forward(self, x, class_id):
                 torch.cat((noise_emb[0], class_emb), dim=1)
             )  # B x (2 x ld x ch)
         else:
-            y = self.affine_transfrom(
-                torch.cat((noise_emb, class_emb), dim=1)
-            )  # B x (2 x ld x ch)
+            y = self.affine_transfrom(torch.cat((noise_emb, class_emb), dim=1))  # B x (2 x ld x ch)
 
         y = y.view(-1, 8 * self.ch, self.latent_dim, self.latent_dim)  # B x ch x ld x ld
 
diff --git a/satflow/models/losses.py b/satflow/models/losses.py
deleted file mode 100644
index 07171552..00000000
--- a/satflow/models/losses.py
+++ /dev/null
@@ -1,298 +0,0 @@
-import numpy as np
-import torch
-import torch.nn as nn
-from pytorch_msssim import SSIM, MS_SSIM
-from torch.nn import functional as F
-
-
-class SSIMLoss(nn.Module):
-    def __init__(self, convert_range: bool = False, **kwargs):
-        """
-        SSIM Loss, optionally converting input range from [-1,1] to [0,1]
-        Args:
-            convert_range:
-            **kwargs:
-        """
-        super(SSIMLoss, self).__init__()
-        self.convert_range = convert_range
-        self.ssim_module = SSIM(**kwargs)
-
-    def forward(self, x: torch.Tensor, y: torch.Tensor):
-        if self.convert_range:
-            x = torch.div(torch.add(x, 1), 2)
-            y = torch.div(torch.add(y, 1), 2)
-        return 1.0 - self.ssim_module(x, y)
-
-
-class MS_SSIMLoss(nn.Module):
-    def __init__(self, convert_range: bool = False, **kwargs):
-        """
-        Multi-Scale SSIM Loss, optionally converting input range from [-1,1] to [0,1]
-        Args:
-            convert_range:
-            **kwargs:
-        """
-        super(MS_SSIMLoss, self).__init__()
-        self.convert_range = convert_range
-        self.ssim_module = MS_SSIM(**kwargs)
-
-    def forward(self, x: torch.Tensor, y: torch.Tensor):
-        if self.convert_range:
-            x = torch.div(torch.add(x, 1), 2)
-            y = torch.div(torch.add(y, 1), 2)
-        return 1.0 - self.ssim_module(x, y)
-
-
-class SSIMLossDynamic(nn.Module):
-    def __init__(self, convert_range: bool = False, **kwargs):
-        """
-        SSIM Loss on only dynamic part of the images, optionally converting input range from [-1,1] to [0,1]
-
-        In Mathieu et al. to stop SSIM regressing towards the mean and predicting only the background, they only
-        run SSIM on the dynamic parts of the image. We can accomplish that by subtracting the current image from the future ones
-
-        Args:
-            convert_range:
-            **kwargs:
-        """
-        super(SSIMLossDynamic, self).__init__()
-        self.convert_range = convert_range
-        self.ssim_module = MS_SSIM(**kwargs)
-
-    def forward(self, curr_image: torch.Tensor, x: torch.Tensor, y: torch.Tensor):
-        if self.convert_range:
-            curr_image = torch.div(torch.add(curr_image, 1), 2)
-            x = torch.div(torch.add(x, 1), 2)
-            y = torch.div(torch.add(y, 1), 2)
-        # Subtract 'now' image to get what changes for both x and y
-        x = x - curr_image
-        y = y - curr_image
-        # TODO: Mask out loss from pixels that don't change
-        return 1.0 - self.ssim_module(x, y)
-
-
-def tv_loss(img, tv_weight):
-    """
-    Taken from https://github.com/chongyangma/cs231n/blob/master/assignments/assignment3/style_transfer_pytorch.py
-    Compute total variation loss.
-    Inputs:
-    - img: PyTorch Variable of shape (1, 3, H, W) holding an input image.
-    - tv_weight: Scalar giving the weight w_t to use for the TV loss.
-    Returns:
-    - loss: PyTorch Variable holding a scalar giving the total variation loss
-      for img weighted by tv_weight.
-    """
-    w_variance = torch.sum(torch.pow(img[:, :, :, :-1] - img[:, :, :, 1:], 2))
-    h_variance = torch.sum(torch.pow(img[:, :, :-1, :] - img[:, :, 1:, :], 2))
-    loss = tv_weight * (h_variance + w_variance)
-    return loss
-
-
-class TotalVariationLoss(nn.Module):
-    def __init__(self, tv_weight: float = 1.0):
-        super(TotalVariationLoss, self).__init__()
-        self.tv_weight = tv_weight
-
-    def forward(self, x: torch.Tensor):
-        return tv_loss(x, self.tv_weight)
-
-
-class GradientDifferenceLoss(nn.Module):
-    """"""
-
-    def __init__(self, alpha: int = 2):
-        super(GradientDifferenceLoss, self).__init__()
-        self.alpha = alpha
-
-    def forward(self, x: torch.Tensor, y: torch.Tensor):
-        t1 = torch.pow(
-            torch.abs(
-                torch.abs(x[:, :, :, 1:, :] - x[:, :, :, :-1, :])
-                - torch.abs(y[:, :, :, 1:, :] - y[:, :, :, :-1, :])
-            ),
-            self.alpha,
-        )
-        t2 = torch.pow(
-            torch.abs(
-                torch.abs(x[:, :, :, :, :-1] - x[:, :, :, :, 1:])
-                - torch.abs(y[:, :, :, :, :-1] - y[:, :, :, :, 1:])
-            ),
-            self.alpha,
-        )
-        loss = t1 + t2
-        print(loss.shape)
-        return loss
-
-
-class GridCellLoss(nn.Module):
-    """
-    Grid Cell Regularizer loss from Skillful Nowcasting, see https://arxiv.org/pdf/2104.00954.pdf
-
-    """
-
-    def __init__(self, weight_fn=None):
-        super().__init__()
-        self.weight_fn = weight_fn  # In Paper, weight_fn is max(y+1,24)
-
-    def forward(self, generated_images, targets):
-        """
-        Calculates the grid cell regularizer value, assumes generated images are the mean predictions from
-        6 calls to the generater (Monte Carlo estimation of the expectations for the latent variable)
-        Args:
-            generated_images: Mean generated images from the generator
-            targets: Ground truth future frames
-
-        Returns:
-            Grid Cell Regularizer term
-        """
-        difference = generated_images - targets
-        if self.weight_fn is not None:
-            difference *= self.weight_fn(targets)
-        difference /= targets.size(1) * targets.size(3) * targets.size(4)  # 1/HWN
-        return difference.mean()
-
-
-class NowcastingLoss(nn.Module):
-    def __init__(self):
-        super().__init__()
-
-    def forward(self, x, real_flag):
-        if real_flag is True:
-            x = -x
-        return F.relu(1.0 + x).mean()
-
-
-class FocalLoss(nn.Module):
-    """
-    copy from: https://github.com/Hsuxu/Loss_ToolBox-PyTorch/blob/master/FocalLoss/FocalLoss.py
-    This is a implementation of Focal Loss with smooth label cross entropy supported which is proposed in
-    'Focal Loss for Dense Object Detection. (https://arxiv.org/abs/1708.02002)'
-        Focal_Loss= -1*alpha*(1-pt)*log(pt)
-    :param num_class:
-    :param alpha: (tensor) 3D or 4D the scalar factor for this criterion
-    :param gamma: (float,double) gamma > 0 reduces the relative loss for well-classified examples (p>0.5) putting more
-                    focus on hard misclassified example
-    :param smooth: (float,double) smooth value when cross entropy
-    :param balance_index: (int) balance class index, should be specific when alpha is float
-    :param size_average: (bool, optional) By default, the losses are averaged over each loss element in the batch.
-    """
-
-    def __init__(
-        self,
-        apply_nonlin=None,
-        alpha=None,
-        gamma=2,
-        balance_index=0,
-        smooth=1e-5,
-        size_average=True,
-    ):
-        super(FocalLoss, self).__init__()
-        self.apply_nonlin = apply_nonlin
-        self.alpha = alpha
-        self.gamma = gamma
-        self.balance_index = balance_index
-        self.smooth = smooth
-        self.size_average = size_average
-
-        if self.smooth is not None:
-            if self.smooth < 0 or self.smooth > 1.0:
-                raise ValueError("smooth value should be in [0,1]")
-
-    def forward(self, logit, target):
-        if self.apply_nonlin is not None:
-            logit = self.apply_nonlin(logit)
-        num_class = logit.shape[1]
-
-        if logit.dim() > 2:
-            # N,C,d1,d2 -> N,C,m (m=d1*d2*...)
-            logit = logit.view(logit.size(0), logit.size(1), -1)
-            logit = logit.permute(0, 2, 1).contiguous()
-            logit = logit.view(-1, logit.size(-1))
-        target = torch.squeeze(target, 1)
-        target = target.view(-1, 1)
-        #
-        alpha = self.alpha
-
-        if alpha is None:
-            alpha = torch.ones(num_class, 1)
-        elif isinstance(alpha, (list, np.ndarray)):
-            assert len(alpha) == num_class
-            alpha = torch.FloatTensor(alpha).view(num_class, 1)
-            alpha = alpha / alpha.sum()
-        elif isinstance(alpha, float):
-            alpha = torch.ones(num_class, 1)
-            alpha = alpha * (1 - self.alpha)
-            alpha[self.balance_index] = self.alpha
-
-        else:
-            raise TypeError("Not support alpha type")
-
-        if alpha.device != logit.device:
-            alpha = alpha.to(logit.device)
-
-        idx = target.cpu().long()
-
-        one_hot_key = torch.FloatTensor(target.size(0), num_class).zero_()
-        one_hot_key = one_hot_key.scatter_(1, idx, 1)
-        if one_hot_key.device != logit.device:
-            one_hot_key = one_hot_key.to(logit.device)
-
-        if self.smooth:
-            one_hot_key = torch.clamp(
-                one_hot_key, self.smooth / (num_class - 1), 1.0 - self.smooth
-            )
-        pt = (one_hot_key * logit).sum(1) + self.smooth
-        logpt = pt.log()
-
-        gamma = self.gamma
-
-        alpha = alpha[idx]
-        alpha = torch.squeeze(alpha)
-        loss = -1 * alpha * torch.pow((1 - pt), gamma) * logpt
-
-        if self.size_average:
-            loss = loss.mean()
-        else:
-            loss = loss.sum()
-        return loss
-
-
-def get_loss(loss: str = "mse", **kwargs) -> torch.nn.Module:
-    if isinstance(loss, torch.nn.Module):
-        return loss
-    assert loss in [
-        "mse",
-        "bce",
-        "binary_crossentropy",
-        "crossentropy",
-        "focal",
-        "ssim",
-        "ms_ssim",
-        "l1",
-        "tv",
-        "total_variation",
-        "ssim_dynamic",
-        "gdl",
-        "gradient_difference_loss",
-    ]
-    if loss == "mse":
-        criterion = F.mse_loss
-    elif loss in ["bce", "binary_crossentropy", "crossentropy"]:
-        criterion = F.nll_loss
-    elif loss in ["focal"]:
-        criterion = FocalLoss()
-    elif loss in ["ssim"]:
-        criterion = SSIMLoss(data_range=1.0, size_average=True, **kwargs)
-    elif loss in ["ms_ssim"]:
-        criterion = MS_SSIMLoss(data_range=1.0, size_average=True, **kwargs)
-    elif loss in ["ssim_dynamic"]:
-        criterion = SSIMLossDynamic(data_range=1.0, size_average=True, **kwargs)
-    elif loss in ["l1"]:
-        criterion = torch.nn.L1Loss()
-    elif loss in ["tv", "total_variation"]:
-        criterion = TotalVariationLoss(tv_weight=kwargs.get("tv_weight", 1))
-    elif loss in ["gdl", "gradient_difference_loss"]:
-        criterion = GradientDifferenceLoss(alpha=kwargs.get("alpha", 2))
-    else:
-        raise ValueError(f"loss {loss} not recognized")
-    return criterion
diff --git a/satflow/models/perceiver.py b/satflow/models/perceiver.py
index 7740f188..129dc301 100644
--- a/satflow/models/perceiver.py
+++ b/satflow/models/perceiver.py
@@ -7,10 +7,10 @@
 from math import prod
 from torch.distributions import uniform
 from typing import Iterable, Dict, Optional, Any, Union, Tuple
-from satflow.models.base import register_model, BaseModel
+from nowcasting_utils.models.base import register_model, BaseModel
 from einops import rearrange, repeat
 from pl_bolts.optimizers.lr_scheduler import LinearWarmupCosineAnnealingLR
-from satflow.models.losses import get_loss
+from nowcasting_utils.models.loss import get_loss
 import torch_optimizer as optim
 import logging
 
diff --git a/satflow/models/pix2pix.py b/satflow/models/pix2pix.py
index 5b4ba36b..56e95cbb 100644
--- a/satflow/models/pix2pix.py
+++ b/satflow/models/pix2pix.py
@@ -4,7 +4,7 @@
 from collections import OrderedDict
 from torch.optim import lr_scheduler
 import pytorch_lightning as pl
-from satflow.models.base import register_model
+from nowcasting_utils.models.base import register_model
 from satflow.models.gan.discriminators import define_discriminator, GANLoss
 from satflow.models.gan import define_generator
 
diff --git a/satflow/models/pixel_cnn.py b/satflow/models/pixel_cnn.py
index 3e328d30..4583fd4d 100644
--- a/satflow/models/pixel_cnn.py
+++ b/satflow/models/pixel_cnn.py
@@ -1,7 +1,7 @@
 import torch
 import torch.nn.functional as F
 import pytorch_lightning as pl
-from satflow.models.base import register_model
+from nowcasting_utils.models.base import register_model
 from pl_bolts.models.vision import PixelCNN as Pixcnn
 
 
diff --git a/satflow/models/pl_metnet.py b/satflow/models/pl_metnet.py
index e0006d58..096db1db 100644
--- a/satflow/models/pl_metnet.py
+++ b/satflow/models/pl_metnet.py
@@ -2,8 +2,8 @@
 import torch.nn as nn
 from typing import Any
 
-from satflow.models.base import register_model, BaseModel
-from satflow.models.losses import get_loss
+from nowcasting_utils.models.base import register_model, BaseModel
+from nowcasting_utils.models.loss import get_loss
 from pl_bolts.optimizers.lr_scheduler import LinearWarmupCosineAnnealingLR
 from metnet import MetNet
 
diff --git a/satflow/models/runet.py b/satflow/models/runet.py
index f805bb51..81a70286 100644
--- a/satflow/models/runet.py
+++ b/satflow/models/runet.py
@@ -3,8 +3,8 @@
 import pytorch_lightning as pl
 import torchvision
 from typing import Union
-from satflow.models.base import register_model
-from satflow.models.losses import get_loss
+from nowcasting_utils.models.base import register_model
+from nowcasting_utils.models.loss import get_loss
 import numpy as np
 
 
@@ -122,19 +122,13 @@ def __init__(self, img_ch=3, output_ch=1, t=2, conv_type: str = "standard"):
         self.RRCNN1 = RRCNN_block(ch_in=img_ch, ch_out=64, t=t, conv_type=conv_type)
         self.Blur1 = antialiased_cnns.BlurPool(64, stride=2) if self.antialiased else nn.Identity()
         self.RRCNN2 = RRCNN_block(ch_in=64, ch_out=128, t=t, conv_type=conv_type)
-        self.Blur2 = (
-            antialiased_cnns.BlurPool(128, stride=2) if self.antialiased else nn.Identity()
-        )
+        self.Blur2 = antialiased_cnns.BlurPool(128, stride=2) if self.antialiased else nn.Identity()
 
         self.RRCNN3 = RRCNN_block(ch_in=128, ch_out=256, t=t, conv_type=conv_type)
-        self.Blur3 = (
-            antialiased_cnns.BlurPool(256, stride=2) if self.antialiased else nn.Identity()
-        )
+        self.Blur3 = antialiased_cnns.BlurPool(256, stride=2) if self.antialiased else nn.Identity()
 
         self.RRCNN4 = RRCNN_block(ch_in=256, ch_out=512, t=t, conv_type=conv_type)
-        self.Blur4 = (
-            antialiased_cnns.BlurPool(512, stride=2) if self.antialiased else nn.Identity()
-        )
+        self.Blur4 = antialiased_cnns.BlurPool(512, stride=2) if self.antialiased else nn.Identity()
 
         self.RRCNN5 = RRCNN_block(ch_in=512, ch_out=1024, t=t, conv_type=conv_type)
 
diff --git a/satflow/models/timegan.py b/satflow/models/timegan.py
deleted file mode 100644
index 01628808..00000000
--- a/satflow/models/timegan.py
+++ /dev/null
@@ -1,610 +0,0 @@
-# -*- coding: UTF-8 -*-
-# Taken from https://github.com/d9n13lt4n/timegan-pytorch
-import torch
-import pytorch_lightning as pl
-
-
-class TimeGAN(pl.LightningModule):
-    def __init__(self):
-        super().__init__()
-
-    def forward(self, x):
-        return x
-
-    def training_step(self, batch, batch_idx):
-        return
-
-    def validation_step(self, batch, batch_idx):
-        return
-
-    def visualze(self):
-        return
-
-    def configure_optimizers(self):
-        return
-
-
-class EmbeddingNetwork(torch.nn.Module):
-    """The embedding network (encoder) for TimeGAN"""
-
-    def __init__(self, args):
-        super(EmbeddingNetwork, self).__init__()
-        self.feature_dim = args.feature_dim
-        self.hidden_dim = args.hidden_dim
-        self.num_layers = args.num_layers
-        self.padding_value = args.padding_value
-        self.max_seq_len = args.max_seq_len
-
-        # Embedder Architecture
-        self.emb_rnn = torch.nn.GRU(
-            input_size=self.feature_dim,
-            hidden_size=self.hidden_dim,
-            num_layers=self.num_layers,
-            batch_first=True,
-        )
-        self.emb_linear = torch.nn.Linear(self.hidden_dim, self.hidden_dim)
-        self.emb_sigmoid = torch.nn.Sigmoid()
-
-        # Init weights
-        # Default weights of TensorFlow is Xavier Uniform for W and 1 or 0 for b
-        # Reference:
-        # - https://www.tensorflow.org/api_docs/python/tf/compat/v1/get_variable
-        # - https://github.com/tensorflow/tensorflow/blob/v2.3.1/tensorflow/python/keras/layers/legacy_rnn/rnn_cell_impl.py#L484-L614
-        with torch.no_grad():
-            for name, param in self.emb_rnn.named_parameters():
-                if "weight_ih" in name:
-                    torch.nn.init.xavier_uniform_(param.data)
-                elif "weight_hh" in name:
-                    torch.nn.init.xavier_uniform_(param.data)
-                elif "bias_ih" in name:
-                    param.data.fill_(1)
-                elif "bias_hh" in name:
-                    param.data.fill_(0)
-            for name, param in self.emb_linear.named_parameters():
-                if "weight" in name:
-                    torch.nn.init.xavier_uniform_(param)
-                elif "bias" in name:
-                    param.data.fill_(0)
-
-    def forward(self, X, T):
-        """Forward pass for embedding features from original space into latent space
-        Args:
-            - X: input time-series features (B x S x F)
-            - T: input temporal information (B)
-        Returns:
-            - H: latent space embeddings (B x S x H)
-        """
-        # Dynamic RNN input for ignoring paddings
-        X_packed = torch.nn.utils.rnn.pack_padded_sequence(
-            input=X, lengths=T, batch_first=True, enforce_sorted=False
-        )
-
-        # 128 x 100 x 71
-        H_o, H_t = self.emb_rnn(X_packed)
-
-        # Pad RNN output back to sequence length
-        H_o, T = torch.nn.utils.rnn.pad_packed_sequence(
-            sequence=H_o,
-            batch_first=True,
-            padding_value=self.padding_value,
-            total_length=self.max_seq_len,
-        )
-
-        # 128 x 100 x 10
-        logits = self.emb_linear(H_o)
-        # 128 x 100 x 10
-        H = self.emb_sigmoid(logits)
-        return H
-
-
-class RecoveryNetwork(torch.nn.Module):
-    """The recovery network (decoder) for TimeGAN"""
-
-    def __init__(self, args):
-        super(RecoveryNetwork, self).__init__()
-        self.hidden_dim = args.hidden_dim
-        self.feature_dim = args.feature_dim
-        self.num_layers = args.num_layers
-        self.padding_value = args.padding_value
-        self.max_seq_len = args.max_seq_len
-
-        # Recovery Architecture
-        self.rec_rnn = torch.nn.GRU(
-            input_size=self.hidden_dim,
-            hidden_size=self.hidden_dim,
-            num_layers=self.num_layers,
-            batch_first=True,
-        )
-        self.rec_linear = torch.nn.Linear(self.hidden_dim, self.feature_dim)
-
-        # Init weights
-        # Default weights of TensorFlow is Xavier Uniform for W and 1 or 0 for b
-        # Reference:
-        # - https://www.tensorflow.org/api_docs/python/tf/compat/v1/get_variable
-        # - https://github.com/tensorflow/tensorflow/blob/v2.3.1/tensorflow/python/keras/layers/legacy_rnn/rnn_cell_impl.py#L484-L614
-        with torch.no_grad():
-            for name, param in self.rec_rnn.named_parameters():
-                if "weight_ih" in name:
-                    torch.nn.init.xavier_uniform_(param.data)
-                elif "weight_hh" in name:
-                    torch.nn.init.xavier_uniform_(param.data)
-                elif "bias_ih" in name:
-                    param.data.fill_(1)
-                elif "bias_hh" in name:
-                    param.data.fill_(0)
-            for name, param in self.rec_linear.named_parameters():
-                if "weight" in name:
-                    torch.nn.init.xavier_uniform_(param)
-                elif "bias" in name:
-                    param.data.fill_(0)
-
-    def forward(self, H, T):
-        """Forward pass for the recovering features from latent space to original space
-        Args:
-            - H: latent representation (B x S x E)
-            - T: input temporal information (B)
-        Returns:
-            - X_tilde: recovered data (B x S x F)
-        """
-        # Dynamic RNN input for ignoring paddings
-        H_packed = torch.nn.utils.rnn.pack_padded_sequence(
-            input=H, lengths=T, batch_first=True, enforce_sorted=False
-        )
-
-        # 128 x 100 x 10
-        H_o, H_t = self.rec_rnn(H_packed)
-
-        # Pad RNN output back to sequence length
-        H_o, T = torch.nn.utils.rnn.pad_packed_sequence(
-            sequence=H_o,
-            batch_first=True,
-            padding_value=self.padding_value,
-            total_length=self.max_seq_len,
-        )
-
-        # 128 x 100 x 71
-        X_tilde = self.rec_linear(H_o)
-        return X_tilde
-
-
-class SupervisorNetwork(torch.nn.Module):
-    """The Supervisor network (decoder) for TimeGAN"""
-
-    def __init__(self, args):
-        super(SupervisorNetwork, self).__init__()
-        self.hidden_dim = args.hidden_dim
-        self.num_layers = args.num_layers
-        self.padding_value = args.padding_value
-        self.max_seq_len = args.max_seq_len
-
-        # Supervisor Architecture
-        self.sup_rnn = torch.nn.GRU(
-            input_size=self.hidden_dim,
-            hidden_size=self.hidden_dim,
-            num_layers=self.num_layers - 1,
-            batch_first=True,
-        )
-        self.sup_linear = torch.nn.Linear(self.hidden_dim, self.hidden_dim)
-        self.sup_sigmoid = torch.nn.Sigmoid()
-
-        # Init weights
-        # Default weights of TensorFlow is Xavier Uniform for W and 1 or 0 for b
-        # Reference:
-        # - https://www.tensorflow.org/api_docs/python/tf/compat/v1/get_variable
-        # - https://github.com/tensorflow/tensorflow/blob/v2.3.1/tensorflow/python/keras/layers/legacy_rnn/rnn_cell_impl.py#L484-L614
-        with torch.no_grad():
-            for name, param in self.sup_rnn.named_parameters():
-                if "weight_ih" in name:
-                    torch.nn.init.xavier_uniform_(param.data)
-                elif "weight_hh" in name:
-                    torch.nn.init.xavier_uniform_(param.data)
-                elif "bias_ih" in name:
-                    param.data.fill_(1)
-                elif "bias_hh" in name:
-                    param.data.fill_(0)
-            for name, param in self.sup_linear.named_parameters():
-                if "weight" in name:
-                    torch.nn.init.xavier_uniform_(param)
-                elif "bias" in name:
-                    param.data.fill_(0)
-
-    def forward(self, H, T):
-        """Forward pass for the supervisor for predicting next step
-        Args:
-            - H: latent representation (B x S x E)
-            - T: input temporal information (B)
-        Returns:
-            - H_hat: predicted next step data (B x S x E)
-        """
-        # Dynamic RNN input for ignoring paddings
-        H_packed = torch.nn.utils.rnn.pack_padded_sequence(
-            input=H, lengths=T, batch_first=True, enforce_sorted=False
-        )
-
-        # 128 x 100 x 10
-        H_o, H_t = self.sup_rnn(H_packed)
-
-        # Pad RNN output back to sequence length
-        H_o, T = torch.nn.utils.rnn.pad_packed_sequence(
-            sequence=H_o,
-            batch_first=True,
-            padding_value=self.padding_value,
-            total_length=self.max_seq_len,
-        )
-
-        # 128 x 100 x 10
-        logits = self.sup_linear(H_o)
-        # 128 x 100 x 10
-        H_hat = self.sup_sigmoid(logits)
-        return H_hat
-
-
-class GeneratorNetwork(torch.nn.Module):
-    """The generator network (encoder) for TimeGAN"""
-
-    def __init__(self, args):
-        super(GeneratorNetwork, self).__init__()
-        self.Z_dim = args.Z_dim
-        self.hidden_dim = args.hidden_dim
-        self.num_layers = args.num_layers
-        self.padding_value = args.padding_value
-        self.max_seq_len = args.max_seq_len
-
-        # Generator Architecture
-        self.gen_rnn = torch.nn.GRU(
-            input_size=self.Z_dim,
-            hidden_size=self.hidden_dim,
-            num_layers=self.num_layers,
-            batch_first=True,
-        )
-        self.gen_linear = torch.nn.Linear(self.hidden_dim, self.hidden_dim)
-        self.gen_sigmoid = torch.nn.Sigmoid()
-
-        # Init weights
-        # Default weights of TensorFlow is Xavier Uniform for W and 1 or 0 for b
-        # Reference:
-        # - https://www.tensorflow.org/api_docs/python/tf/compat/v1/get_variable
-        # - https://github.com/tensorflow/tensorflow/blob/v2.3.1/tensorflow/python/keras/layers/legacy_rnn/rnn_cell_impl.py#L484-L614
-        with torch.no_grad():
-            for name, param in self.gen_rnn.named_parameters():
-                if "weight_ih" in name:
-                    torch.nn.init.xavier_uniform_(param.data)
-                elif "weight_hh" in name:
-                    torch.nn.init.xavier_uniform_(param.data)
-                elif "bias_ih" in name:
-                    param.data.fill_(1)
-                elif "bias_hh" in name:
-                    param.data.fill_(0)
-            for name, param in self.gen_linear.named_parameters():
-                if "weight" in name:
-                    torch.nn.init.xavier_uniform_(param)
-                elif "bias" in name:
-                    param.data.fill_(0)
-
-    def forward(self, Z, T):
-        """Takes in random noise (features) and generates synthetic features within the latent space
-        Args:
-            - Z: input random noise (B x S x Z)
-            - T: input temporal information
-        Returns:
-            - H: embeddings (B x S x E)
-        """
-        # Dynamic RNN input for ignoring paddings
-        Z_packed = torch.nn.utils.rnn.pack_padded_sequence(
-            input=Z, lengths=T, batch_first=True, enforce_sorted=False
-        )
-
-        # 128 x 100 x 71
-        H_o, H_t = self.gen_rnn(Z_packed)
-
-        # Pad RNN output back to sequence length
-        H_o, T = torch.nn.utils.rnn.pad_packed_sequence(
-            sequence=H_o,
-            batch_first=True,
-            padding_value=self.padding_value,
-            total_length=self.max_seq_len,
-        )
-
-        # 128 x 100 x 10
-        logits = self.gen_linear(H_o)
-        # B x S
-        H = self.gen_sigmoid(logits)
-        return H
-
-
-class DiscriminatorNetwork(torch.nn.Module):
-    """The Discriminator network (decoder) for TimeGAN"""
-
-    def __init__(self, args):
-        super(DiscriminatorNetwork, self).__init__()
-        self.hidden_dim = args.hidden_dim
-        self.num_layers = args.num_layers
-        self.padding_value = args.padding_value
-        self.max_seq_len = args.max_seq_len
-
-        # Discriminator Architecture
-        self.dis_rnn = torch.nn.GRU(
-            input_size=self.hidden_dim,
-            hidden_size=self.hidden_dim,
-            num_layers=self.num_layers,
-            batch_first=True,
-        )
-        self.dis_linear = torch.nn.Linear(self.hidden_dim, 1)
-
-        # Init weights
-        # Default weights of TensorFlow is Xavier Uniform for W and 1 or 0 for b
-        # Reference:
-        # - https://www.tensorflow.org/api_docs/python/tf/compat/v1/get_variable
-        # - https://github.com/tensorflow/tensorflow/blob/v2.3.1/tensorflow/python/keras/layers/legacy_rnn/rnn_cell_impl.py#L484-L614
-        with torch.no_grad():
-            for name, param in self.dis_rnn.named_parameters():
-                if "weight_ih" in name:
-                    torch.nn.init.xavier_uniform_(param.data)
-                elif "weight_hh" in name:
-                    torch.nn.init.xavier_uniform_(param.data)
-                elif "bias_ih" in name:
-                    param.data.fill_(1)
-                elif "bias_hh" in name:
-                    param.data.fill_(0)
-            for name, param in self.dis_linear.named_parameters():
-                if "weight" in name:
-                    torch.nn.init.xavier_uniform_(param)
-                elif "bias" in name:
-                    param.data.fill_(0)
-
-    def forward(self, H, T):
-        """Forward pass for predicting if the data is real or synthetic
-        Args:
-            - H: latent representation (B x S x E)
-            - T: input temporal information
-        Returns:
-            - logits: predicted logits (B x S x 1)
-        """
-        # Dynamic RNN input for ignoring paddings
-        H_packed = torch.nn.utils.rnn.pack_padded_sequence(
-            input=H, lengths=T, batch_first=True, enforce_sorted=False
-        )
-
-        # 128 x 100 x 10
-        H_o, H_t = self.dis_rnn(H_packed)
-
-        # Pad RNN output back to sequence length
-        H_o, T = torch.nn.utils.rnn.pad_packed_sequence(
-            sequence=H_o,
-            batch_first=True,
-            padding_value=self.padding_value,
-            total_length=self.max_seq_len,
-        )
-
-        # 128 x 100
-        logits = self.dis_linear(H_o).squeeze(-1)
-        return logits
-
-
-class TimeGANModel(torch.nn.Module):
-    """Implementation of TimeGAN (Yoon et al., 2019) using PyTorch
-    Reference:
-    - https://papers.nips.cc/paper/2019/hash/c9efe5f26cd17ba6216bbe2a7d26d490-Abstract.html
-    - https://github.com/jsyoon0823/TimeGAN
-    """
-
-    def __init__(self, args):
-        super(TimeGANModel, self).__init__()
-        self.device = args.device
-        self.feature_dim = args.feature_dim
-        self.Z_dim = args.Z_dim
-        self.hidden_dim = args.hidden_dim
-        self.max_seq_len = args.max_seq_len
-        self.batch_size = args.batch_size
-
-        self.embedder = EmbeddingNetwork(args)
-        self.recovery = RecoveryNetwork(args)
-        self.generator = GeneratorNetwork(args)
-        self.supervisor = SupervisorNetwork(args)
-        self.discriminator = DiscriminatorNetwork(args)
-
-    def _recovery_forward(self, X, T):
-        """The embedding network forward pass and the embedder network loss
-        Args:
-            - X: the original input features
-            - T: the temporal information
-        Returns:
-            - E_loss: the reconstruction loss
-            - X_tilde: the reconstructed features
-        """
-        # Forward Pass
-        H = self.embedder(X, T)
-        X_tilde = self.recovery(H, T)
-
-        # For Joint training
-        H_hat_supervise = self.supervisor(H, T)
-        G_loss_S = torch.nn.functional.mse_loss(
-            H_hat_supervise[:, :-1, :], H[:, 1:, :]
-        )  # Teacher forcing next output
-
-        # Reconstruction Loss
-        E_loss_T0 = torch.nn.functional.mse_loss(X_tilde, X)
-        E_loss0 = 10 * torch.sqrt(E_loss_T0)
-        E_loss = E_loss0 + 0.1 * G_loss_S
-        return E_loss, E_loss0, E_loss_T0
-
-    def _supervisor_forward(self, X, T):
-        """The supervisor training forward pass
-        Args:
-            - X: the original feature input
-        Returns:
-            - S_loss: the supervisor's loss
-        """
-        # Supervision Forward Pass
-        H = self.embedder(X, T)
-        H_hat_supervise = self.supervisor(H, T)
-
-        # Supervised loss
-        S_loss = torch.nn.functional.mse_loss(
-            H_hat_supervise[:, :-1, :], H[:, 1:, :]
-        )  # Teacher forcing next output
-        return S_loss
-
-    def _discriminator_forward(self, X, T, Z, gamma=1):
-        """The discriminator forward pass and adversarial loss
-        Args:
-            - X: the input features
-            - T: the temporal information
-            - Z: the input noise
-        Returns:
-            - D_loss: the adversarial loss
-        """
-        # Discriminator
-        H = self.embedder(X, T).detach()
-        H_hat = self.supervisor(H, T).detach()
-        E_hat = self.generator(Z, T).detach()
-
-        # Forward Pass
-        Y_real = self.discriminator(H, T)  # Encoded original data
-        Y_fake = self.discriminator(H_hat, T)  # Output of supervisor
-        Y_fake_e = self.discriminator(E_hat, T)  # Output of generator
-
-        D_loss_real = torch.nn.functional.binary_cross_entropy_with_logits(
-            Y_real, torch.ones_like(Y_real)
-        )
-        D_loss_fake = torch.nn.functional.binary_cross_entropy_with_logits(
-            Y_fake, torch.zeros_like(Y_fake)
-        )
-        D_loss_fake_e = torch.nn.functional.binary_cross_entropy_with_logits(
-            Y_fake_e, torch.zeros_like(Y_fake_e)
-        )
-
-        D_loss = D_loss_real + D_loss_fake + gamma * D_loss_fake_e
-
-        return D_loss
-
-    def _generator_forward(self, X, T, Z, gamma=1):
-        """The generator forward pass
-        Args:
-            - X: the original feature input
-            - T: the temporal information
-            - Z: the noise for generator input
-        Returns:
-            - G_loss: the generator's loss
-        """
-        # Supervised Forward Pass
-        H = self.embedder(X, T)
-        H_hat_supervise = self.supervisor(H, T)
-        X_tilde = self.recovery(H, T)
-
-        # Generator Forward Pass
-        E_hat = self.generator(Z, T)
-        H_hat = self.supervisor(E_hat, T)
-
-        # Synthetic data generated
-        X_hat = self.recovery(H_hat, T)
-
-        # Generator Loss
-        # 1. Adversarial loss
-        Y_fake = self.discriminator(H_hat, T)  # Output of supervisor
-        Y_fake_e = self.discriminator(E_hat, T)  # Output of generator
-
-        G_loss_U = torch.nn.functional.binary_cross_entropy_with_logits(
-            Y_fake, torch.ones_like(Y_fake)
-        )
-        G_loss_U_e = torch.nn.functional.binary_cross_entropy_with_logits(
-            Y_fake_e, torch.ones_like(Y_fake_e)
-        )
-
-        # 2. Supervised loss
-        G_loss_S = torch.nn.functional.mse_loss(
-            H_hat_supervise[:, :-1, :], H[:, 1:, :]
-        )  # Teacher forcing next output
-
-        # 3. Two Momments
-        G_loss_V1 = torch.mean(
-            torch.abs(
-                torch.sqrt(X_hat.var(dim=0, unbiased=False) + 1e-6)
-                - torch.sqrt(X.var(dim=0, unbiased=False) + 1e-6)
-            )
-        )
-        G_loss_V2 = torch.mean(torch.abs((X_hat.mean(dim=0)) - (X.mean(dim=0))))
-
-        G_loss_V = G_loss_V1 + G_loss_V2
-
-        # 4. Summation
-        G_loss = G_loss_U + gamma * G_loss_U_e + 100 * torch.sqrt(G_loss_S) + 100 * G_loss_V
-
-        return G_loss
-
-    def _inference(self, Z, T):
-        """Inference for generating synthetic data
-        Args:
-            - Z: the input noise
-            - T: the temporal information
-        Returns:
-            - X_hat: the generated data
-        """
-        # Generator Forward Pass
-        E_hat = self.generator(Z, T)
-        H_hat = self.supervisor(E_hat, T)
-
-        # Synthetic data generated
-        X_hat = self.recovery(H_hat, T)
-        return X_hat
-
-    def forward(self, X, T, Z, obj, gamma=1):
-        """
-        Args:
-            - X: the input features (B, H, F)
-            - T: the temporal information (B)
-            - Z: the sampled noise (B, H, Z)
-            - obj: the network to be trained (`autoencoder`, `supervisor`, `generator`, `discriminator`)
-            - gamma: loss hyperparameter
-        Returns:
-            - loss: The loss for the forward pass
-            - X_hat: The generated data
-        """
-        if obj != "inference":
-            if X is None:
-                raise ValueError("`X` should be given")
-
-            X = torch.FloatTensor(X)
-            X = X.to(self.device)
-
-        if Z is not None:
-            Z = torch.FloatTensor(Z)
-            Z = Z.to(self.device)
-
-        if obj == "autoencoder":
-            # Embedder & Recovery
-            loss = self._recovery_forward(X, T)
-
-        elif obj == "supervisor":
-            # Supervisor
-            loss = self._supervisor_forward(X, T)
-
-        elif obj == "generator":
-            if Z is None:
-                raise ValueError("`Z` is not given")
-
-            # Generator
-            loss = self._generator_forward(X, T, Z)
-
-        elif obj == "discriminator":
-            if Z is None:
-                raise ValueError("`Z` is not given")
-
-            # Discriminator
-            loss = self._discriminator_forward(X, T, Z)
-
-            return loss
-
-        elif obj == "inference":
-
-            X_hat = self._inference(Z, T)
-            X_hat = X_hat.cpu().detach()
-
-            return X_hat
-
-        else:
-            raise ValueError(
-                "`obj` should be either `autoencoder`, `supervisor`, `generator`, or `discriminator`"
-            )
-
-        return loss
diff --git a/satflow/models/unet.py b/satflow/models/unet.py
index 7f936e7b..29473c4d 100644
--- a/satflow/models/unet.py
+++ b/satflow/models/unet.py
@@ -1,11 +1,11 @@
 import torch
 import pytorch_lightning as pl
-from satflow.models.base import register_model
+from nowcasting_utils.models.base import register_model
 from pl_bolts.models.vision import UNet
 import numpy as np
 from typing import Union
 import torchvision
-from satflow.models.losses import get_loss
+from nowcasting_utils.models.loss import get_loss
 
 
 @register_model
diff --git a/satflow/resources/areas.yaml b/satflow/resources/areas.yaml
deleted file mode 100644
index ee37f822..00000000
--- a/satflow/resources/areas.yaml
+++ /dev/null
@@ -1,42 +0,0 @@
-europe_tmerc:
-  description: Euro 3.0km area - Europe
-  projection:
-    proj: tmerc
-    ellps: WGS84
-    lat_0: 90.0
-    lon_0: 0.0
-    lat_ts: 60.0
-  shape:
-    height: 2048
-    width: 2560
-  area_extent:
-    lower_left_xy: [-3780000.0, -7644000.0]
-    upper_right_xy: [3900000.0, -1500000.0]
-europe_tmerc1:
-  description: Euro 1.0km area - Europe
-  projection:
-    proj: tmerc
-    ellps: WGS84
-    lat_0: 90.0
-    lon_0: 0.0
-    lat_ts: 60.0
-  shape:
-    height: 8000
-    width: 10000
-  area_extent:
-    lower_left_xy: [-3780000.0, -7644000.0]
-    upper_right_xy: [3900000.0, -1500000.0]
-europe_tmerc5:
-  description: Euro 0.5km area - Europe
-  projection:
-    proj: tmerc
-    ellps: WGS84
-    lat_0: 90.0
-    lon_0: 0.0
-    lat_ts: 60.0
-  shape:
-    height: 16000
-    width: 20000
-  area_extent:
-    lower_left_xy: [-3780000.0, -7644000.0]
-    upper_right_xy: [3900000.0, -1500000.0]
diff --git a/satflow/resources/cutdown_europe_dem.npy b/satflow/resources/cutdown_europe_dem.npy
deleted file mode 100644
index 387f1727..00000000
Binary files a/satflow/resources/cutdown_europe_dem.npy and /dev/null differ
diff --git a/satflow/resources/location_array.npy b/satflow/resources/location_array.npy
deleted file mode 100644
index 558eae18..00000000
Binary files a/satflow/resources/location_array.npy and /dev/null differ
diff --git a/satflow/run.py b/satflow/run.py
index 793e876d..828112d9 100644
--- a/satflow/run.py
+++ b/satflow/run.py
@@ -16,7 +16,7 @@ def main(config: DictConfig):
     # Imports should be nested inside @hydra.main to optimize tab completion
     # Read more here: https://github.com/facebookresearch/hydra/issues/934
     from satflow.core import utils
-    from satflow.experiments.pl_train import train
+    from satflow.experiments.train import train
 
     # A couple of optional utilities:
     # - disabling python warnings
diff --git a/satflow/version.py b/satflow/version.py
index b5fdc753..493f7415 100644
--- a/satflow/version.py
+++ b/satflow/version.py
@@ -1,1 +1,1 @@
-__version__ = "0.2.2"
+__version__ = "0.3.0"
diff --git a/satflow/visualizer/__init__.py b/satflow/visualizer/__init__.py
deleted file mode 100644
index e69de29b..00000000
diff --git a/satflow/visualizer/utils.py b/satflow/visualizer/utils.py
deleted file mode 100644
index dee8b2f8..00000000
--- a/satflow/visualizer/utils.py
+++ /dev/null
@@ -1,2 +0,0 @@
-import numpy as np
-import matplotlib.pyplot as plt
diff --git a/setup.py b/setup.py
index e80c9a33..66816a72 100644
--- a/setup.py
+++ b/setup.py
@@ -1,7 +1,9 @@
 from distutils.core import setup
+from pathlib import Path
 
-with open("requirements.txt") as f:
-    install_requires = f.read().splitlines()
+this_directory = Path(__file__).parent
+install_requires = (this_directory / "requirements.txt").read_text().splitlines()
+long_description = (this_directory / "README.md").read_text()
 
 exec(open("satflow/version.py").read())
 setup(
@@ -13,6 +15,8 @@
     company="Open Climate Fix Ltd",
     author="Jacob Bieker",
     install_requires=install_requires,
+    long_description=long_description,
+    ong_description_content_type="text/markdown",
     author_email="jacob@openclimatefix.org",
     description="Satellite Optical Flow",
 )


Test patch:
diff --git a/datasets/satflow-multi-test.tar b/datasets/satflow-multi-test.tar
deleted file mode 100644
index 42cf94a1..00000000
Binary files a/datasets/satflow-multi-test.tar and /dev/null differ
diff --git a/datasets/satflow-test.tar b/datasets/satflow-test.tar
deleted file mode 100644
index 7b33f1c9..00000000
Binary files a/datasets/satflow-test.tar and /dev/null differ
diff --git a/satflow/tests/0.nc b/satflow/tests/0.nc
new file mode 100644
index 00000000..d5b9b6ad
Binary files /dev/null and b/satflow/tests/0.nc differ
diff --git a/satflow/tests/configs/satflow.yaml b/satflow/tests/configs/satflow.yaml
deleted file mode 100644
index 88de4439..00000000
--- a/satflow/tests/configs/satflow.yaml
+++ /dev/null
@@ -1,41 +0,0 @@
-batch_size: 12
-data_dir: "datasets/"
-shuffle: 0
-sources:
-  train: "satflow-test.tar"
-  val: "satflow-test.tar"
-  test: "satflow-test.tar"
-num_workers: 12
-pin_memory: True
-config:
-  visualize: False
-  num_timesteps: 12
-  skip_timesteps: 1
-  forecast_times: 24
-  output_shape: 128
-  target_type: "cloudmask"
-  num_crops: 10
-  use_topo: False
-  use_latlon: False
-  use_time: False
-  time_aux: False
-  use_mask: False
-  use_image: False
-  add_pixel_coords: False
-  time_as_channels: False
-  bands:
-    [
-      "HRV",
-      "IR016",
-      "IR039",
-      "IR087",
-      "IR097",
-      "IR108",
-      "IR120",
-      "IR134",
-      "VIS006",
-      "VIS008",
-      "WV062",
-      "WV073",
-    ]
-  transforms: {}
diff --git a/satflow/tests/configs/satflow_all.yaml b/satflow/tests/configs/satflow_all.yaml
deleted file mode 100644
index 9ca21d1d..00000000
--- a/satflow/tests/configs/satflow_all.yaml
+++ /dev/null
@@ -1,42 +0,0 @@
-batch_size: 12
-data_dir: "datasets/"
-shuffle: 0
-sources:
-  train: "satflow-test.tar"
-  val: "satflow-test.tar"
-  test: "satflow-test.tar"
-num_workers: 12
-pin_memory: True
-config:
-  visualize: False
-  num_timesteps: 12
-  skip_timesteps: 1
-  forecast_times: 24
-  output_shape: 128
-  target_type: "cloudmask"
-  num_crops: 10
-  use_topo: True
-  use_latlon: True
-  use_time: True
-  time_aux: False
-  use_mask: True
-  use_image: True
-  add_pixel_coords: True
-  add_polar_coords: True
-  time_as_channels: False
-  bands:
-    [
-      "HRV",
-      "IR016",
-      "IR039",
-      "IR087",
-      "IR097",
-      "IR108",
-      "IR120",
-      "IR134",
-      "VIS006",
-      "VIS008",
-      "WV062",
-      "WV073",
-    ]
-  transforms: {}
diff --git a/satflow/tests/configs/satflow_channels.yaml b/satflow/tests/configs/satflow_channels.yaml
deleted file mode 100644
index 7a7989ce..00000000
--- a/satflow/tests/configs/satflow_channels.yaml
+++ /dev/null
@@ -1,27 +0,0 @@
-batch_size: 12
-data_dir: "datasets/"
-shuffle: 0
-sources:
-  train: "satflow-test.tar"
-  val: "satflow-test.tar"
-  test: "satflow-test.tar"
-num_workers: 12
-pin_memory: True
-config:
-  visualize: False
-  num_timesteps: 12
-  skip_timesteps: 1
-  forecast_times: 24
-  output_shape: 128
-  target_type: "cloudmask"
-  num_crops: 10
-  use_topo: False
-  use_latlon: False
-  use_time: False
-  time_aux: False
-  use_mask: False
-  use_image: False
-  add_pixel_coords: False
-  time_as_channels: False
-  bands: ["IR016", "IR097", "IR134", "VIS008", "WV073"]
-  transforms: {}
diff --git a/satflow/tests/configs/satflow_clouds.yaml b/satflow/tests/configs/satflow_clouds.yaml
deleted file mode 100644
index b0a5ce99..00000000
--- a/satflow/tests/configs/satflow_clouds.yaml
+++ /dev/null
@@ -1,27 +0,0 @@
-batch_size: 12
-data_dir: "datasets/"
-shuffle: 0
-sources:
-  train: "satflow-test.tar"
-  val: "satflow-test.tar"
-  test: "satflow-test.tar"
-num_workers: 12
-pin_memory: True
-config:
-  visualize: False
-  num_timesteps: 12
-  skip_timesteps: 1
-  forecast_times: 24
-  output_shape: 128
-  target_type: "cloudmask"
-  num_crops: 10
-  use_topo: True
-  use_latlon: True
-  use_time: True
-  time_aux: False
-  use_mask: True
-  use_image: False
-  add_pixel_coords: True
-  time_as_channels: False
-  bands: ["IR016", "VIS006", "VIS008"]
-  transforms: {}
diff --git a/satflow/tests/configs/satflow_crop.yaml b/satflow/tests/configs/satflow_crop.yaml
deleted file mode 100644
index 23ec90b2..00000000
--- a/satflow/tests/configs/satflow_crop.yaml
+++ /dev/null
@@ -1,42 +0,0 @@
-batch_size: 12
-data_dir: "datasets/"
-shuffle: 0
-sources:
-  train: "satflow-test.tar"
-  val: "satflow-test.tar"
-  test: "satflow-test.tar"
-num_workers: 12
-pin_memory: True
-config:
-  visualize: False
-  num_timesteps: 12
-  skip_timesteps: 1
-  forecast_times: 24
-  output_shape: 256
-  output_target: 64
-  target_type: "cloudmask"
-  num_crops: 10
-  use_topo: True
-  use_latlon: True
-  use_time: True
-  time_aux: False
-  use_mask: True
-  use_image: True
-  add_pixel_coords: True
-  time_as_channels: False
-  bands:
-    [
-      "HRV",
-      "IR016",
-      "IR039",
-      "IR087",
-      "IR097",
-      "IR108",
-      "IR120",
-      "IR134",
-      "VIS006",
-      "VIS008",
-      "WV062",
-      "WV073",
-    ]
-  transforms: {}
diff --git a/satflow/tests/configs/satflow_large.yaml b/satflow/tests/configs/satflow_large.yaml
deleted file mode 100644
index 43e64970..00000000
--- a/satflow/tests/configs/satflow_large.yaml
+++ /dev/null
@@ -1,41 +0,0 @@
-batch_size: 12
-data_dir: "datasets/"
-shuffle: 0
-sources:
-  train: "satflow-test.tar"
-  val: "satflow-test.tar"
-  test: "satflow-test.tar"
-num_workers: 12
-pin_memory: True
-config:
-  visualize: False
-  num_timesteps: 12
-  skip_timesteps: 1
-  forecast_times: 24
-  output_shape: 256
-  target_type: "cloudmask"
-  num_crops: 10
-  use_topo: True
-  use_latlon: True
-  use_time: True
-  time_aux: False
-  use_mask: True
-  use_image: True
-  add_pixel_coords: True
-  time_as_channels: False
-  bands:
-    [
-      "HRV",
-      "IR016",
-      "IR039",
-      "IR087",
-      "IR097",
-      "IR108",
-      "IR120",
-      "IR134",
-      "VIS006",
-      "VIS008",
-      "WV062",
-      "WV073",
-    ]
-  transforms: {}
diff --git a/satflow/tests/configs/satflow_time_aux.yaml b/satflow/tests/configs/satflow_time_aux.yaml
deleted file mode 100644
index 576bdd7c..00000000
--- a/satflow/tests/configs/satflow_time_aux.yaml
+++ /dev/null
@@ -1,41 +0,0 @@
-batch_size: 12
-data_dir: "datasets/"
-shuffle: 0
-sources:
-  train: "satflow-test.tar"
-  val: "satflow-test.tar"
-  test: "satflow-test.tar"
-num_workers: 12
-pin_memory: True
-config:
-  visualize: False
-  num_timesteps: 12
-  skip_timesteps: 1
-  forecast_times: 24
-  output_shape: 128
-  target_type: "cloudmask"
-  num_crops: 10
-  use_topo: True
-  use_latlon: True
-  use_time: False
-  time_aux: True
-  use_mask: True
-  use_image: True
-  add_pixel_coords: True
-  time_as_channels: False
-  bands:
-    [
-      "HRV",
-      "IR016",
-      "IR039",
-      "IR087",
-      "IR097",
-      "IR108",
-      "IR120",
-      "IR134",
-      "VIS006",
-      "VIS008",
-      "WV062",
-      "WV073",
-    ]
-  transforms: {}
diff --git a/satflow/tests/configs/satflow_time_channels.yaml b/satflow/tests/configs/satflow_time_channels.yaml
deleted file mode 100644
index b147bdf2..00000000
--- a/satflow/tests/configs/satflow_time_channels.yaml
+++ /dev/null
@@ -1,41 +0,0 @@
-batch_size: 12
-data_dir: "datasets/"
-shuffle: 0
-sources:
-  train: "satflow-test.tar"
-  val: "satflow-test.tar"
-  test: "satflow-test.tar"
-num_workers: 12
-pin_memory: True
-config:
-  visualize: False
-  num_timesteps: 12
-  skip_timesteps: 1
-  forecast_times: 24
-  output_shape: 128
-  target_type: "cloudmask"
-  num_crops: 10
-  use_topo: False
-  use_latlon: False
-  use_time: False
-  time_aux: False
-  use_mask: False
-  use_image: False
-  add_pixel_coords: False
-  time_as_channels: True
-  bands:
-    [
-      "HRV",
-      "IR016",
-      "IR039",
-      "IR087",
-      "IR097",
-      "IR108",
-      "IR120",
-      "IR134",
-      "VIS006",
-      "VIS008",
-      "WV062",
-      "WV073",
-    ]
-  transforms: {}
diff --git a/satflow/tests/configs/satflow_time_channels_all.yaml b/satflow/tests/configs/satflow_time_channels_all.yaml
deleted file mode 100644
index 0d9f52d2..00000000
--- a/satflow/tests/configs/satflow_time_channels_all.yaml
+++ /dev/null
@@ -1,42 +0,0 @@
-batch_size: 12
-data_dir: "datasets/"
-shuffle: 0
-sources:
-  train: "satflow-test.tar"
-  val: "satflow-test.tar"
-  test: "satflow-test.tar"
-num_workers: 12
-pin_memory: True
-config:
-  visualize: False
-  num_timesteps: 12
-  skip_timesteps: 1
-  forecast_times: 24
-  output_shape: 128
-  target_type: "cloudmask"
-  num_crops: 10
-  use_topo: True
-  use_latlon: True
-  use_time: True
-  time_aux: False
-  use_mask: True
-  use_image: True
-  add_pixel_coords: True
-  add_polar_coords: True
-  time_as_channels: True
-  bands:
-    [
-      "HRV",
-      "IR016",
-      "IR039",
-      "IR087",
-      "IR097",
-      "IR108",
-      "IR120",
-      "IR134",
-      "VIS006",
-      "VIS008",
-      "WV062",
-      "WV073",
-    ]
-  transforms: {}
diff --git a/satflow/tests/test_datamodule.py b/satflow/tests/test_datamodule.py
deleted file mode 100644
index 5e0cc17c..00000000
--- a/satflow/tests/test_datamodule.py
+++ /dev/null
@@ -1,128 +0,0 @@
-# NIR1.6, VIS0.8 and VIS0.6 RGB for near normal view
-import numpy as np
-from satflow.data.datamodules import SatFlowDataModule
-import yaml
-
-
-def load_config(config_file):
-    with open(config_file, "r") as cfg:
-        return yaml.load(cfg, Loader=yaml.FullLoader)
-
-
-def check_channels(config):
-    channels = len(config.get("bands", []))
-    print(f"Channels: (Bands) {channels}")
-    channels = channels + 1 if config.get("use_mask", False) else channels
-    print(f"Channels: (Use Mask) {channels}")
-    channels = (
-        channels + 3
-        if config.get("use_time", False) and not config.get("time_aux", False)
-        else channels
-    )
-    print(f"Channels: (Use Time) {channels}")
-    if config.get("time_as_channels", False):
-        # Calc number of channels + inital ones
-        channels = channels * (config["num_timesteps"] + 1)
-    print(f"Channels: (Time as Channels) {channels}")
-    channels = channels + 1 if config.get("use_topo", False) else channels
-    print(f"Channels: (Use Topo) {channels}")
-    channels = channels + 3 if config.get("use_latlon", False) else channels
-    print(f"Channels: (Use Latlon) {channels}")
-    channels = channels + 2 if config.get("add_pixel_coords", False) else channels
-    print(f"Channels: (Add Pixel Coordinates) {channels}")
-    channels = channels + 1 if config.get("add_polar_coords", False) else channels
-    print(f"Channels: (Add Polar Coordinates) {channels}")
-    return channels
-
-
-def test_satflow_cloudmask():
-    config = load_config("satflow/tests/configs/satflow.yaml")
-    cloudflow = SatFlowDataModule(**config)
-    cloudflow.setup()
-    data = next(iter(cloudflow.train_dataloader()))
-    x, y = data
-    channels = check_channels(config["config"])
-    assert x.shape == (12, 13, channels, 128, 128)
-    assert y.shape == (12, 24, 1, 128, 128)
-
-
-def test_satflow_all():
-    config = load_config("satflow/tests/configs/satflow_all.yaml")
-    cloudflow = SatFlowDataModule(**config)
-    cloudflow.setup()
-    data = next(iter(cloudflow.train_dataloader()))
-    x, image = data
-    channels = check_channels(config["config"])
-    assert x.shape == (12, 13, channels, 128, 128)
-    assert image.shape == (12, 24, 12, 128, 128)
-    assert not np.allclose(image[0].numpy(), image[-1].numpy())
-
-
-def test_satflow_large():
-    config = load_config("satflow/tests/configs/satflow_large.yaml")
-    cloudflow = SatFlowDataModule(**config)
-    cloudflow.setup()
-    data = next(iter(cloudflow.train_dataloader()))
-    x, image = data
-    channels = check_channels(config["config"])
-    assert x.shape == (12, 13, channels, 256, 256)
-    assert image.shape == (12, 24, 12, 256, 256)
-    assert not np.allclose(image[0].numpy(), image[-1].numpy())
-
-
-def test_satflow_crop():
-    config = load_config("satflow/tests/configs/satflow_crop.yaml")
-    cloudflow = SatFlowDataModule(**config)
-    cloudflow.setup()
-    data = next(iter(cloudflow.train_dataloader()))
-    x, image = data
-    channels = check_channels(config["config"])
-    assert x.shape == (12, 13, channels, 256, 256)
-    assert image.shape == (12, 24, 12, 64, 64)
-    assert not np.allclose(image[0].numpy(), image[-1].numpy())
-    assert not np.allclose(x[0].numpy(), x[-1].numpy())
-
-
-def test_satflow_channels():
-    config = load_config("satflow/tests/configs/satflow_channels.yaml")
-    cloudflow = SatFlowDataModule(**config)
-    cloudflow.setup()
-    data = next(iter(cloudflow.train_dataloader()))
-    x, y = data
-    channels = check_channels(config["config"])
-    assert x.shape == (12, 13, channels, 128, 128)
-    assert y.shape == (12, 24, 1, 128, 128)
-
-
-def test_satflow_time_channels():
-    config = load_config("satflow/tests/configs/satflow_time_channels.yaml")
-    cloudflow = SatFlowDataModule(**config)
-    cloudflow.setup()
-    data = next(iter(cloudflow.train_dataloader()))
-    x, y = data
-    channels = check_channels(config["config"])
-    assert x.shape[1] == channels
-    assert x.shape == (12, 156, 128, 128)
-    assert y.shape == (12, 24, 128, 128)
-
-
-def test_satflow_time_channels_all():
-    config = load_config("satflow/tests/configs/satflow_time_channels_all.yaml")
-    cloudflow = SatFlowDataModule(**config)
-    cloudflow.setup()
-    data = next(iter(cloudflow.train_dataloader()))
-    x, image = data
-    channels = check_channels(config["config"])
-    assert x.numpy().shape == (12, channels, 128, 128)
-    assert image.shape == (12, 12 * 24, 128, 128)
-    assert not np.allclose(image[0].numpy(), image[-1].numpy())
-
-
-def test_cloudflow():
-    config = load_config("satflow/tests/configs/satflow.yaml")
-    cloudflow = SatFlowDataModule(**config)
-    cloudflow.setup()
-    data = next(iter(cloudflow.train_dataloader()))
-    x, y = data
-    assert x.numpy().shape == (12, 13, 12, 128, 128)
-    assert y.numpy().shape == (12, 24, 1, 128, 128)
diff --git a/satflow/tests/test_datamodules.py b/satflow/tests/test_datamodules.py
new file mode 100644
index 00000000..c0d2601a
--- /dev/null
+++ b/satflow/tests/test_datamodules.py
@@ -0,0 +1,10 @@
+from satflow.data.datamodules import SatFlowDataModule
+
+
+def test_datamodule_subsetting():
+    dataset = SatFlowDataModule(fake_data=True)
+    dataset.setup()
+    train_dset = dataset.train_dataloader()
+    sample, target = next(iter(train_dset))
+    assert sample["sat_data"].shape == (32, 7, 16, 16, 12)
+    assert target["sat_data"].shape == (32, 48, 16, 16, 12)
diff --git a/satflow/tests/test_dataset.py b/satflow/tests/test_dataset.py
deleted file mode 100644
index 46f3766c..00000000
--- a/satflow/tests/test_dataset.py
+++ /dev/null
@@ -1,157 +0,0 @@
-# NIR1.6, VIS0.8 and VIS0.6 RGB for near normal view
-import numpy as np
-from satflow.data.datasets import SatFlowDataset, CloudFlowDataset
-import webdataset as wds
-import yaml
-
-
-def load_config(config_file):
-    with open(config_file, "r") as cfg:
-        return yaml.load(cfg, Loader=yaml.FullLoader)["config"]
-
-
-def check_channels(config):
-    channels = len(config.get("bands", []))
-    print(f"Channels: (Bands) {channels}")
-    channels = channels + 1 if config.get("use_mask", False) else channels
-    print(f"Channels: (Use Mask) {channels}")
-    channels = (
-        channels + 3
-        if config.get("use_time", False) and not config.get("time_aux", False)
-        else channels
-    )
-    print(f"Channels: (Use Time) {channels}")
-    if config.get("time_as_channels", False):
-        # Calc number of channels + inital ones
-        channels = channels * (config["num_timesteps"] + 1)
-    print(f"Channels: (Time as Channels) {channels}")
-    channels = channels + 1 if config.get("use_topo", False) else channels
-    print(f"Channels: (Use Topo) {channels}")
-    channels = channels + 3 if config.get("use_latlon", False) else channels
-    print(f"Channels: (Use Latlon) {channels}")
-    channels = channels + 2 if config.get("add_pixel_coords", False) else channels
-    print(f"Channels: (Add Pixel Coordinates) {channels}")
-    channels = channels + 1 if config.get("add_polar_coords", False) else channels
-    print(f"Channels: (Add Polar Coordinates) {channels}")
-    return channels
-
-
-def test_satflow_cloudmask():
-    dataset = wds.WebDataset("datasets/satflow-test.tar")
-    config = load_config("satflow/tests/configs/satflow.yaml")
-    channels = check_channels(config)
-    cloudflow = SatFlowDataset([dataset], config)
-    data = next(iter(cloudflow))
-    x, y = data
-    assert x.shape == (13, channels, 128, 128)
-    assert y.shape == (24, 1, 128, 128)
-
-
-def test_satflow_all():
-    dataset = wds.WebDataset("datasets/satflow-test.tar")
-    config = load_config("satflow/tests/configs/satflow_all.yaml")
-    cloudflow = SatFlowDataset([dataset], config)
-    data = next(iter(cloudflow))
-    x, image = data
-    channels = check_channels(config)
-    assert x.shape == (13, channels, 128, 128)
-    assert image.shape == (24, 12, 128, 128)
-    assert not np.allclose(image[0], image[-1])
-
-
-def test_satflow_large():
-    dataset = wds.WebDataset("datasets/satflow-test.tar")
-    config = load_config("satflow/tests/configs/satflow_large.yaml")
-    cloudflow = SatFlowDataset([dataset], config)
-    data = next(iter(cloudflow))
-    x, image = data
-    channels = check_channels(config)
-    assert x.shape == (13, channels, 256, 256)
-    assert image.shape == (24, 12, 256, 256)
-    assert not np.allclose(image[0], image[-1])
-
-
-def test_satflow_crop():
-    dataset = wds.WebDataset("datasets/satflow-test.tar")
-    config = load_config("satflow/tests/configs/satflow_crop.yaml")
-    cloudflow = SatFlowDataset([dataset], config)
-    data = next(iter(cloudflow))
-    x, image = data
-    channels = check_channels(config)
-    assert x.shape == (13, channels, 256, 256)
-    assert image.shape == (24, 12, 64, 64)
-    assert not np.allclose(image[0], image[-1])
-    assert not np.allclose(x[0], x[-1])
-
-
-def test_satflow_channels():
-    dataset = wds.WebDataset("datasets/satflow-test.tar")
-    config = load_config("satflow/tests/configs/satflow_channels.yaml")
-    cloudflow = SatFlowDataset([dataset], config)
-    data = next(iter(cloudflow))
-    x, y = data
-    channels = check_channels(config)
-    assert x.shape == (13, channels, 128, 128)
-    assert y.shape == (24, 1, 128, 128)
-
-
-def test_satflow_time_channels():
-    dataset = wds.WebDataset("datasets/satflow-test.tar")
-    config = load_config("satflow/tests/configs/satflow_time_channels.yaml")
-    cloudflow = SatFlowDataset([dataset], config)
-    data = next(iter(cloudflow))
-    x, y = data
-    channels = check_channels(config)
-    assert x.shape[0] == channels
-    assert x.shape == (156, 128, 128)
-    assert y.shape == (24, 128, 128)
-
-
-def test_satflow_time_channels_all():
-    dataset = wds.WebDataset("datasets/satflow-test.tar")
-    config = load_config("satflow/tests/configs/satflow_time_channels_all.yaml")
-    cloudflow = SatFlowDataset([dataset], config)
-    data = next(iter(cloudflow))
-    x, image = data
-    channels = check_channels(config)
-    assert x.shape == (channels, 128, 128)
-    assert image.shape == (12 * 24, 128, 128)
-    assert not np.allclose(image[0], image[-1])
-
-
-def test_cloudflow():
-    dataset = wds.WebDataset("datasets/satflow-test.tar")
-    config = load_config("satflow/tests/configs/satflow.yaml")
-    cloudflow = CloudFlowDataset([dataset], config)
-    data = next(iter(cloudflow))
-    x, y = data
-    assert x.shape == (13, 1, 128, 128)
-    assert y.shape == (24, 1, 128, 128)
-
-
-def test_satflow_all_deterministic_validation():
-    dataset = wds.WebDataset("datasets/satflow-test.tar")
-    config = load_config("satflow/tests/configs/satflow_all.yaml")
-    cloudflow = SatFlowDataset([dataset], config, train=False)
-    data = next(iter(cloudflow))
-    x, image = data
-    dataset2 = wds.WebDataset("datasets/satflow-test.tar")
-    cloudflow2 = SatFlowDataset([dataset2], config, train=False)
-    data = next(iter(cloudflow2))
-    x2, image2 = data
-    np.testing.assert_almost_equal(x, x2)
-    np.testing.assert_almost_equal(image, image2)
-    assert not np.allclose(image[0], image[-1])
-
-
-def test_satflow_all_deterministic_validation_restart():
-    dataset = wds.WebDataset("datasets/satflow-test.tar")
-    config = load_config("satflow/tests/configs/satflow_all.yaml")
-    cloudflow = SatFlowDataset([dataset], config, train=False)
-    data = next(iter(cloudflow))
-    x, image = data
-    data = next(iter(cloudflow))
-    x2, image2 = data
-    np.testing.assert_almost_equal(x, x2)
-    np.testing.assert_almost_equal(image, image2)
-    assert not np.allclose(image[0], image[-1])
diff --git a/satflow/tests/test_datasets.py b/satflow/tests/test_datasets.py
new file mode 100644
index 00000000..0e3828c6
--- /dev/null
+++ b/satflow/tests/test_datasets.py
@@ -0,0 +1,33 @@
+from satflow.data.datasets import SatFlowDataset
+from nowcasting_dataset.consts import (
+    SATELLITE_DATA,
+    SATELLITE_X_COORDS,
+    SATELLITE_Y_COORDS,
+    SATELLITE_DATETIME_INDEX,
+    NWP_DATA,
+    NWP_Y_COORDS,
+    NWP_X_COORDS,
+    NWP_TARGET_TIME,
+    DATETIME_FEATURE_NAMES,
+)
+
+
+def test_dataset():
+    train_dataset = SatFlowDataset(
+        1,
+        "tests/",
+        "tests/",
+        cloud="local",
+        required_keys=[
+            NWP_DATA,
+            SATELLITE_DATA,
+            SATELLITE_DATETIME_INDEX,
+            NWP_TARGET_TIME,
+        ],
+        history_minutes=10,
+        forecast_minutes=10,
+    )
+
+    sample, target = next(iter(train_dataset))
+    assert sample[SATELLITE_DATA].shape[1] == 3
+    assert target[SATELLITE_DATA].shape[1] == 2
diff --git a/satflow/tests/test_e2e.py b/satflow/tests/test_e2e.py
deleted file mode 100644
index 0ca5e24d..00000000
--- a/satflow/tests/test_e2e.py
+++ /dev/null
@@ -1,52 +0,0 @@
-import numpy as np
-from satflow.models import LitMetNet, Perceiver
-from satflow.models.base import list_models, create_model
-from satflow.data.datamodules import SatFlowDataModule
-import yaml
-import torch
-
-
-def load_config(config_file):
-    with open(config_file, "r") as cfg:
-        return yaml.load(cfg, Loader=yaml.FullLoader)
-
-
-TRAIN_SOURCES = {
-    "train": "satflow-test.tar",
-    "val": "satflow-test.tar",
-    "test": "satflow-test.tar",
-}
-DATA_DIR = "datasets/"
-
-
-def test_metnet_e2e():
-    config = load_config("satflow/configs/model/metnet.yaml")
-    config.pop("_target_")  # This is only for Hydra
-    model = LitMetNet(**config)
-    num_params = sum([x.numel() for x in model.parameters()])
-
-    data_config = load_config("satflow/configs/datamodule/metnet.yaml")
-    data_config.pop("_target_")
-    data_config["sources"] = TRAIN_SOURCES
-    data_config["data_dir"] = DATA_DIR
-    datamodule = SatFlowDataModule(**data_config)
-    datamodule.setup()
-    x, y = next(iter(datamodule.train_dataloader()))
-    assert not torch.isnan(x).any(), "Input included NaNs"
-    assert not torch.isnan(y).any(), "Target included NaNs"
-    model.train()
-    out = model(x)
-    out.mean().backward()
-    for n, x in model.named_parameters():
-        assert x.grad is not None, f"No gradient for {n}"
-    num_grad = sum([x.grad.numel() for x in model.parameters() if x.grad is not None])
-    # MetNet creates predictions for the center 1/4th
-    assert out.size() == (
-        2,
-        config["forecast_steps"],
-        config["output_channels"],
-        config["input_size"] // 4,
-        config["input_size"] // 4,
-    )
-    assert num_params == num_grad, "Some parameters are missing gradients"
-    assert not torch.isnan(out).any(), "Output included NaNs"
diff --git a/satflow/tests/test_hub.py b/satflow/tests/test_hub.py
deleted file mode 100644
index a0136897..00000000
--- a/satflow/tests/test_hub.py
+++ /dev/null
@@ -1,38 +0,0 @@
-import os
-import tempfile
-
-from huggingface_hub import CONFIG_NAME, PYTORCH_WEIGHTS_NAME
-
-from satflow.models.base import BaseModel
-from satflow.models.hub import SatFlowModelHubMixin
-
-
-class DummyModel(BaseModel, SatFlowModelHubMixin):
-    # Define hyperparameters of various types to test serialisation / deserialisation
-    # This follows the pattern adopted in LitMetNet and Perceiver
-    def __init__(
-        self, loss: str = "mse", forecast_steps: int = 42, pretrained: bool = True, lr: float = 0.05
-    ):
-        super(BaseModel, self).__init__()
-        self.loss = loss
-        self.forecast_steps = forecast_steps
-        self.pretrained = pretrained
-        self.lr = lr
-        self.save_hyperparameters()
-
-
-def test_satflow_mixin():
-    # Override some of the hyperparameters
-    config = {"loss": "rmse", "forecast_steps": 123}
-    model = DummyModel(**config)
-    with tempfile.TemporaryDirectory() as storage_folder:
-        # Save model weights, config, and model card
-        model.save_pretrained(storage_folder)
-        folder_contents = os.listdir(storage_folder)
-        assert len(folder_contents) == 3
-        assert "README.md" in folder_contents
-        assert PYTORCH_WEIGHTS_NAME in folder_contents
-        assert CONFIG_NAME in folder_contents
-        # Load new model and compare hyperparameters
-        new_model = DummyModel.from_pretrained(storage_folder)
-        assert new_model.hparams == model.hparams
diff --git a/satflow/tests/test_models.py b/satflow/tests/test_models.py
index 6b6524b7..44c63245 100644
--- a/satflow/tests/test_models.py
+++ b/satflow/tests/test_models.py
@@ -1,6 +1,5 @@
-import numpy as np
 from satflow.models import LitMetNet, Perceiver
-from satflow.models.base import list_models, create_model
+from nowcasting_utils.models.base import list_models, create_model
 import yaml
 import torch
 import pytest