NowcastNet¶

Model Training CommandModel Evaluation CommandModel Export CommandModel Inference Command

None

# linux
wget -c https://paddle-org.bj.bcebos.com/paddlescience/datasets/nowcastnet/mrms.tar
# windows
# curl https://paddle-org.bj.bcebos.com/paddlescience/datasets/nowcastnet/mrms.tar -o mrms.tar
mkdir ./datasets
tar -xvf mrms.tar -C ./datasets/
python nowcastnet.py mode=eval EVAL.pretrained_model_path=https://paddle-org.bj.bcebos.com/paddlescience/models/nowcastnet/nowcastnet_pretrained.pdparams

python nowcastnet.py mode=export

# linux
wget -c https://paddle-org.bj.bcebos.com/paddlescience/datasets/nowcastnet/mrms.tar
# windows
# curl https://paddle-org.bj.bcebos.com/paddlescience/datasets/nowcastnet/mrms.tar -o mrms.tar
mkdir ./datasets
tar -xvf mrms.tar -C ./datasets/
python nowcastnet.py mode=infer

1. Background Introduction¶

Deep learning has recently emerged as a powerful tool for weather forecasting, particularly for precipitation nowcasting using radar data. These methods leverage vast amounts of radar composite observations to train end-to-end neural networks, often without explicit reliance on physical laws.

Here, we reproduce NowcastNet, a nonlinear model designed for extreme precipitation nowcasting. NowcastNet unifies physical evolution schemes with conditional learning within a neural network framework, enabling effective end-to-end optimization.

2. Model Principle¶

This chapter only briefly introduces the model principle of NowcastNet. For detailed theoretical derivation, please read Skilful nowcasting of extreme precipitation with NowcastNet.

The model architecture is illustrated below:

nowcastnet-arch — NowcastNet Network Model

The model utilizes pre-trained weights for inference. We detail the inference process below.

3. Model Construction¶

The PaddleScience implementation is as follows:

examples/nowcastnet/nowcastnet.py
if cfg.CASE_TYPE == "large":
    dataset_path = cfg.LARGE_DATASET_PATH
    model_cfg = cfg.MODEL.large
    output_dir = osp.join(cfg.output_dir, "large")
elif cfg.CASE_TYPE == "normal":
    dataset_path = cfg.NORMAL_DATASET_PATH
    model_cfg = cfg.MODEL.normal
    output_dir = osp.join(cfg.output_dir, "normal")
else:
    raise ValueError(
        f"cfg.CASE_TYPE should in ['normal', 'large'], but got '{cfg.mode}'"
    )
model = ppsci.arch.NowcastNet(**model_cfg)

examples/nowcastnet/conf/nowcastnet.yaml
# model settings
MODEL:
  normal:
    input_keys: ["input"]
    output_keys: ["output"]
    input_length: 9
    total_length: 29
    image_width: 512
    image_height: 512
    image_ch: 2
    ngf: 32
  large:
    input_keys: ["input"]
    output_keys: ["output"]
    input_length: 9
    total_length: 29
    image_width: 1024
    image_height: 1024

Here, input_keys and output_keys denote the input and output variable names of the network model.

4. Model Evaluation Visualization¶

After configuration, pass the instantiated objects to ppsci.solver.Solver:

examples/nowcastnet/nowcastnet.py
solver = ppsci.solver.Solver(
    model,
    output_dir=output_dir,
    pretrained_model_path=cfg.EVAL.pretrained_model_path,
)

Next, initialize VisualizerRadar to generate visualization results:

examples/nowcastnet/nowcastnet.py
visualizer = {
    "v_nowcastnet": ppsci.visualize.VisualizerRadar(
        {"input": frames_tensor},
        {
            "output": lambda out: out["output"],
        },
        prefix="v_nowcastnet",
        case_type=cfg.CASE_TYPE,
        total_length=model_cfg.total_length,
    )
}
solver.visualizer = visualizer
# visualize prediction
solver.visualize(batch_id)

5. Complete Code¶

examples/nowcastnet/nowcastnet.py
"""
Reference: https://codeocean.com/capsule/3935105/tree/v1
"""
from os import path as osp

import hydra
import paddle
from omegaconf import DictConfig

import ppsci
from ppsci.utils import logger


def train(cfg: DictConfig):
    print("Not supported.")


