training fix, docs overhaul

Author: simon-donike

.github/workflows/docs.yml

@@ -4,7 +4,7 @@ on:
     branches:
       - main
     paths:
-      - "paper/**"
+      - "docs/**"
   workflow_dispatch:     # allow manual runs
 
 permissions:
@@ -31,7 +31,6 @@ jobs:
           pip install --upgrade pip
           pip install mkdocs-material mkdocstrings[python] pymdown-extensions
 
-      # Make your package importable for mkdocstrings (ok if no setup.py/pyproject)
       - name: Make repo importable
         run: |
           pip install -e . || true

docs/architecture.md

@@ -1,12 +1,10 @@
 # Architecture
 
-This document outlines how ESA OpenSR organises its super-resolution GAN, the major components that make up the model, and how
-each piece interacts during training and inference.
+This document outlines how ESA OpenSR organises its super-resolution GAN, the major components that make up the model, and how each piece interacts during training and inference.
 
 ## SRGAN Lightning module
 
-`opensr_srgan/model/SRGAN.py` defines `SRGAN_model`, a `pytorch_lightning.LightningModule` that encapsulates the full adversarial workflow.
-The module is initialised from a YAML configuration file and provides the following responsibilities:
+`opensr_srgan/model/SRGAN.py` defines `SRGAN_model`, a `pytorch_lightning.LightningModule` that encapsulates the full adversarial workflow. The module is initialised from a YAML configuration file and provides the following responsibilities:
 
 * **Configuration ingestion.** Uses OmegaConf to load hyperparameters, dataset choices, and logging options. Convenience helpers
   such as `_pretrain_check()` and `_compute_adv_loss_weight()` translate config values into runtime behaviour.
@@ -47,8 +45,7 @@ The generator zoo lives under `opensr_srgan/model/generators/` and can be select
 * **Advanced variants (`SRGAN_advanced.py`).** Provides additional block implementations and compatibility aliases exposed in
   `__init__.py` for backwards compatibility.
 
-Common traits across generators include configurable input channel counts (`Model.in_bands`), support for upscaling factors from
-2× to 8×, and residual scaling to stabilise deeper networks.
+Common traits across generators include configurable input channel counts (`Model.in_bands`), support for upscaling factors from 2× to 8×, and residual scaling to stabilise deeper networks.
 
 ## Discriminator options
 
