ecmwf

/

aifs-ens-1.0

like 12

Follow

ECMWF 156

Hello!

Are there plans to publish the O96 model used in the Subseasonal section of the paper?

I've tried running the N320 model, published in this repo, for 46 days but the forecast becomes unstable around day 25-30.
In the plot below you can see the "2t" averaged over the globe and how it drops unexpectectly with increasing lead time.

This is the code that I used to create the plot

import argparse
import datetime
from collections import defaultdict

import earthkit.data as ekd
import earthkit.regrid as ekr
import mlflow
import numpy as np
import pandas as pd
from anemoi.inference.runners.simple import SimpleRunner
from ecmwf.opendata import Client as OpendataClient
from matplotlib import pyplot as plt

PARAM_SFC = ["10u", "10v", "2d", "2t", "msl", "skt", "sp", "tcw"]
PARAM_SFC_FC = ["lsm", "z", "slor", "sdor"]
PARAM_SOIL = ["sot"]
PARAM_PL = ["gh", "t", "u", "v", "w", "q"]
LEVELS = [1000, 925, 850, 700, 600, 500, 400, 300, 250, 200, 150, 100, 50]
SOIL_LEVELS = [1, 2]


def get_open_data(param, date, levelist=[], number=None):
    fields = defaultdict(list)
    # Get the data for the current date and the previous date
    for date in [date - datetime.timedelta(hours=6), date]:
        if number is None:
            data = ekd.from_source(
                "ecmwf-open-data", date=date, param=param, levelist=levelist
            )
        else:
            data = ekd.from_source(
                "ecmwf-open-data",
                date=date,
                param=param,
                levelist=levelist,
                number=[number],
                stream="enfo",
            )

        for f in data:
            # Open data is between -180 and 180, we need to shift it to 0-360
            assert f.to_numpy().shape == (721, 1440)
            values = np.roll(f.to_numpy(), -f.shape[1] // 2, axis=1)
            # Interpolate the data to from 0.25 to N320
            values = ekr.interpolate(values, {"grid": (0.25, 0.25)}, {"grid": "N320"})
            # Add the values to the list
            name = (
                f"{f.metadata('param')}_{f.metadata('levelist')}"
                if levelist
                else f.metadata("param")
            )
            fields[name].append(values)

    # Create a single matrix for each parameter
    for param, values in fields.items():
        fields[param] = np.stack(values)

    return fields


def main(
    experiment_name: str,
    in_grid: str,
    out_grid: float,
):
    # Setup Mlflow
    try:
        experiment = mlflow.get_experiment_by_name(experiment_name)
        experiment_id = experiment.experiment_id
    except AttributeError:
        experiment_id = mlflow.create_experiment(experiment_name)
    with mlflow.start_run(experiment_id=experiment_id) as _:
        # setup initial state
        DATE = OpendataClient("ecmwf").latest()
        print("Initial date is", DATE)
        fields = {}
        number = None
        fields.update(get_open_data(param=PARAM_SFC, number=number, date=DATE))
        fields.update(get_open_data(param=PARAM_SFC_FC, date=DATE))
        soil = get_open_data(
            param=PARAM_SOIL, levelist=SOIL_LEVELS, number=number, date=DATE
        )
        mapping = {"sot_1": "stl1", "sot_2": "stl2", "vsw_1": "swvl1", "vsw_2": "swvl2"}
        for k, v in soil.items():
            fields[mapping[k]] = v
        fields.update(
            get_open_data(param=PARAM_PL, levelist=LEVELS, number=number, date=DATE)
        )
        # Transform GH to Z
        for level in LEVELS:
            gh = fields.pop(f"gh_{level}")
            fields[f"z_{level}"] = gh * 9.80665
        input_state = dict(date=DATE, fields=fields)
        # Run single forecast
        checkpoint = {"huggingface": "ecmwf/aifs-ens-1.0"}
        runner = SimpleRunner(checkpoint, device="cuda")
        rows = []
        for output_state in runner.run(input_state=input_state, lead_time=int(24 * 46)):
            lat_lon = ekr.interpolate(
                output_state["fields"]["2t"],
                {"grid": in_grid},
                {"grid": (out_grid, out_grid)},
            )
            t2m_mean = lat_lon.mean().item()
            curr_date = output_state["date"]
            lead_time_days = (curr_date - DATE).total_seconds() / 86400.0
            rows.append({"LeadTimeDays": lead_time_days, "AverageT2M": t2m_mean})
        output_pdf = pd.DataFrame(rows)
        fig, ax = plt.subplots(dpi=300, figsize=(6, 3))
        output_pdf.plot(x="LeadTimeDays", y="AverageT2M", ax=ax)
        ax.set_title(f"{DATE} forecast")
        plt.tight_layout()
        mlflow.log_figure(fig, "forecast_plot.png")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Run ECMWF repo AIFS-ENS min example")
    parser.add_argument("--experiment_name", default="weather-dev", type=str)
    parser.add_argument("--in_grid", type=str, default="N320")
    parser.add_argument("--out_grid", type=float, default=0.25)
    args = parser.parse_args()
    print(f"Input arguments: {args}")
    main(**vars(args))