Fully-connected neural network retrieval#
This notebook implements a precipitation retrieval for the IPWG SatRain dataset based on a fully-connected neural network implemented using PyTorch and Lightning.
NOTE: This notebook can be run on Google Colab. To install the necessary dependencies uncomment the following cell and execute it.
#!pip install satrain[complete]
The training data#
The satrain package provides dataset classes that take care of downloading, preprocessing and loading of the SatRain training data. The satrain.pytorch.dataset.SatRainTabular class implements the PyTorch dataset interface for the SatRain data in tabular format, i.e., for pixel-based retrievals.
The retrieval input data to load and the quality criteria for the reference surface precipitation data can be configured using the InputConfig classes from the satrain.input module and the TargetConfig class from the ipgml.target module, respectively.
For this example notebook we use only GMI observations as input. We configure the GMI input to not include the earth-incidence angles. Moreover, enable minimum-maximum normalization of the input and replace NAN values with -1.5.
NOTE: Both input and reference data in the SatRain dataset may contain NANs because observations or precipitation estimates may be missing or of insufficient quality. It is up to the user to handle those.
We choose the on_swath geometry for the retrieval, which is a natural choice for pixel-based retrievals. We also set up batching in the dataset, which is more efficient for tabular data than leaving it to the PyTorch data loader to perform the batching.
from satrain.input import GMI, Ancillary, Geo, GeoIR
from satrain.target import TargetConfig
inputs = [GMI(normalize="minmax", nan=-1.5, include_angles=False)]
target_config = TargetConfig(min_rqi=0.5)
geometry = "on_swath"
batch_size = 1024
With these settings, we can instantiate the training dataset. By setting stack=True we also tell the retrieval to stack all input tensors instead of loading the input data as a dictionary.
Note: The current implementation of the dataset loads all training data into memory upon instantiation. Therefore, executing the cell below may take some time.
import xarray as xr
data = xr.load_dataset("/home/simon/data/ipwgml/satrain/gmi/training/xs/gridded/2019/01/06/gmi_20190106182924.nc")
from torch.utils.data import DataLoader
from satrain.pytorch.datasets import SatRainTabular
training_data = SatRainTabular(
base_sensor="gmi",
geometry=geometry,
split="training",
subset="s",
retrieval_input=inputs,
batch_size=batch_size,
target_config=target_config,
stack=True,
download=True,
)
training_loader = DataLoader(training_data, shuffle=True, batch_size=None, num_workers=1)
INFO Loading training data from 5484 training scenes.
We also create a validation loader with the same configuration.
validation_data = SatRainTabular(
base_sensor="gmi",
geometry="on_swath",
split="validation",
subset="s",
retrieval_input=inputs,
batch_size=batch_size,
target_config=target_config,
stack=True,
download=True,
shuffle=False
)
validation_loader = DataLoader(validation_data, shuffle=False, batch_size=None)
INFO Loading validation data from 838 training scenes.
Defining the neural network#
from typing import Any, Callable, Dict
import torch
from torch import optim
from torch import nn
from torch.nn.functional import binary_cross_entropy_with_logits
import lightning as L
OUTPUTS = [
"surface_precip",
"probability_of_precip",
"probability_of_heavy_precip"
]
N_EPOCHS = 20
class MLP(L.LightningModule):
"""
Lightning module implementing a multi-layer perceptron (MLP) for retrieving precipitation from satellite
observations.
"""
def __init__(
self,
n_input_features: int,
n_hidden_layers: int,
n_neurons: int,
activation_fn: Callable[[], nn.Module] = nn.GELU,
normalization_layer: Callable[[int], nn.Module] = nn.LayerNorm,
n_epochs: int = N_EPOCHS
):
"""
Args:
n_input_features: The number of features in the input
n_hidden_layers: The number of hidden layers in the MLP
n_neurons: The number of neurons in the hidden layers
activation_fn: A callable to create activation function layers.
normalization_layer: A callable to create normalization layers.
n_epochs: The numebr of epochs the model will be trained for.
"""
super().__init__()
blocks = [
nn.Linear(n_input_features, n_neurons),
normalization_layer(n_neurons),
activation_fn(),
]
for _ in range(n_hidden_layers):
blocks += [
nn.Linear(n_neurons, n_neurons),
normalization_layer(n_neurons),
activation_fn(),
]
self.body = nn.Sequential(*blocks)
heads = {}
for output in OUTPUTS:
heads[output] = nn.Sequential(
nn.Linear(n_neurons, n_neurons),
normalization_layer(n_neurons),
activation_fn(),
nn.Linear(n_neurons, 1)
)
self.heads = nn.ModuleDict(heads)
self.n_epochs = N_EPOCHS
def forward(self, retrieval_input: torch.Tensor) -> Dict[str, torch.Tensor]:
"""
Forward retrieval input through network and produce dictionary with predictions.
Args:
retrieval_input: The retrieval input as a single torch.Tensor.
Return:
A dictionary containing the predictions for 'surface_precip', 'probability_of_precip',
and 'probability_of_heavy_precip'.
"""
y = self.body(retrieval_input)
return {
name: module(y) for name, module in self.heads.items()
}
def training_step(self, batch, batch_idx) -> torch.Tensor:
"""
Calculates the loss-function gradients for the MLP.
The loss is calculated as the sum of the MSE for 'surface_precip' and the binary cross-entropy loss
for precipitation detection and heavy precipitation detection.
Args:
batch: A tuple containing the training data loaded from the data loader.
batch_idx: The index of the batch in the current epoch. Not used.
Return:
A scalar torch.Tensor containing the total loss.
"""
inpt, target = batch
pred = self(inpt)
surface_precip = target["surface_precip"]
precip_mask = target["precip_mask"]
heavy_precip_mask = target["heavy_precip_mask"]
surface_precip_pred = pred["surface_precip"]
pop = pred["probability_of_precip"]
pohp = pred["probability_of_heavy_precip"]
# MSE loss for QPE
loss_quant = ((surface_precip_pred[..., 0] - surface_precip) ** 2).mean()
# BCE loss for detection targets
loss_detect = binary_cross_entropy_with_logits(pop[..., 0], precip_mask)
loss_detect_heavy = binary_cross_entropy_with_logits(pohp[..., 0], heavy_precip_mask)
tot_loss = loss_quant + loss_detect + loss_detect_heavy
return tot_loss
def validation_step(self, batch, batch_idx) -> None:
"""
Calculates the loss-function values on validation data.
Args:
batch: A tuple containing the training data loaded from the data loader.
batch_idx: The index of the batch in the current epoch. Not used.
"""
inpt, target = batch
pred = self(inpt)
surface_precip = target["surface_precip"]
precip_mask = target["precip_mask"]
heavy_precip_mask = target["heavy_precip_mask"]
surface_precip_pred = pred["surface_precip"]
pop = pred["probability_of_precip"]
pohp = pred["probability_of_heavy_precip"]
# MSE loss for QPE
loss_quant = ((surface_precip_pred[..., 0] - surface_precip) ** 2).mean()
# BCE loss for detection targets
loss_detect = binary_cross_entropy_with_logits(pop[..., 0], precip_mask)
loss_detect_heavy = binary_cross_entropy_with_logits(pohp[..., 0], heavy_precip_mask)
tot_loss = loss_quant + loss_detect + loss_detect_heavy
opt = self.optimizers()
learning_rate = opt.param_groups[0]['lr']
self.log_dict(
{
"val_loss": loss_quant + loss_detect + loss_detect_heavy,
"val_loss_quant": loss_quant,
"val_loss_detect": loss_detect,
"val_loss_detect_heavy": loss_detect_heavy,
"learning_rate": learning_rate
},
on_epoch=True,
prog_bar=True
)
def configure_optimizers(self) -> Dict[str, Any]:
"""
We use the Adam optimizer with a cosine annealing learning rate schedule.
"""
optimizer = optim.Adam(self.parameters(), lr=5e-4)
scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=self.n_epochs)
return {
"optimizer": optimizer,
"lr_scheduler": scheduler
}
from satrain.input import calculate_input_features
input_features = calculate_input_features(inputs, stack=True)
mlp = MLP(n_input_features=input_features, n_hidden_layers=4, n_neurons=256)
Running the training#
from lightning.pytorch.callbacks import TQDMProgressBar
trainer = L.Trainer(
max_epochs=N_EPOCHS,
callbacks=[TQDMProgressBar(refresh_rate=100)]
)
trainer.fit(model=mlp, train_dataloaders=training_loader, val_dataloaders=validation_loader)
INFO: GPU available: True (cuda), used: True
INFO GPU available: True (cuda), used: True
INFO: TPU available: False, using: 0 TPU cores
INFO TPU available: False, using: 0 TPU cores
INFO: HPU available: False, using: 0 HPUs
INFO HPU available: False, using: 0 HPUs
INFO: You are using a CUDA device ('NVIDIA RTX 6000 Ada Generation') that has Tensor Cores. To properly utilize them, you should set `torch.set_float32_matmul_precision('medium' | 'high')` which will trade-off precision for performance. For more details, read https://pytorch.org/docs/stable/generated/torch.set_float32_matmul_precision.html#torch.set_float32_matmul_precision
INFO You are using a CUDA device ('NVIDIA RTX 6000 Ada Generation') that has Tensor Cores. To properly utilize them, you should set `torch.set_float32_matmul_precision('medium' | 'high')` which will trade-off precision for performance. For more details, read https://pytorch.org/docs/stable/generated/torch.set_float32_matmul_precision.html#torch.set_float32_matmul _precision
INFO: LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]
INFO LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]
INFO:
| Name | Type | Params | Mode
---------------------------------------------
0 | body | Sequential | 269 K | train
1 | heads | ModuleDict | 199 K | train
---------------------------------------------
468 K Trainable params
0 Non-trainable params
468 K Total params
1.876 Total estimated model params size (MB)
INFO | Name | Type | Params | Mode --------------------------------------------- 0 | body | Sequential | 269 K | train 1 | heads | ModuleDict | 199 K | train --------------------------------------------- 468 K Trainable params 0 Non-trainable params 468 K Total params 1.876 Total estimated model params size (MB)
/home/simon/miniconda3/envs/ipwgml/lib/python3.10/site-packages/lightning/pytorch/trainer/connectors/data_connector.py:424: The 'val_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=127` in the `DataLoader` to improve performance.
/home/simon/miniconda3/envs/ipwgml/lib/python3.10/site-packages/lightning/pytorch/trainer/connectors/data_connector.py:424: The 'train_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=127` in the `DataLoader` to improve performance.
INFO: `Trainer.fit` stopped: `max_epochs=20` reached.
INFO `Trainer.fit` stopped: `max_epochs=20` reached.
Plotting validation samples#
To get an idea of the performance of the model, let’s create a scatter plot of retrieved and reference precipitation from the validation datasets.
mlp.to(device="cpu", dtype=torch.float32)
surface_precip = []
surface_precip_ref = []
n_batches = 100
for batch_ind, (x, y) in enumerate(iter(validation_loader)):
if 100 < batch_ind:
break
with torch.no_grad():
pred = mlp(x)
surface_precip += [pred["surface_precip"]]
surface_precip_ref += [y["surface_precip"]]
surface_precip = torch.cat(surface_precip).numpy().ravel()
surface_precip_ref = torch.cat(surface_precip_ref).numpy().ravel()
from satrain.plotting import set_style, cmap_tbs
from matplotlib.gridspec import GridSpec
from matplotlib.colors import LogNorm
import matplotlib.pyplot as plt
import numpy as np
set_style()
fig = plt.figure(figsize=(5, 4))
gs = GridSpec(1, 2, width_ratios=[1.0, 0.075])
ax = fig.add_subplot(gs[0, 0])
bins = np.logspace(-2, 2, 41)
scatter = np.histogram2d(surface_precip_ref, surface_precip, bins=bins)[0]
mappable = ax.pcolormesh(bins, bins, scatter.T, norm=LogNorm(), cmap=cmap_tbs)
ax.plot(bins, bins, color="grey", linestyle="--")
ax.set_xlabel("Reference [mm h$^{-1}$]")
ax.set_ylabel("Retrieved [mm h$^{-1}$]")
ax.set_xscale("log")
ax.set_yscale("log")
ax.set_aspect(1.0)
bias = (surface_precip - surface_precip_ref).mean() / surface_precip_ref.mean()
mse = ((surface_precip - surface_precip_ref) ** 2).mean()
corr_coef = np.corrcoef(surface_precip, surface_precip_ref)[0, 1]
ax.text(0.015, 30, f"Bias = {bias:.2f}%\nMSE = {mse:.2} (mm h$^{{-1}}$)$^2$\nCorrelation coef.: {corr_coef:.2f}", fontsize=7, color="grey")
cax = fig.add_subplot(gs[:, -1])
plt.colorbar(mappable, label="Counts", cax=cax)
<matplotlib.colorbar.Colorbar at 0x7f3160a83880>
Evaluating the retrieval#
In oder to evaluate the fully-connected precipitation retrieval, we need a function to conventional precipitation retrievals, we can use the satrain.evaluation.Evaluator. The evaluates the retrieval using the exact same data used to evaluate the IMERG and GPROF retrievals and thus ensures that the results are comparable. To ensure consistency between training and evaluation data, we instantiate the evaluator with the same values for geometry and retrieval_input as the training and validation dataset objects.
from satrain.evaluation import Evaluator
evaluator = Evaluator(
base_sensor="gmi",
domain="korea",
geometry=geometry,
retrieval_input=inputs,
download=True
)
The retrieval callback function#
Evaluating the MLP using the satrain.evaluation.Evaluator requires implementing a retrieval callback function that the evaluator can call to obtain the rerieval results for a given collocation scene. The satrain.pytorch module provides a wrapper class that turns a given Pytorch retrieval into such a callback function. The PytorchRetrieval class simply takes the data from the evaluator and converts it to torch.Tensor objects and puts the results back into an xarray.Dataset.
from satrain.pytorch import PytorchRetrieval
mlp_retrieval = PytorchRetrieval(mlp, retrieval_input=inputs, stack=True, device=torch.device("cuda"))
Case study#
fig = evaluator.plot_retrieval_results(516, mlp_retrieval, input_data_format="tabular", batch_size=1024)
Evaluation#
evaluator.evaluate(retrieval_fn=mlp_retrieval, input_data_format="tabular", batch_size=4048, n_processes=1)
Results#
Precipitation quantification#
evaluator.get_precip_quantification_results(name="MLP (GMI)").T
| MLP (GMI) | ERA5 | GPROF V7 (GMI) | PERSIANN CCS | PERSIANN PDIR-Now | PU-Net | GPROF-IR | |
|---|---|---|---|---|---|---|---|
| Bias [$\%$] | -19.272751 | -12.580348 | -27.236600 | -31.790741 | -21.130333 | -6.975220 | -19.987970 |
| MAE [$mm h^{-1}$] | 0.137514 | 0.212744 | 0.338000 | 0.242497 | 0.247181 | 0.208386 | 0.172630 |
| MSE [$(mm h^{-1})^2$] | 1.631004 | 2.465700 | 5.175863 | 3.129146 | 2.917237 | 2.523134 | 1.996170 |
| SMAPE$_{0.1}$ [$\%$] | 78.110466 | 105.643559 | 83.221078 | 174.468526 | 158.385099 | 124.117068 | 87.419827 |
| Correlation coeff. [] | 0.690980 | 0.336344 | 0.650835 | 0.194232 | 0.257086 | 0.410891 | 0.566848 |
| Effective resolution [$^\circ$] | 0.273716 | inf | 0.279000 | inf | inf | inf | 1.061470 |
sc = evaluator.precip_quantification_metrics[-1].compute()
Precipitation detection#
evaluator.get_precip_detection_results(name="MLP (GMI)").T
| MLP (GMI) | ERA5 | GPROF V7 (GMI) | PERSIANN CCS | PERSIANN PDIR-Now | PU-Net | GPROF-IR | |
|---|---|---|---|---|---|---|---|
| POD [] | 0.660506 | NaN | 0.756630 | 0.252253 | 0.370663 | 0.634386 | 0.662734 |
| FAR [] | 0.215590 | NaN | 0.351667 | 0.607171 | 0.658323 | 0.519121 | 0.364343 |
| HSS [] | 0.769311 | NaN | 0.574699 | 0.360400 | 0.282276 | 0.423919 | 0.604062 |
Probabilistic precipitation detection#
evaluator.get_prob_precip_detection_results(name="MLP (GMI)").T
| MLP (GMI) | ERA5 | GPROF V7 (GMI) | PERSIANN CCS | PERSIANN PDIR-Now | PU-Net | GPROF-IR | |
|---|---|---|---|---|---|---|---|
| AUC [] | 0.695487 | 0.0 | 0.499888 | NaN | 0.0 | NaN | 0.613103 |
Heavy precipitation detection#
evaluator.get_heavy_precip_detection_results(name="MLP (GMI)").T
| MLP (GMI) | ERA5 | GPROF V7 (GMI) | PERSIANN CCS | PERSIANN PDIR-Now | PU-Net | GPROF-IR | |
|---|---|---|---|---|---|---|---|
| POD [] | 0.092844 | NaN | 0.756630 | 0.252253 | 0.370663 | 0.634386 | 0.662734 |
| FAR [] | 0.256023 | NaN | 0.351667 | 0.607171 | 0.658323 | 0.519121 | 0.364343 |
| HSS [] | 0.743862 | NaN | 0.574699 | 0.360400 | 0.282276 | 0.423919 | 0.604062 |
Heavy probabilistic precipitation detection#
evaluator.get_prob_heavy_precip_detection_results(name="MLP (GMI)").T
| MLP (GMI) | ERA5 | GPROF V7 (GMI) | PERSIANN CCS | PERSIANN PDIR-Now | PU-Net | GPROF-IR | |
|---|---|---|---|---|---|---|---|
| AUC [] | 0.503429 | 0.0 | 0.506518 | NaN | 0.0 | NaN | 0.292481 |