@@ -59,13 +56,11 @@ Common traits across generators include configurable input channel counts (`Mode
 * **PatchGAN discriminator (`patchgan.py`).** Operates on local patches, which can improve high-frequency fidelity when training
   with large images. The depth is controlled by `n_blocks` and defaults to three layers.
 
-Both discriminators use LeakyReLU activations and strided convolutions to progressively downsample the input until a real/fake
-logit map is produced.
+Both discriminators use LeakyReLU activations and strided convolutions to progressively downsample the input until a real/fake logit map is produced.
 
 ## Loss suite and metrics
 
-`opensr_srgan/model/loss` contains the perceptual and pixel-based criteria applied to the generator outputs. The primary entry point is
-`GeneratorContentLoss`, which supports:
+`opensr_srgan/model/loss` contains the perceptual and pixel-based criteria applied to the generator outputs. The primary entry point is `GeneratorContentLoss`, which supports:
 
 * **L1 reconstruction** over all spectral bands.
 * **Spectral Angle Mapper (SAM)** to preserve spectral signatures.
@@ -95,5 +90,4 @@ These helpers allow the generator to operate in a normalised space while still r
 5. When exported, `predict_step()` can be called directly or wrapped in a Lightning `Trainer.predict()` loop for large-scale
    inference.
 
-This modular design keeps the research workflow flexible: swap components with configuration changes, extend the factories with
-new architectures, or plug in custom losses without touching the training loop itself.
+This modular design keeps the research workflow flexible: swap components with configuration changes, extend the factories with new architectures, or plug in custom losses without touching the training loop itself.

docs/assets/adv_loss_warmup.png

@@ -1,7 +1,6 @@
 # Configuration
 
-ESA OpenSR relies on YAML files to control every aspect of the training pipeline. This page documents the available keys and how
-they influence the underlying code. Use `opensr_srgan/configs/config_20m.yaml` and `opensr_srgan/configs/config_10m.yaml` as starting points.
+ESA OpenSR relies on YAML files to control every aspect of the training pipeline. This page documents the available keys and how they influence the underlying code. Use `opensr_srgan/configs/config_20m.yaml` and `opensr_srgan/configs/config_10m.yaml` as starting points.
 
 ## File structure
 
@@ -119,16 +118,14 @@ performing sweeps:
 
 ### Discriminator presets
 
-Tune discriminator depth to match the generator capacity—too shallow and adversarial loss underfits, too deep and the
-training loop destabilises. These starting points mirror the architectures bundled with the repo:
+Tune discriminator depth to match the generator capacity—too shallow and adversarial loss underfits, too deep and the training loop destabilises. These starting points mirror the architectures bundled with the repo:
 
 | Discriminator type | Recommended depth parameter | Additional notes |
 | --- | --- | --- |
 | `standard` | `n_blocks = 8` | Mirrors the original SRGAN CNN with alternating stride-1/stride-2 blocks before the dense head.】 |
 | `patchgan` | `n_blocks = 3` | Maps to the 3-layer PatchGAN (a.k.a. `n_layers`); increase to 4–5 for larger crops or when the generator is particularly sharp. |
 
-When adjusting these presets, scale generator and discriminator together and monitor adversarial loss ramps defined in
-`Training.Losses` to keep training stable.
+When adjusting these presets, scale generator and discriminator together and monitor adversarial loss ramps defined in `Training.Losses` to keep training stable.
 
 ## Optimisers
 
@@ -170,5 +167,4 @@ available for experiments that prefer a steady rise.
 * **Override selectively.** When launching through scripts or notebooks, you can load a base config and override specific fields at
   runtime using `OmegaConf.merge`.
 
-With a clear understanding of these fields, you can rapidly iterate on architectures, datasets, and training strategies without
-modifying the underlying code.
+With a clear understanding of these fields, you can rapidly iterate on architectures, datasets, and training strategies without modifying the underlying code.

docs/assets/discr_x_prob.png

@@ -1,37 +1,31 @@
 # Data
 
-The training stack ships with a single, self-contained example dataset that you can download in seconds to verify that the
-pipeline works end to end. This page explains how to fetch the sample data, how it is structured, and what you need to do when
+The training stack ships with a single, self-contained example dataset that you can download in seconds to verify that the pipeline works end to end. This page explains how to fetch the sample data, how it is structured, and what you need to do when
 you are ready to plug in your own collections.
 
 ## Example dataset
 
-The example dataset is a small Sentinel-2 crop bundle hosted on the Hugging Face Hub. Each `.npz` archive contains a
-high-resolution chip stored under the key `hr`. Low-resolution counterparts are generated on the fly by bicubic
-interpolation inside the dataset class.
+The example dataset is a small Sentinel-2 crop bundle hosted on the Hugging Face Hub. Each `.npz` archive contains a high-resolution chip stored under the key `hr`. Low-resolution counterparts are generated on the fly by bicubic interpolation inside the dataset class.
 
 * **Scale factor:** 4× upsampling between the generated LR inputs and provided HR targets.
 * **Splits:** All files except the final 20 samples are used for training; the last 20 form the validation split.
 * **Normalisation:** Values above 1.5 are assumed to be Sentinel-2 reflectance and are normalised by `1/10000`.
 
 ### Downloading the files
 
-`opensr_srgan/data/example_data/download_example_dataset.py` exposes a helper that downloads and extracts the archive
-into `example_dataset/` relative to your working directory. Run it from a Python shell or a small script:
+`opensr_srgan/data/example_data/download_example_dataset.py` exposes a helper that downloads and extracts the archive into `example_dataset/` relative to your working directory. Run it from a Python shell or a small script:
 
 ```python
 from opensr_srgan.data.example_data.download_example_dataset import get_example_dataset
 
 get_example_dataset()
 ```
 
-The helper pulls `example_dataset.zip` from the [`simon-donike/SR-GAN`](https://huggingface.co/simon-donike/SR-GAN)
-repository, extracts it, and removes the temporary archive once the copy completes.
+The helper pulls `example_dataset.zip` from the [`simon-donike/SR-GAN`](https://huggingface.co/simon-donike/SR-GAN) repository, extracts it, and removes the temporary archive once the copy completes.
 
 ### Directory layout
 
-After extraction the folder contains `.npz` chips named `hr_*.npz`. No further preparation is required. Simply point the
-configuration to the dataset by setting:
+After extraction the folder contains `.npz` chips named `hr_*.npz`. No further preparation is required. Simply point the configuration to the dataset by setting:
 
 ```yaml
 Data:
@@ -43,16 +37,11 @@ training and validation dataloaders.
 
 ## Adding new datasets
 
-When you are ready to move beyond the bundled sample data you can register a custom dataset. The repository uses a single
-factory function—`opensr_srgan.data.dataset_selector.select_dataset`—to keep the training script agnostic of individual
+When you are ready to move beyond the bundled sample data you can register a custom dataset. The repository uses a single factory function `opensr_srgan.data.dataset_selector.select_dataset` to keep the training script agnostic of individual
 collections. To integrate a new source:
 
-1. **Implement a dataset class.** Create a `torch.utils.data.Dataset` that returns `(lr, hr)` tensors and performs any
-   normalisation your sensor requires. Place the implementation somewhere under `opensr_srgan/data/`.
-2. **Update the selector.** Add a new `elif` branch to `select_dataset` that imports your dataset class, instantiates the
-   training and validation splits, and returns them.
-3. **Expose configuration hooks.** Introduce a new `Data.dataset_type` key (for example `MySensor`) and any additional fields
-   you need (paths, augmentation flags, scale factors, …). Document the required entries in your configuration file.
+1. **Implement a dataset class.** Create a `torch.utils.data.Dataset` that returns `(lr, hr)` tensors and performs any normalisation your sensor requires. Place the implementation somewhere under `opensr_srgan/data/`.
+2. **Update the selector.** Add a new `elif` branch to `select_dataset` that imports your dataset class, instantiates the training and validation splits, and returns them.
+3. **Expose configuration hooks.** Introduce a new `Data.dataset_type` key (for example `MyDataset`) and any additional fields you need (paths, augmentation flags, scale factors, …). Document the required entries in your configuration file.
 
-Following this pattern keeps `opensr_srgan.train` untouched: once the selector knows about your dataset you can launch
-training through the CLI or the Python API without further changes.
+Following this pattern keeps `opensr_srgan.train` untouched: once the selector knows about your dataset you can launch training through the CLI or the Python API without further changes.