File size: 2,365 Bytes
87b732d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 | """Copyright (c) Microsoft Corporation. Licensed under the MIT license."""
from typing import Callable
import torch
import torch.nn as nn
from .aurora_normalisation import level_to_str
__all__ = ["LevelConditioned"]
class LevelConditioned(nn.Module):
"""A module with pressure-level-specific parameters."""
def __init__(
self,
construct_module: Callable[[], nn.Module],
levels: tuple[int | float, ...],
levels_dim: int,
) -> None:
"""Instantiate.
Args:
construct_module (Callable[[], :class:`nn.Module`]): Function that construct a new
instance of the module that should have pressure-level-specific parameters.
levels (tuple[int | float, ...]): All possible pressure levels
levels_dim (int): Dimension of the input that ranges of pressure levels.
"""
super().__init__()
self.levels_dim = levels_dim
self.layers = torch.nn.ParameterDict(
{level_to_str(level): construct_module() for level in levels}
)
def forward(
self, x: torch.Tensor, *args, levels: tuple[int | float, ...], **kw_args
) -> torch.Tensor:
"""Run the module.
Args:
x (:class:`torch.Tensor`): Input.
*args (object): Further arguments.
levels (tuple[int | float, ...]): Pressure levels in input `x`.
**kw_args (dict): Further keyword arguments.
Returns:
:class:`torch.Tensor`: Output of applying the module to `x`, where the appropriate
modules with pressure-level-specific parameters are applied to the appropriate
elements in `x` along dimension `self.levels_dim`.
"""
# Resolve `self.levels_dim` to a normal index.
levels_dim = self.levels_dim
while levels_dim < 0:
levels_dim += len(x.shape)
if x.shape[levels_dim] != len(levels):
raise ValueError("Incorrect number of pressure levels.")
def index(i: int) -> tuple[slice | int, ...]:
return levels_dim * (slice(None),) + (i,)
return torch.stack(
[
self.layers[level_to_str(level)](x[index(i)], *args, **kw_args)
for i, level in enumerate(levels)
],
dim=levels_dim,
)
|