def evaluate(cfg: DictConfig):
    # set random seed for reproducibility
    ppsci.utils.misc.set_random_seed(cfg.seed)
    # initialize logger
    logger.init_logger("ppsci", osp.join(cfg.output_dir, "train.log"), "info")

    if cfg.CASE_TYPE == "large":
        dataset_path = cfg.LARGE_DATASET_PATH
        model_cfg = cfg.MODEL.large
        output_dir = osp.join(cfg.output_dir, "large")
    elif cfg.CASE_TYPE == "normal":
        dataset_path = cfg.NORMAL_DATASET_PATH
        model_cfg = cfg.MODEL.normal
        output_dir = osp.join(cfg.output_dir, "normal")
    else:
        raise ValueError(
            f"cfg.CASE_TYPE should in ['normal', 'large'], but got '{cfg.mode}'"
        )
    model = ppsci.arch.NowcastNet(**model_cfg)

    input_keys = ("radar_frames",)
    dataset_param = {
        "input_keys": input_keys,
        "label_keys": (),
        "image_width": model_cfg.image_width,
        "image_height": model_cfg.image_height,
        "total_length": model_cfg.total_length,
        "dataset_path": dataset_path,
        "data_type": paddle.get_default_dtype(),
    }
    test_data_loader = paddle.io.DataLoader(
        ppsci.data.dataset.RadarDataset(**dataset_param),
        batch_size=1,
        shuffle=False,
        num_workers=cfg.CPU_WORKER,
        drop_last=True,
    )

    # initialize solver
    solver = ppsci.solver.Solver(
        model,
        output_dir=output_dir,
        pretrained_model_path=cfg.EVAL.pretrained_model_path,
    )

    for batch_id, test_ims in enumerate(test_data_loader):
        test_ims = test_ims[0][input_keys[0]].numpy()
        frames_tensor = paddle.to_tensor(
            data=test_ims, dtype=paddle.get_default_dtype()
        )
        if batch_id <= cfg.NUM_SAVE_SAMPLES:
            visualizer = {
                "v_nowcastnet": ppsci.visualize.VisualizerRadar(
                    {"input": frames_tensor},
                    {
                        "output": lambda out: out["output"],
                    },
                    prefix="v_nowcastnet",
                    case_type=cfg.CASE_TYPE,
                    total_length=model_cfg.total_length,
                )
            }
            solver.visualizer = visualizer
            # visualize prediction
            solver.visualize(batch_id)


def export(cfg: DictConfig):
    from paddle.static import InputSpec

    # set models
    if cfg.CASE_TYPE == "large":
        model_cfg = cfg.MODEL.large
    elif cfg.CASE_TYPE == "normal":
        model_cfg = cfg.MODEL.normal
    else:
        raise ValueError(
            f"cfg.CASE_TYPE should in ['normal', 'large'], but got '{cfg.mode}'"
        )
    model = ppsci.arch.NowcastNet(**model_cfg)

    # load pretrained model
    solver = ppsci.solver.Solver(
        model=model, pretrained_model_path=cfg.INFER.pretrained_model_path
    )
    # export models
    input_spec = [
        {
            key: InputSpec(
                [None, 29, model_cfg.image_width, model_cfg.image_height, 2],
                "float32",
                name=key,
            )
            for key in model_cfg.input_keys
        },
    ]
    solver.export(input_spec, cfg.INFER.export_path)


def inference(cfg: DictConfig):
    import os.path as osp

    from deploy.python_infer import pinn_predictor

    # set model predictor
    predictor = pinn_predictor.PINNPredictor(cfg)

    if cfg.CASE_TYPE == "large":
        dataset_path = cfg.LARGE_DATASET_PATH
        model_cfg = cfg.MODEL.large
        output_dir = osp.join(cfg.output_dir, "large")
    elif cfg.CASE_TYPE == "normal":
        dataset_path = cfg.NORMAL_DATASET_PATH
        model_cfg = cfg.MODEL.normal
        output_dir = osp.join(cfg.output_dir, "normal")
    else:
        raise ValueError(
            f"cfg.CASE_TYPE should in ['normal', 'large'], but got '{cfg.mode}'"
        )

    input_keys = ("radar_frames",)
    dataset_param = {
        "input_keys": input_keys,
        "label_keys": (),
        "image_width": model_cfg.image_width,
        "image_height": model_cfg.image_height,
        "total_length": model_cfg.total_length,
        "dataset_path": dataset_path,
    }
    test_data_loader = paddle.io.DataLoader(
        ppsci.data.dataset.RadarDataset(**dataset_param),
        batch_size=cfg.INFER.batch_size,
        num_workers=cfg.CPU_WORKER,
        drop_last=True,
    )
    for batch_id, test_ims in enumerate(test_data_loader):
        if batch_id > cfg.NUM_SAVE_SAMPLES:
            break
        test_ims = {"input": test_ims[0][input_keys[0]].numpy()}
        output_dict = predictor.predict(test_ims, cfg.INFER.batch_size)
        # mapping data to model_cfg.output_keys
        output_dict = {
            store_key: output_dict[infer_key]
            for store_key, infer_key in zip(model_cfg.output_keys, output_dict.keys())
        }

        visualizer = ppsci.visualize.VisualizerRadar(
            test_ims,
            {
                "output": lambda out: out["output"],
            },
            prefix="v_nowcastnet",
            case_type=cfg.CASE_TYPE,
            total_length=model_cfg.total_length,
        )
        test_ims.update(output_dict)
        visualizer.save(osp.join(output_dir, f"epoch_{batch_id}"), test_ims)


@hydra.main(version_base=None, config_path="./conf", config_name="nowcastnet.yaml")
def main(cfg: DictConfig):
    if cfg.mode == "train":
        train(cfg)
    elif cfg.mode == "eval":
        evaluate(cfg)
    elif cfg.mode == "export":
        export(cfg)
    elif cfg.mode == "infer":
        inference(cfg)
    else:
        raise ValueError(
            f"cfg.mode should in ['train', 'eval', 'export', 'infer'], but got '{cfg.mode}'"
        )


if __name__ == "__main__":
    main()

6. Result Display¶

The figures below display the model's predictions compared to the ground truth.

7. References¶

Skilful nowcasting of extreme precipitation with NowcastNet