Weather API

Auto-generated from source — see Weather for narrative descriptions.

ConvLSTM

keras_climate.weather.convlstm.ConvLSTMNowcaster

ConvLSTMNowcaster(input_shape=(10, 128, 128, 1), pred_steps=10, filters=(64, 64, 64), kernel_size=3, out_channels=1, name='convlstm_nowcaster')

Source code in keras_climate/weather/convlstm.py

def ConvLSTMNowcaster(
    input_shape=(10, 128, 128, 1),
    pred_steps=10,
    filters=(64, 64, 64),
    kernel_size=3,
    out_channels=1,
    name="convlstm_nowcaster",
):

    inputs = keras.Input(shape=input_shape, name="frames")
    x = inputs

    states = []
    for i, f in enumerate(filters):
        x, h, c = layers.ConvLSTM2D(
            f,
            kernel_size,
            padding="same",
            return_sequences=True,
            return_state=True,
            name=f"encoder_convlstm{i}",
        )(x)
        x = layers.BatchNormalization(name=f"encoder_bn{i}")(x)
        states.append([h, c])

    decoder_convlstms = [
        layers.ConvLSTM2D(
            f,
            kernel_size,
            padding="same",
            return_sequences=True,
            return_state=True,
            name=f"decoder_convlstm{i}",
        )
        for i, f in enumerate(filters)
    ]
    decoder_bns = [
        layers.BatchNormalization(name=f"decoder_bn{i}") for i in range(len(filters))
    ]
    head = layers.Conv3D(out_channels, 1, activation="sigmoid", name="frame_head")

    H, W = input_shape[1], input_shape[2]
    cur = layers.Lambda(
        lambda t: ops.zeros((ops.shape(t)[0], 1, H, W, out_channels), dtype=t.dtype),
        name="zero_seed",
    )(inputs)

    outputs = []
    cur_states = states
    for step in range(pred_steps):
        new_states = []
        h = cur
        for i in range(len(filters)):
            h, hs, cs = decoder_convlstms[i](h, initial_state=cur_states[i])
            h = decoder_bns[i](h)
            new_states.append([hs, cs])
        frame = head(h)
        outputs.append(frame)
        cur = frame
        cur_states = new_states

    out = layers.Concatenate(axis=1, name="stack_frames")(outputs)
    return keras.Model(inputs, out, name=name)

keras_climate.weather.convlstm.convlstm_config

convlstm_config(variant='base')

Source code in keras_climate/weather/convlstm.py

def convlstm_config(variant="base"):
    presets = {
        "small": dict(filters=(32, 32)),
        "base": dict(filters=(64, 64, 64)),
        "large": dict(filters=(128, 96, 64)),
    }
    return presets[variant]

Earthformer

keras_climate.weather.earthformer.Earthformer

Earthformer(input_shape=(10, 128, 128, 1), pred_steps=10, base_dim=64, stage_depths=(2, 2, 2), num_heads=4, cuboid_size=(2, 4, 4), name='earthformer')

Source code in keras_climate/weather/earthformer.py

def Earthformer(
    input_shape=(10, 128, 128, 1),
    pred_steps=10,
    base_dim=64,
    stage_depths=(2, 2, 2),
    num_heads=4,
    cuboid_size=(2, 4, 4),
    name="earthformer",
):
    T_in, H, W, C_in = input_shape
    inputs = keras.Input(shape=input_shape, name="frames")

    x = layers.TimeDistributed(
        layers.Conv2D(base_dim, 3, padding="same", name="conv"), name="stem"
    )(inputs)

    skips = []
    dim = base_dim
    for stage_i, depth in enumerate(stage_depths):
        x = _stage(x, dim, num_heads, depth, [cuboid_size], name=f"enc_stage{stage_i}")
        skips.append(x)
        if stage_i < len(stage_depths) - 1:
            x = layers.TimeDistributed(
                layers.Conv2D(dim * 2, 3, strides=2, padding="same", name="conv"),
                name=f"downsample{stage_i}",
            )(x)
            dim *= 2

    for stage_i in reversed(range(len(stage_depths) - 1)):
        dim //= 2
        x = layers.TimeDistributed(
            layers.Conv2DTranspose(dim, 3, strides=2, padding="same", name="conv"),
            name=f"upsample{stage_i}",
        )(x)
        x = layers.Concatenate(axis=-1, name=f"skip_concat{stage_i}")(
            [x, skips[stage_i]]
        )
        x = layers.TimeDistributed(
            layers.Conv2D(dim, 1, name="conv"), name=f"skip_proj{stage_i}"
        )(x)
        x = _stage(
            x,
            dim,
            num_heads,
            stage_depths[stage_i],
            [cuboid_size],
            name=f"dec_stage{stage_i}",
        )

    x = layers.Permute((2, 3, 1, 4), name="move_time_last")(x)
    x = layers.Reshape((H, W, T_in * dim), name="merge_time_channel")(x)
    x = layers.Dense(pred_steps * C_in, name="time_channel_proj")(x)
    x = layers.Reshape((H, W, pred_steps, C_in), name="split_time_channel")(x)
    out = layers.Permute((3, 1, 2, 4), name="move_time_first")(x)

    return keras.Model(inputs, out, name=name)

