Creating meps_example_reduced

# Standard library
import argparse
from pathlib import Path

# Third-party
import marimo as mo
import numpy as np
import torch

Creating meps_example_reduced#

This notebook outlines how the small test dataset meps_example_reduced was created from meps_example. The original zipped datasets were 263 MB and 2.6 GB, respectively. See the README for instructions to download and extract meps_example.

The dataset was reduced by cropping the grid and keeping fewer variables. Run this notebook from the repository root with:

python docs/notebooks/create_reduced_meps_dataset.py

The default input is data/meps_example; output directories are created under data/meps_example_reduced. Use --input-dir and --output-dir to select other locations. Rerunning overwrites matching output files. The input and output directories must differ.

The final configuration edit described below remains a manual step. Exporting this notebook to Markdown does not run the data processing.

_parser = argparse.ArgumentParser(description="Reduce the MEPS dataset")
_parser.add_argument(
    "--input-dir", type=Path, default=Path("data/meps_example")
)
_parser.add_argument(
    "--output-dir", type=Path, default=Path("data/meps_example_reduced")
)
_args = _parser.parse_args()
input_dir = _args.input_dir.resolve()
output_dir = _args.output_dir.resolve()
if not (input_dir / "static" / "nwp_xy.npy").is_file():
    raise FileNotFoundError(
        f"MEPS source dataset not found at {input_dir}. "
        "Download and extract meps_example as described in the README, "
        "or pass --input-dir."
    )
if output_dir == input_dir:
    raise ValueError("Input and output directories must differ")
(output_dir / "static").mkdir(parents=True, exist_ok=True)
for _split in ("train", "test", "val"):
    (output_dir / "samples" / _split).mkdir(parents=True, exist_ok=True)

Reduce the grid#

The number of grid points was reduced to one quarter by halving both spatial dimensions. Removing a quarter of the coordinates along each outer edge keeps the central grid points centered.

grid_xy = np.load(input_dir / "static" / "nwp_xy.npy")
num_x, num_y = grid_xy.shape[1:]
x_slice = slice(num_x // 4, 3 * num_x // 4)
y_slice = slice(num_y // 4, 3 * num_y // 4)
grid_xy_reduced = grid_xy[:, x_slice, y_slice]
np.save(output_dir / "static" / "nwp_xy.npy", grid_xy_reduced)

Cropping removes the original border. Establish a new perimeter of 10 grid points, matching the border width in meps_example.

_old_border_mask = np.load(input_dir / "static" / "border_mask.npy")
assert not np.any(_old_border_mask[10:-10, 10:-10])
assert np.all(_old_border_mask[:10, :])
assert np.all(_old_border_mask[:, :10])
assert np.all(_old_border_mask[-10:, :])
assert np.all(_old_border_mask[:, -10:])

_border_mask = np.zeros_like(grid_xy_reduced[0], dtype=bool)
_border_mask[:10] = True
_border_mask[:, :10] = True
_border_mask[-10:] = True
_border_mask[:, -10:] = True
np.save(output_dir / "static" / "border_mask.npy", _border_mask)

Crop surface geopotential and grid features to the same grid. The two-element flux_stats.pt vector does not depend on the grid shape or the number of variables, so copy it unchanged.

_surface_geopotential = np.load(
    input_dir / "static" / "surface_geopotential.npy"
)
np.save(
    output_dir / "static" / "surface_geopotential.npy",
    _surface_geopotential[x_slice, y_slice],
)

_grid_features = torch.load(input_dir / "static" / "grid_features.pt")
print(_grid_features.shape)
# Reshape to the grid, crop, then flatten the spatial dimensions again.
_grid_features_reduced = _grid_features.reshape(num_x, num_y, 4)[
    x_slice, y_slice, :
].reshape((-1, 4))
torch.save(
    _grid_features_reduced, output_dir / "static" / "grid_features.pt"
)
torch.save(
    torch.load(input_dir / "static" / "flux_stats.pt"),
    output_dir / "static" / "flux_stats.pt",
)

Reduce the variables#

Keep the first eight variables, including their weights and normalization statistics.

num_vars = 8
_parameter_weights = np.load(input_dir / "static" / "parameter_weights.npy")
np.save(
    output_dir / "static" / "parameter_weights.npy",
    _parameter_weights[:num_vars],
)
for _name in ("diff_mean", "diff_std", "parameter_mean", "parameter_std"):
    _old_file = torch.load(input_dir / "static" / f"{_name}.pt")
    torch.save(_old_file[:num_vars], output_dir / "static" / f"{_name}.pt")

Reduce the samples#

Each of samples/train, samples/test, and samples/val contains:

  • nwp_YYYYMMDDHH_mbrXXX.npy: forecast time, two spatial dimensions, then variables. Crop the spatial dimensions and keep eight variables.

  • wtr_YYYYMMDDHH.npy: two spatial dimensions, both cropped.

  • nwp_toa_downwelling_shortwave_flux_YYYYMMDDHH.npy: forecast time followed by two spatial dimensions. Keep every forecast time.

Here YYYYMMDDHH is a date and hour, and XXX is a three-digit member index. In the original example, the state array had shape (65, 268, 238, 18) and the shortwave-flux array had shape (65, 268, 238).

for _sample in ("train", "test", "val"):
    for _file in sorted((input_dir / "samples" / _sample).iterdir()):
        _data = np.load(_file)
        if "mbr" in _file.name:
            _data = _data[:, x_slice, y_slice, :num_vars]
        elif "wtr" in _file.name:
            _data = _data[x_slice, y_slice]
        else:
            _data = _data[:, x_slice, y_slice]
        np.save(output_dir / "samples" / _sample / _file.name, _data)

Update the configuration manually#

Finally, copy data_config.yaml from the source dataset into the output directory. Truncate the variable units, long names and short names to the first eight entries, and set the new grid shape.

In the original reduction, unit descriptions containing ^ were changed to ** to avoid requiring LaTeX in GitHub CI. The resulting data/meps_example_reduced directory was then archived and uploaded to a European Weather Cloud S3 bucket.