keras_climate.weather.earthformer.CuboidAttention

CuboidAttention(dim, num_heads, cuboid_size=(2, 4, 4), strategy='local', **kwargs)

Bases: Layer

Source code in keras_climate/weather/earthformer.py

def __init__(
    self, dim, num_heads, cuboid_size=(2, 4, 4), strategy="local", **kwargs
):
    super().__init__(**kwargs)
    self.dim = dim
    self.num_heads = num_heads
    self.head_dim = dim // num_heads
    self.scale = self.head_dim**-0.5
    self.cuboid_size = cuboid_size
    self.strategy = strategy

keras_climate.weather.earthformer.CuboidTransformerBlock

CuboidTransformerBlock(dim, num_heads, cuboid_size=(2, 4, 4), strategy='local', mlp_ratio=4.0, drop_path=0.0, **kwargs)

Bases: Layer

Source code in keras_climate/weather/earthformer.py

def __init__(
    self,
    dim,
    num_heads,
    cuboid_size=(2, 4, 4),
    strategy="local",
    mlp_ratio=4.0,
    drop_path=0.0,
    **kwargs,
):
    super().__init__(**kwargs)
    self.norm1 = layers.LayerNormalization(epsilon=1e-6, name="norm1")
    self.attn = CuboidAttention(dim, num_heads, cuboid_size, strategy, name="attn")
    self.drop_path = DropPath(drop_path)
    self.norm2 = layers.LayerNormalization(epsilon=1e-6, name="norm2")
    self.mlp = MLP(int(dim * mlp_ratio), dim, name="mlp")

MetNet

keras_climate.weather.metnet.MetNet

MetNet(input_shape=(7, 256, 256, 1), lead_times=12, base_filters=96, attn_dim=256, attn_layers=4, num_heads=8, downsample_factor=4, num_bins=1, name='metnet')

Source code in keras_climate/weather/metnet.py

def MetNet(
    input_shape=(7, 256, 256, 1),
    lead_times=12,
    base_filters=96,
    attn_dim=256,
    attn_layers=4,
    num_heads=8,
    downsample_factor=4,
    num_bins=1,
    name="metnet",
):
    T_in, H, W, C_in = input_shape
    inputs = keras.Input(shape=input_shape, name="frames")
    lead_time_input = keras.Input(shape=(), dtype="int32", name="lead_time_idx")

    x = layers.TimeDistributed(layers.ZeroPadding2D(1), name="stem_pad")(inputs)
    x = layers.TimeDistributed(
        layers.Conv2D(
            base_filters, 3, strides=downsample_factor, padding="valid", name="conv"
        ),
        name="stem",
    )(x)
    x = layers.TimeDistributed(layers.BatchNormalization(name="bn"), name="stem_bn")(x)
    x = layers.TimeDistributed(layers.Activation("relu", name="act"), name="stem_act")(
        x
    )

    x = layers.ConvLSTM2D(
        base_filters, 3, padding="same", return_sequences=False, name="temporal_encoder"
    )(x)

    x = _context_aggregating_cnn(x, base_filters * 2, name="context_tower")

    lead_embed = layers.Embedding(lead_times, base_filters * 2, name="lead_time_embed")(
        lead_time_input
    )
    lead_embed = layers.Reshape((1, 1, base_filters * 2), name="lead_time_reshape")(
        lead_embed
    )
    x = layers.Add(name="add_lead_time")([x, lead_embed])

    x = layers.Conv2D(attn_dim, 1, name="attn_proj_in")(x)
    for i in range(attn_layers):
        x = AxialAttention2D(attn_dim, num_heads, name=f"axial_attn{i}")(x)

    for i in range(int(downsample_factor).bit_length() - 1):
        x = layers.Conv2DTranspose(
            attn_dim // 2, 2, strides=2, padding="same", name=f"upsample{i}"
        )(x)
        x = layers.BatchNormalization(name=f"upsample_bn{i}")(x)
        x = layers.Activation("relu", name=f"upsample_relu{i}")(x)

    activation = "softmax" if num_bins > 1 else "sigmoid"
    outputs = layers.Conv2D(num_bins, 1, activation=activation, name="precip_head")(x)

    return keras.Model([inputs, lead_time_input], outputs, name=name)

keras_climate.weather.metnet.AxialAttention2D

AxialAttention2D(dim, num_heads=8, **kwargs)

Bases: Layer

Source code in keras_climate/weather/metnet.py

def __init__(self, dim, num_heads=8, **kwargs):
    super().__init__(**kwargs)
    self.dim = dim
    self.num_heads = num_heads
    self.head_dim = dim // num_heads
    self.scale = self.head_dim**-0.5

FourCastNet

keras_climate.weather.fourcastnet.FourCastNet

FourCastNet(img_size=(720, 1440), patch_size=8, in_chans=20, out_chans=20, embed_dim=768, depth=12, num_blocks=8, mlp_ratio=4.0, name='fourcastnet')

Source code in keras_climate/weather/fourcastnet.py

def FourCastNet(
    img_size=(720, 1440),
    patch_size=8,
    in_chans=20,
    out_chans=20,
    embed_dim=768,
    depth=12,
    num_blocks=8,
    mlp_ratio=4.0,
    name="fourcastnet",
):
    H, W = img_size
    grid_h, grid_w = H // patch_size, W // patch_size

    inputs = keras.Input(shape=(H, W, in_chans), name="state")
    x = layers.Conv2D(
        embed_dim, kernel_size=patch_size, strides=patch_size, name="patch_embed_proj"
    )(inputs)
    x = LearnedPositionEmbedding2D(grid_h, grid_w, embed_dim, name="pos_embed_layer")(x)

    for i in range(depth):
        x = AFNOBlock(embed_dim, num_blocks, mlp_ratio, name=f"block{i}")(x)

    x = layers.LayerNormalization(epsilon=1e-6, name="norm")(x)
    x = layers.Dense(patch_size * patch_size * out_chans, use_bias=False, name="head")(
        x
    )
    x = unpatchify_2d(x, grid_h, grid_w, patch_size, out_chans)

    return keras.Model(inputs, x, name=name)

ClimaX

keras_climate.weather.climax.ClimaX

ClimaX(img_size=(128, 256), patch_size=4, num_vars=48, embed_dim=1024, depth=8, decoder_depth=2, num_heads=16, mlp_ratio=4.0, name='climax')

Source code in keras_climate/weather/climax.py

def ClimaX(
    img_size=(128, 256),
    patch_size=4,
    num_vars=48,
    embed_dim=1024,
    depth=8,
    decoder_depth=2,
    num_heads=16,
    mlp_ratio=4.0,
    name="climax",
):
    H, W = img_size
    grid_h, grid_w = H // patch_size, W // patch_size
    num_patches = grid_h * grid_w

    field_in = keras.Input(shape=(H, W, num_vars), name="fields")
    lead_time_in = keras.Input(shape=(1,), name="lead_time")

    x = MultiVariablePatchEmbed(num_vars, patch_size, embed_dim, name="token_embeds")(
        field_in
    )
    x = VariableEmbedding(num_vars, embed_dim, name="var_embed_layer")(x)
    x = VariableAggregation(embed_dim, num_heads, name="var_agg")(x)
    x = LearnedPositionEmbedding1D(num_patches, embed_dim, name="pos_embed_layer")(x)

    lead_embed = layers.Dense(embed_dim, name="lead_time_embed")(lead_time_in)
    x = AddLeadTime(name="add_lead_time")([x, lead_embed])

    for i in range(depth):
        x = TransformerEncoderBlock(embed_dim, num_heads, mlp_ratio, name=f"blocks{i}")(
            x
        )
    x = layers.LayerNormalization(name="norm")(x)

    for i in range(decoder_depth):
        x = layers.Dense(embed_dim, activation="gelu", name=f"head{2 * i}")(x)
    out = layers.Dense(num_vars * patch_size**2, name=f"head{2 * decoder_depth}")(x)

    out = layers.Reshape(
        (grid_h, grid_w, patch_size * patch_size * num_vars), name="to_grid"
    )(out)
    out = unpatchify_2d(out, grid_h, grid_w, patch_size, num_vars)

    return keras.Model([field_in, lead_time_in], out, name=name)

Pangu-Weather

keras_climate.weather.pangu_weather.PanguWeather

PanguWeather(depths=(2, 6, 6, 2), num_heads=(6, 12, 12, 6), dims=(192, 384, 384, 192), name='pangu_weather')

Source code in keras_climate/weather/pangu_weather.py

def PanguWeather(
    depths=(2, 6, 6, 2),
    num_heads=(6, 12, 12, 6),
    dims=(192, 384, 384, 192),
    name="pangu_weather",
):
    upper_in = keras.Input(shape=(13, 721, 1440, 6), name="upper_air")
    surface_in = keras.Input(shape=(721, 1440, 7), name="surface")

    x = PatchEmbedding(dims[0], name="_input_layer")([upper_in, surface_in])

    x = EarthSpecificLayer(
        depths[0],
        dims[0],
        num_heads[0],
        (8, 181, 360),
        name="layers_EarthSpecificLayer0",
    )(x)
    skip = x

    x = DownSample(dims[0], (8, 181, 360), name="downsample")(x)

    x = EarthSpecificLayer(
        depths[1],
        dims[1],
        num_heads[1],
        (8, 91, 180),
        name="layers_EarthSpecificLayer1",
    )(x)
    x = EarthSpecificLayer(
        depths[2],
        dims[2],
        num_heads[2],
        (8, 91, 180),
        name="layers_EarthSpecificLayer2",
    )(x)

    x = UpSample(dims[3], name="upsample")(x)

    x = EarthSpecificLayer(
        depths[3],
        dims[3],
        num_heads[3],
        (8, 181, 360),
        name="layers_EarthSpecificLayer3",
    )(x)

    x = layers.Concatenate(axis=-1, name="skip_concat")([skip, x])

    upper_out, surface_out = PatchRecovery(name="_output_layer")(x)

    return keras.Model([upper_in, surface_in], [upper_out, surface_out], name=name)