Arch(网络模型) 模块

ppsci.arch

Arch

Bases: Layer

Base class for Network.

Source code in ppsci/arch/base.py

class Arch(nn.Layer):
    """Base class for Network."""

    input_keys: Tuple[str, ...]
    output_keys: Tuple[str, ...]

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._input_transform: Callable[
            [Dict[str, paddle.Tensor]], Dict[str, paddle.Tensor]
        ] = None

        self._output_transform: Callable[
            [Dict[str, paddle.Tensor], Dict[str, paddle.Tensor]],
            Dict[str, paddle.Tensor],
        ] = None

    def forward(self, *args, **kwargs):
        raise NotImplementedError("Arch.forward is not implemented")

    @property
    def num_params(self) -> int:
        """Return number of parameters within network.

        Returns:
            int: Number of parameters.
        """
        num = 0
        for name, param in self.named_parameters():
            if hasattr(param, "shape"):
                num += np.prod(list(param.shape), dtype="int")
            else:
                logger.warning(f"{name} has no attribute 'shape'")
        return num

    @property
    def num_buffers(self) -> int:
        """Return number of buffers within network.

        Returns:
            int: Number of buffers.
        """
        num = 0
        for name, buffer in self.named_buffers():
            if hasattr(buffer, "shape"):
                num += np.prod(list(buffer.shape), dtype="int")
            else:
                logger.warning(f"{name} has no attribute 'shape'")
        return num

    @staticmethod
    def concat_to_tensor(
        data_dict: Dict[str, paddle.Tensor], keys: Tuple[str, ...], axis=-1
    ) -> Tuple[paddle.Tensor, ...]:
        """Concatenate tensors from dict in the order of given keys.

        Args:
            data_dict (Dict[str, paddle.Tensor]): Dict contains tensor.
            keys (Tuple[str, ...]): Keys tensor fetched from.
            axis (int, optional): Axis concatenate at. Defaults to -1.

        Returns:
            Tuple[paddle.Tensor, ...]: Concatenated tensor.

        Examples:
            >>> import paddle
            >>> import ppsci
            >>> model = ppsci.arch.Arch()
            >>> # fetch one tensor
            >>> out = model.concat_to_tensor({'x':paddle.rand([64, 64, 1])}, ('x',))
            >>> print(out.dtype, out.shape)
            paddle.float32 [64, 64, 1]
            >>> # fetch more tensors
            >>> out = model.concat_to_tensor(
            ...     {'x1':paddle.rand([64, 64, 1]), 'x2':paddle.rand([64, 64, 1])},
            ...     ('x1', 'x2'),
            ...     axis=2)
            >>> print(out.dtype, out.shape)
            paddle.float32 [64, 64, 2]

        """
        if len(keys) == 1:
            return data_dict[keys[0]]
        data = [data_dict[key] for key in keys]
        return paddle.concat(data, axis)

    @staticmethod
    def split_to_dict(
        data_tensor: paddle.Tensor, keys: Tuple[str, ...], axis=-1
    ) -> Dict[str, paddle.Tensor]:
        """Split tensor and wrap into a dict by given keys.

        Args:
            data_tensor (paddle.Tensor): Tensor to be split.
            keys (Tuple[str, ...]): Keys tensor mapping to.
            axis (int, optional): Axis split at. Defaults to -1.

        Returns:
            Dict[str, paddle.Tensor]: Dict contains tensor.

        Examples:
            >>> import paddle
            >>> import ppsci
            >>> model = ppsci.arch.Arch()
            >>> # split one tensor
            >>> out = model.split_to_dict(paddle.rand([64, 64, 1]), ('x',))
            >>> for k, v in out.items():
            ...     print(f"{k} {v.dtype} {v.shape}")
            x paddle.float32 [64, 64, 1]
            >>> # split more tensors
            >>> out = model.split_to_dict(paddle.rand([64, 64, 2]), ('x1', 'x2'), axis=2)
            >>> for k, v in out.items():
            ...     print(f"{k} {v.dtype} {v.shape}")
            x1 paddle.float32 [64, 64, 1]
            x2 paddle.float32 [64, 64, 1]

        """
        if len(keys) == 1:
            return {keys[0]: data_tensor}
        data = paddle.split(data_tensor, len(keys), axis=axis)
        return {key: data[i] for i, key in enumerate(keys)}

    def register_input_transform(
        self,
        transform: Callable[[Dict[str, paddle.Tensor]], Dict[str, paddle.Tensor]],
    ):
        """Register input transform.

        Args:
            transform (Callable[[Dict[str, paddle.Tensor]], Dict[str, paddle.Tensor]]):
                Input transform of network, receive a single tensor dict and return a single tensor dict.

        Examples:
            >>> import ppsci
            >>> def transform_in(in_):
            ...     x = in_["x"]
            ...     # transform input
            ...     x_ = 2.0 * x
            ...     input_trans = {"2x": x_}
            ...     return input_trans
            >>> # `MLP` inherits from `Arch`
            >>> model = ppsci.arch.MLP(
            ...     input_keys=("2x",),
            ...     output_keys=("y",),
            ...     num_layers=5,
            ...     hidden_size=32)
            >>> model.register_input_transform(transform_in)
            >>> out = model({"x":paddle.rand([64, 64, 1])})
            >>> for k, v in out.items():
            ...     print(f"{k} {v.dtype} {v.shape}")
            y paddle.float32 [64, 64, 1]

        """
        self._input_transform = transform

    def register_output_transform(
        self,
        transform: Callable[
            [Dict[str, paddle.Tensor], Dict[str, paddle.Tensor]],
            Dict[str, paddle.Tensor],
        ],
    ):
        """Register output transform.

        Args:
            transform (Callable[[Dict[str, paddle.Tensor], Dict[str, paddle.Tensor]], Dict[str, paddle.Tensor]]):
                Output transform of network, receive two single tensor dict(raw input
                and raw output) and return a single tensor dict(transformed output).

        Examples:
            >>> import ppsci
            >>> def transform_out(in_, out):
            ...     x = in_["x"]
            ...     y = out["y"]
            ...     u = 2.0 * x * y
            ...     output_trans = {"u": u}
            ...     return output_trans
            >>> # `MLP` inherits from `Arch`
            >>> model = ppsci.arch.MLP(
            ...     input_keys=("x",),
            ...     output_keys=("y",),
            ...     num_layers=5,
            ...     hidden_size=32)
            >>> model.register_output_transform(transform_out)
            >>> out = model({"x":paddle.rand([64, 64, 1])})
            >>> for k, v in out.items():
            ...     print(f"{k} {v.dtype} {v.shape}")
            u paddle.float32 [64, 64, 1]

        """
        self._output_transform = transform

    def freeze(self):
        """Freeze all parameters.

        Examples:
            >>> import ppsci
            >>> model = ppsci.arch.Arch()
            >>> # freeze all parameters and make model `eval`
            >>> model.freeze()
            >>> assert not model.training
            >>> for p in model.parameters():
            ...     assert p.stop_gradient

        """
        for param in self.parameters():
            param.stop_gradient = True

        self.eval()

    def unfreeze(self):
        """Unfreeze all parameters.

        Examples:
            >>> import ppsci
            >>> model = ppsci.arch.Arch()
            >>> # unfreeze all parameters and make model `train`
            >>> model.unfreeze()
            >>> assert model.training
            >>> for p in model.parameters():
            ...     assert not p.stop_gradient

        """
        for param in self.parameters():
            param.stop_gradient = False

        self.train()

    def __str__(self):
        num_fc = 0
        num_conv = 0
        num_bn = 0
        for layer in self.sublayers(include_self=True):
            if isinstance(layer, nn.Linear):
                num_fc += 1
            elif isinstance(layer, (nn.Conv2D, nn.Conv3D, nn.Conv1D)):
                num_conv += 1
            elif isinstance(layer, (nn.BatchNorm, nn.BatchNorm2D, nn.BatchNorm3D)):
                num_bn += 1

        return ", ".join(
            [
                self.__class__.__name__,
                f"input_keys = {self.input_keys}",
                f"output_keys = {self.output_keys}",
                f"num_fc = {num_fc}",
                f"num_conv = {num_conv}",
                f"num_bn = {num_bn}",
                f"num_params = {self.num_params}",
                f"num_buffers = {self.num_buffers}",
            ]
        )

num_buffers property

Return number of buffers within network.

Returns:

Name	Type	Description
int	int	Number of buffers.

num_params property

Return number of parameters within network.

Returns:

Name	Type	Description
int	int	Number of parameters.

concat_to_tensor(data_dict, keys, axis=-1) staticmethod

Concatenate tensors from dict in the order of given keys.

Parameters:

Name	Type	Description	Default
data_dict	Dict[str, Tensor]	Dict contains tensor.	required
keys	Tuple[str, ...]	Keys tensor fetched from.	required
axis	int	Axis concatenate at. Defaults to -1.	-1

Returns:

Type	Description
Tuple[Tensor, ...]	Tuple[paddle.Tensor, ...]: Concatenated tensor.

Examples:

>>> import paddle
>>> import ppsci
>>> model = ppsci.arch.Arch()
>>> # fetch one tensor
>>> out = model.concat_to_tensor({'x':paddle.rand([64, 64, 1])}, ('x',))
>>> print(out.dtype, out.shape)
paddle.float32 [64, 64, 1]
>>> # fetch more tensors
>>> out = model.concat_to_tensor(
...     {'x1':paddle.rand([64, 64, 1]), 'x2':paddle.rand([64, 64, 1])},
...     ('x1', 'x2'),
...     axis=2)
>>> print(out.dtype, out.shape)
paddle.float32 [64, 64, 2]

Source code in ppsci/arch/base.py

@staticmethod
def concat_to_tensor(
    data_dict: Dict[str, paddle.Tensor], keys: Tuple[str, ...], axis=-1
) -> Tuple[paddle.Tensor, ...]:
    """Concatenate tensors from dict in the order of given keys.

    Args:
        data_dict (Dict[str, paddle.Tensor]): Dict contains tensor.
        keys (Tuple[str, ...]): Keys tensor fetched from.
        axis (int, optional): Axis concatenate at. Defaults to -1.

    Returns:
        Tuple[paddle.Tensor, ...]: Concatenated tensor.

    Examples:
        >>> import paddle
        >>> import ppsci
        >>> model = ppsci.arch.Arch()
        >>> # fetch one tensor
        >>> out = model.concat_to_tensor({'x':paddle.rand([64, 64, 1])}, ('x',))
        >>> print(out.dtype, out.shape)
        paddle.float32 [64, 64, 1]
        >>> # fetch more tensors
        >>> out = model.concat_to_tensor(
        ...     {'x1':paddle.rand([64, 64, 1]), 'x2':paddle.rand([64, 64, 1])},
        ...     ('x1', 'x2'),
        ...     axis=2)
        >>> print(out.dtype, out.shape)
        paddle.float32 [64, 64, 2]

    """
    if len(keys) == 1:
        return data_dict[keys[0]]
    data = [data_dict[key] for key in keys]
    return paddle.concat(data, axis)

freeze()

Freeze all parameters.

Examples:

>>> import ppsci
>>> model = ppsci.arch.Arch()
>>> # freeze all parameters and make model `eval`
>>> model.freeze()
>>> assert not model.training
>>> for p in model.parameters():
...     assert p.stop_gradient

Source code in ppsci/arch/base.py

def freeze(self):
    """Freeze all parameters.

    Examples:
        >>> import ppsci
        >>> model = ppsci.arch.Arch()
        >>> # freeze all parameters and make model `eval`
        >>> model.freeze()
        >>> assert not model.training
        >>> for p in model.parameters():
        ...     assert p.stop_gradient

    """
    for param in self.parameters():
        param.stop_gradient = True

    self.eval()

register_input_transform(transform)

Register input transform.

Parameters:

Name	Type	Description	Default
transform	Callable[[Dict[str, Tensor]], Dict[str, Tensor]]	Input transform of network, receive a single tensor dict and return a single tensor dict.	required

Examples:

>>> import ppsci
>>> def transform_in(in_):
...     x = in_["x"]
...     # transform input
...     x_ = 2.0 * x
...     input_trans = {"2x": x_}
...     return input_trans
>>> # `MLP` inherits from `Arch`
>>> model = ppsci.arch.MLP(
...     input_keys=("2x",),
...     output_keys=("y",),
...     num_layers=5,
...     hidden_size=32)
>>> model.register_input_transform(transform_in)
>>> out = model({"x":paddle.rand([64, 64, 1])})
>>> for k, v in out.items():
...     print(f"{k} {v.dtype} {v.shape}")
y paddle.float32 [64, 64, 1]

Source code in ppsci/arch/base.py

def register_input_transform(
    self,
    transform: Callable[[Dict[str, paddle.Tensor]], Dict[str, paddle.Tensor]],
):
    """Register input transform.

    Args:
        transform (Callable[[Dict[str, paddle.Tensor]], Dict[str, paddle.Tensor]]):
            Input transform of network, receive a single tensor dict and return a single tensor dict.

    Examples:
        >>> import ppsci
        >>> def transform_in(in_):
        ...     x = in_["x"]
        ...     # transform input
        ...     x_ = 2.0 * x
        ...     input_trans = {"2x": x_}
        ...     return input_trans
        >>> # `MLP` inherits from `Arch`
        >>> model = ppsci.arch.MLP(
        ...     input_keys=("2x",),
        ...     output_keys=("y",),
        ...     num_layers=5,
        ...     hidden_size=32)
        >>> model.register_input_transform(transform_in)
        >>> out = model({"x":paddle.rand([64, 64, 1])})
        >>> for k, v in out.items():
        ...     print(f"{k} {v.dtype} {v.shape}")
        y paddle.float32 [64, 64, 1]

    """
    self._input_transform = transform

register_output_transform(transform)

Register output transform.

Parameters:

Name	Type	Description	Default
transform	Callable[[Dict[str, Tensor], Dict[str, Tensor]], Dict[str, Tensor]]	Output transform of network, receive two single tensor dict(raw input and raw output) and return a single tensor dict(transformed output).	required

Examples:

>>> import ppsci
>>> def transform_out(in_, out):
...     x = in_["x"]
...     y = out["y"]
...     u = 2.0 * x * y
...     output_trans = {"u": u}
...     return output_trans
>>> # `MLP` inherits from `Arch`
>>> model = ppsci.arch.MLP(
...     input_keys=("x",),
...     output_keys=("y",),
...     num_layers=5,
...     hidden_size=32)
>>> model.register_output_transform(transform_out)
>>> out = model({"x":paddle.rand([64, 64, 1])})
>>> for k, v in out.items():
...     print(f"{k} {v.dtype} {v.shape}")
u paddle.float32 [64, 64, 1]

Source code in ppsci/arch/base.py

def register_output_transform(
    self,
    transform: Callable[
        [Dict[str, paddle.Tensor], Dict[str, paddle.Tensor]],
        Dict[str, paddle.Tensor],
    ],
):
    """Register output transform.

    Args:
        transform (Callable[[Dict[str, paddle.Tensor], Dict[str, paddle.Tensor]], Dict[str, paddle.Tensor]]):
            Output transform of network, receive two single tensor dict(raw input
            and raw output) and return a single tensor dict(transformed output).

    Examples:
        >>> import ppsci
        >>> def transform_out(in_, out):
        ...     x = in_["x"]
        ...     y = out["y"]
        ...     u = 2.0 * x * y
        ...     output_trans = {"u": u}
        ...     return output_trans
        >>> # `MLP` inherits from `Arch`
        >>> model = ppsci.arch.MLP(
        ...     input_keys=("x",),
        ...     output_keys=("y",),
        ...     num_layers=5,
        ...     hidden_size=32)
        >>> model.register_output_transform(transform_out)
        >>> out = model({"x":paddle.rand([64, 64, 1])})
        >>> for k, v in out.items():
        ...     print(f"{k} {v.dtype} {v.shape}")
        u paddle.float32 [64, 64, 1]

    """
    self._output_transform = transform

split_to_dict(data_tensor, keys, axis=-1) staticmethod

Split tensor and wrap into a dict by given keys.

Parameters:

Name	Type	Description	Default
data_tensor	Tensor	Tensor to be split.	required
keys	Tuple[str, ...]	Keys tensor mapping to.	required
axis	int	Axis split at. Defaults to -1.	-1

Returns:

Type	Description
Dict[str, Tensor]	Dict[str, paddle.Tensor]: Dict contains tensor.

Examples:

>>> import paddle
>>> import ppsci
>>> model = ppsci.arch.Arch()
>>> # split one tensor
>>> out = model.split_to_dict(paddle.rand([64, 64, 1]), ('x',))
>>> for k, v in out.items():
...     print(f"{k} {v.dtype} {v.shape}")
x paddle.float32 [64, 64, 1]
>>> # split more tensors
>>> out = model.split_to_dict(paddle.rand([64, 64, 2]), ('x1', 'x2'), axis=2)
>>> for k, v in out.items():
...     print(f"{k} {v.dtype} {v.shape}")
x1 paddle.float32 [64, 64, 1]
x2 paddle.float32 [64, 64, 1]

Source code in ppsci/arch/base.py

@staticmethod
def split_to_dict(
    data_tensor: paddle.Tensor, keys: Tuple[str, ...], axis=-1
) -> Dict[str, paddle.Tensor]:
    """Split tensor and wrap into a dict by given keys.

    Args:
        data_tensor (paddle.Tensor): Tensor to be split.
        keys (Tuple[str, ...]): Keys tensor mapping to.
        axis (int, optional): Axis split at. Defaults to -1.

    Returns:
        Dict[str, paddle.Tensor]: Dict contains tensor.

    Examples:
        >>> import paddle
        >>> import ppsci
        >>> model = ppsci.arch.Arch()
        >>> # split one tensor
        >>> out = model.split_to_dict(paddle.rand([64, 64, 1]), ('x',))
        >>> for k, v in out.items():
        ...     print(f"{k} {v.dtype} {v.shape}")
        x paddle.float32 [64, 64, 1]
        >>> # split more tensors
        >>> out = model.split_to_dict(paddle.rand([64, 64, 2]), ('x1', 'x2'), axis=2)
        >>> for k, v in out.items():
        ...     print(f"{k} {v.dtype} {v.shape}")
        x1 paddle.float32 [64, 64, 1]
        x2 paddle.float32 [64, 64, 1]

    """
    if len(keys) == 1:
        return {keys[0]: data_tensor}
    data = paddle.split(data_tensor, len(keys), axis=axis)
    return {key: data[i] for i, key in enumerate(keys)}

unfreeze()

Unfreeze all parameters.

Examples:

>>> import ppsci
>>> model = ppsci.arch.Arch()
>>> # unfreeze all parameters and make model `train`
>>> model.unfreeze()
>>> assert model.training
>>> for p in model.parameters():
...     assert not p.stop_gradient

Source code in ppsci/arch/base.py

def unfreeze(self):
    """Unfreeze all parameters.

    Examples:
        >>> import ppsci
        >>> model = ppsci.arch.Arch()
        >>> # unfreeze all parameters and make model `train`
        >>> model.unfreeze()
        >>> assert model.training
        >>> for p in model.parameters():
        ...     assert not p.stop_gradient

    """
    for param in self.parameters():
        param.stop_gradient = False

    self.train()

AFNONet

Bases: Arch

Adaptive Fourier Neural Network.

Parameters:

Name	Type	Description	Default
input_keys	Tuple[str, ...]	Name of input keys, such as ("input",).	required
output_keys	Tuple[str, ...]	Name of output keys, such as ("output",).	required
img_size	Tuple[int, ...]	Image size. Defaults to (720, 1440).	(720, 1440)
patch_size	Tuple[int, ...]	Path. Defaults to (8, 8).	(8, 8)
in_channels	int	The input tensor channels. Defaults to 20.	20
out_channels	int	The output tensor channels. Defaults to 20.	20
embed_dim	int	The embedding dimension for PatchEmbed. Defaults to 768.	768
depth	int	Number of transformer depth. Defaults to 12.	12
mlp_ratio	float	Number of ratio used in MLP. Defaults to 4.0.	4.0
drop_rate	float	The drop ratio used in MLP. Defaults to 0.0.	0.0
drop_path_rate	float	The drop ratio used in DropPath. Defaults to 0.0.	0.0
num_blocks	int	Number of blocks. Defaults to 8.	8
sparsity_threshold	float	The value of threshold for softshrink. Defaults to 0.01.	0.01
hard_thresholding_fraction	float	The value of threshold for keep mode. Defaults to 1.0.	1.0
num_timestamps	int	Number of timestamp. Defaults to 1.	1

Examples:

>>> import ppsci
>>> model = ppsci.arch.AFNONet(("input", ), ("output", ))
>>> input_data = {"input": paddle.randn([1, 20, 720, 1440])}
>>> output_data = model(input_data)
>>> for k, v in output_data.items():
...     print(k, v.shape)
output [1, 20, 720, 1440]

Source code in ppsci/arch/afno.py

class AFNONet(base.Arch):
    """Adaptive Fourier Neural Network.

    Args:
        input_keys (Tuple[str, ...]): Name of input keys, such as ("input",).
        output_keys (Tuple[str, ...]): Name of output keys, such as ("output",).
        img_size (Tuple[int, ...], optional): Image size. Defaults to (720, 1440).
        patch_size (Tuple[int, ...], optional): Path. Defaults to (8, 8).
        in_channels (int, optional): The input tensor channels. Defaults to 20.
        out_channels (int, optional): The output tensor channels. Defaults to 20.
        embed_dim (int, optional): The embedding dimension for PatchEmbed. Defaults to 768.
        depth (int, optional): Number of transformer depth. Defaults to 12.
        mlp_ratio (float, optional): Number of ratio used in MLP. Defaults to 4.0.
        drop_rate (float, optional): The drop ratio used in MLP. Defaults to 0.0.
        drop_path_rate (float, optional): The drop ratio used in DropPath. Defaults to 0.0.
        num_blocks (int, optional): Number of blocks. Defaults to 8.
        sparsity_threshold (float, optional): The value of threshold for softshrink. Defaults to 0.01.
        hard_thresholding_fraction (float, optional): The value of threshold for keep mode. Defaults to 1.0.
        num_timestamps (int, optional): Number of timestamp. Defaults to 1.

    Examples:
        >>> import ppsci
        >>> model = ppsci.arch.AFNONet(("input", ), ("output", ))
        >>> input_data = {"input": paddle.randn([1, 20, 720, 1440])}
        >>> output_data = model(input_data)
        >>> for k, v in output_data.items():
        ...     print(k, v.shape)
        output [1, 20, 720, 1440]
    """

    def __init__(
        self,
        input_keys: Tuple[str, ...],
        output_keys: Tuple[str, ...],
        img_size: Tuple[int, ...] = (720, 1440),
        patch_size: Tuple[int, ...] = (8, 8),
        in_channels: int = 20,
        out_channels: int = 20,
        embed_dim: int = 768,
        depth: int = 12,
        mlp_ratio: float = 4.0,
        drop_rate: float = 0.0,
        drop_path_rate: float = 0.0,
        num_blocks: int = 8,
        sparsity_threshold: float = 0.01,
        hard_thresholding_fraction: float = 1.0,
        num_timestamps: int = 1,
    ):
        super().__init__()
        self.input_keys = input_keys
        self.output_keys = output_keys

        self.img_size = img_size
        self.patch_size = patch_size
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.embed_dim = embed_dim
        self.num_blocks = num_blocks
        self.num_timestamps = num_timestamps
        norm_layer = partial(nn.LayerNorm, epsilon=1e-6)

        self.patch_embed = PatchEmbed(
            img_size=img_size,
            patch_size=self.patch_size,
            in_channels=self.in_channels,
            embed_dim=embed_dim,
        )
        num_patches = self.patch_embed.num_patches

        data = paddle.zeros((1, num_patches, embed_dim))
        data = initializer.trunc_normal_(data, std=0.02)
        self.pos_embed = paddle.create_parameter(
            shape=data.shape,
            dtype=data.dtype,
            default_initializer=nn.initializer.Assign(data),
        )
        self.pos_drop = nn.Dropout(p=drop_rate)

        dpr = [x.item() for x in paddle.linspace(0, drop_path_rate, depth)]

        self.h = img_size[0] // self.patch_size[0]
        self.w = img_size[1] // self.patch_size[1]

        self.blocks = nn.LayerList(
            [
                Block(
                    dim=embed_dim,
                    mlp_ratio=mlp_ratio,
                    drop=drop_rate,
                    drop_path=dpr[i],
                    norm_layer=norm_layer,
                    num_blocks=self.num_blocks,
                    sparsity_threshold=sparsity_threshold,
                    hard_thresholding_fraction=hard_thresholding_fraction,
                )
                for i in range(depth)
            ]
        )

        self.norm = norm_layer(embed_dim)
        self.head = nn.Linear(
            embed_dim,
            self.out_channels * self.patch_size[0] * self.patch_size[1],
            bias_attr=False,
        )

        self.apply(self._init_weights)

    def _init_weights(self, m):
        if isinstance(m, nn.Linear):
            initializer.trunc_normal_(m.weight, std=0.02)
            if m.bias is not None:
                initializer.zeros_(m.bias)
        elif isinstance(m, nn.LayerNorm):
            initializer.ones_(m.weight)
            initializer.zeros_(m.bias)
        elif isinstance(m, nn.Conv2D):
            initializer.conv_init_(m)

    def forward_tensor(self, x):
        B = x.shape[0]
        x = self.patch_embed(x)
        x = x + self.pos_embed
        x = self.pos_drop(x)

        x = x.reshape((B, self.h, self.w, self.embed_dim))
        for block in self.blocks:
            x = block(x)

        x = self.head(x)

        b = x.shape[0]
        p1 = self.patch_size[0]
        p2 = self.patch_size[1]
        h = self.img_size[0] // self.patch_size[0]
        w = self.img_size[1] // self.patch_size[1]
        c_out = x.shape[3] // (p1 * p2)
        x = x.reshape((b, h, w, p1, p2, c_out))
        x = x.transpose((0, 5, 1, 3, 2, 4))
        x = x.reshape((b, c_out, h * p1, w * p2))

        return x

    @staticmethod
    def split_to_dict(data_tensors: Tuple[paddle.Tensor, ...], keys: Tuple[str, ...]):
        return {key: data_tensors[i] for i, key in enumerate(keys)}

    def forward(self, x):
        if self._input_transform is not None:
            x = self._input_transform(x)

        x_tensor = self.concat_to_tensor(x, self.input_keys)

        y = []
        input = x_tensor
        for _ in range(self.num_timestamps):
            out = self.forward_tensor(input)
            y.append(out)
            input = out
        y = self.split_to_dict(y, self.output_keys)

        if self._output_transform is not None:
            y = self._output_transform(x, y)
        return y

AMGNet

Bases: Layer

A Multi-scale Graph neural Network model based on Encoder-Process-Decoder structure for flow field prediction.

https://doi.org/10.1080/09540091.2022.2131737

Code reference: https://github.com/baoshiaijhin/amgnet

Parameters:

Name	Type	Description	Default
input_keys	Tuple[str, ...]	Name of input keys, such as ("input", ).	required
output_keys	Tuple[str, ...]	Name of output keys, such as ("pred", ).	required
input_dim	int	Number of input dimension.	required
output_dim	int	Number of output dimension.	required
latent_dim	int	Number of hidden(feature) dimension.	required
num_layers	int	Number of layer(s).	required
message_passing_aggregator	Literal['sum']	Message aggregator method in graph. Only "sum" available now.	required
message_passing_steps	int	Message passing steps in graph.	required
speed	str	Whether use vanilla method or fast method for graph_connectivity computation.	required

Examples:

>>> import ppsci
>>> model = ppsci.arch.AMGNet(
...     ("input", ), ("pred", ), 5, 3, 64, 2, "sum", 6, "norm",
... )

Source code in ppsci/arch/amgnet.py

class AMGNet(nn.Layer):
    """A Multi-scale Graph neural Network model
    based on Encoder-Process-Decoder structure for flow field prediction.

    https://doi.org/10.1080/09540091.2022.2131737

    Code reference: https://github.com/baoshiaijhin/amgnet

    Args:
        input_keys (Tuple[str, ...]): Name of input keys, such as ("input", ).
        output_keys (Tuple[str, ...]): Name of output keys, such as ("pred", ).
        input_dim (int): Number of input dimension.
        output_dim (int): Number of output dimension.
        latent_dim (int): Number of hidden(feature) dimension.
        num_layers (int): Number of layer(s).
        message_passing_aggregator (Literal["sum"]): Message aggregator method in graph.
            Only "sum" available now.
        message_passing_steps (int): Message passing steps in graph.
        speed (str): Whether use vanilla method or fast method for graph_connectivity
            computation.

    Examples:
        >>> import ppsci
        >>> model = ppsci.arch.AMGNet(
        ...     ("input", ), ("pred", ), 5, 3, 64, 2, "sum", 6, "norm",
        ... )
    """

    def __init__(
        self,
        input_keys: Tuple[str, ...],
        output_keys: Tuple[str, ...],
        input_dim: int,
        output_dim: int,
        latent_dim: int,
        num_layers: int,
        message_passing_aggregator: Literal["sum"],
        message_passing_steps: int,
        speed: Literal["norm", "fast"],
    ):
        super().__init__()
        self.input_keys = input_keys
        self.output_keys = output_keys
        self._latent_dim = latent_dim
        self.speed = speed
        self._output_dim = output_dim
        self._num_layers = num_layers

        self.encoder = Encoder(input_dim, self._make_mlp, latent_dim=self._latent_dim)
        self.processor = Processor(
            make_mlp=self._make_mlp,
            output_dim=self._latent_dim,
            message_passing_steps=message_passing_steps,
            message_passing_aggregator=message_passing_aggregator,
            use_stochastic_message_passing=False,
        )
        self.post_processor = self._make_mlp(self._latent_dim, 128)
        self.decoder = Decoder(
            make_mlp=functools.partial(self._make_mlp, layer_norm=False),
            output_dim=self._output_dim,
        )

    def forward(self, x: Dict[str, "pgl.Graph"]) -> Dict[str, paddle.Tensor]:
        graphs = x[self.input_keys[0]]
        latent_graph = self.encoder(graphs)
        x, p = self.processor(latent_graph, speed=self.speed)
        node_features = self._spa_compute(x, p)
        pred_field = self.decoder(node_features)
        return {self.output_keys[0]: pred_field}

    def _make_mlp(self, output_dim: int, input_dim: int = 5, layer_norm: bool = True):
        widths = (self._latent_dim,) * self._num_layers + (output_dim,)
        network = FullyConnectedLayer(input_dim, widths)
        if layer_norm:
            network = nn.Sequential(network, nn.LayerNorm(normalized_shape=widths[-1]))
        return network

    def _spa_compute(self, x: List["pgl.Graph"], p):
        j = len(x) - 1
        node_features = x[j].x

        for k in range(1, j + 1):
            pos = p[-k]
            fine_nodes = x[-(k + 1)].pos
            feature = _knn_interpolate(node_features, pos, fine_nodes)
            node_features = x[-(k + 1)].x + feature
            node_features = self.post_processor(node_features)

        return node_features

AutoEncoder

Bases: Arch

AutoEncoder is a class that represents an autoencoder neural network model.

Parameters:

Name	Type	Description	Default
input_keys	Tuple[str, ...]	A tuple of input keys.	required
output_keys	Tuple[str, ...]	A tuple of output keys.	required
input_dim	int	The dimension of the input data.	required
latent_dim	int	The dimension of the latent space.	required
hidden_dim	int	The dimension of the hidden layer.	required

Examples:

>>> import paddle
>>> import ppsci
>>> model = ppsci.arch.AutoEncoder(
...    input_keys=("input1",),
...    output_keys=("mu", "log_sigma", "decoder_z",),
...    input_dim=100,
...    latent_dim=50,
...    hidden_dim=200
... )
>>> input_dict = {"input1": paddle.rand([200, 100]),}
>>> output_dict = model(input_dict)
>>> print(output_dict["mu"].shape)
[200, 50]
>>> print(output_dict["log_sigma"].shape)
[200, 50]
>>> print(output_dict["decoder_z"].shape)
[200, 100]

Source code in ppsci/arch/vae.py

class AutoEncoder(base.Arch):
    """
    AutoEncoder is a class that represents an autoencoder neural network model.

    Args:
        input_keys (Tuple[str, ...]): A tuple of input keys.
        output_keys (Tuple[str, ...]): A tuple of output keys.
        input_dim (int): The dimension of the input data.
        latent_dim (int): The dimension of the latent space.
        hidden_dim (int): The dimension of the hidden layer.

    Examples:
        >>> import paddle
        >>> import ppsci
        >>> model = ppsci.arch.AutoEncoder(
        ...    input_keys=("input1",),
        ...    output_keys=("mu", "log_sigma", "decoder_z",),
        ...    input_dim=100,
        ...    latent_dim=50,
        ...    hidden_dim=200
        ... )
        >>> input_dict = {"input1": paddle.rand([200, 100]),}
        >>> output_dict = model(input_dict)
        >>> print(output_dict["mu"].shape)
        [200, 50]
        >>> print(output_dict["log_sigma"].shape)
        [200, 50]
        >>> print(output_dict["decoder_z"].shape)
        [200, 100]
    """

    def __init__(
        self,
        input_keys: Tuple[str, ...],
        output_keys: Tuple[str, ...],
        input_dim: int,
        latent_dim: int,
        hidden_dim: int,
    ):
        super(AutoEncoder, self).__init__()
        self.input_keys = input_keys
        self.output_keys = output_keys
        # encoder
        self._encoder_linear = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.Tanh(),
        )
        self._encoder_mu = nn.Linear(hidden_dim, latent_dim)
        self._encoder_log_sigma = nn.Linear(hidden_dim, latent_dim)

        self._decoder = nn.Sequential(
            nn.Linear(latent_dim, hidden_dim),
            nn.Tanh(),
            nn.Linear(hidden_dim, input_dim),
        )

    def encoder(self, x):
        h = self._encoder_linear(x)
        mu = self._encoder_mu(h)
        log_sigma = self._encoder_log_sigma(h)
        return mu, log_sigma

    def decoder(self, x):
        return self._decoder(x)

    def forward_tensor(self, x):
        mu, log_sigma = self.encoder(x)
        z = mu + paddle.randn(mu.shape) * paddle.exp(log_sigma)
        return mu, log_sigma, self.decoder(z)

    def forward(self, x):
        x = self.concat_to_tensor(x, self.input_keys, axis=-1)
        mu, log_sigma, decoder_z = self.forward_tensor(x)
        result_dict = {
            self.output_keys[0]: mu,
            self.output_keys[1]: log_sigma,
            self.output_keys[2]: decoder_z,
        }
        return result_dict

ChipDeepONets

Bases: Arch

Multi-branch physics-informed deep operator neural network. The network consists of three branch networks: random heat source, boundary function, and boundary type, as well as a trunk network.

Parameters:

Name	Type	Description	Default
branch_input_keys	Tuple[str, ...]	Name of input data for internal heat source on branch nets.	required
BCtype_input_keys	Tuple[str, ...]	Name of input data for boundary types on branch nets.	required
BC_input_keys	Tuple[str, ...]	Name of input data for boundary on branch nets.	required
trunk_input_keys	Tuple[str, ...]	Name of input data for trunk net.	required
output_keys	Tuple[str, ...]	Output name of predicted temperature.	required
num_loc	int	Number of sampled input data for internal heat source.	required
bctype_loc	int	Number of sampled input data for boundary types.	required
BC_num_loc	int	Number of sampled input data for boundary.	required
num_features	int	Number of features extracted from trunk net, same for all branch nets.	required
branch_num_layers	int	Number of hidden layers of internal heat source on branch nets.	required
BC_num_layers	int	Number of hidden layers of boundary on branch nets.	required
trunk_num_layers	int	Number of hidden layers of trunk net.	required
branch_hidden_size	Union[int, Tuple[int, ...]]	Number of hidden size of internal heat source on branch nets. An integer for all layers, or list of integer specify each layer's size.	required
BC_hidden_size	Union[int, Tuple[int, ...]]	Number of hidden size of boundary on branch nets. An integer for all layers, or list of integer specify each layer's size.	required
trunk_hidden_size	Union[int, Tuple[int, ...]]	Number of hidden size of trunk net. An integer for all layers, or list of integer specify each layer's size.	required
branch_skip_connection	bool	Whether to use skip connection for internal heat source on branch net. Defaults to False.	False
BC_skip_connection	bool	Whether to use skip connection for boundary on branch net. Defaults to False.	False
trunk_skip_connection	bool	Whether to use skip connection for trunk net. Defaults to False.	False
branch_activation	str	Name of activation function for internal heat source on branch net. Defaults to "tanh".	'tanh'
BC_activation	str	Name of activation function for boundary on branch net. Defaults to "tanh".	'tanh'
trunk_activation	str	Name of activation function for trunk net. Defaults to "tanh".	'tanh'
branch_weight_norm	bool	Whether to apply weight norm on parameter(s) for internal heat source on branch net. Defaults to False.	False
BC_weight_norm	bool	Whether to apply weight norm on parameter(s) for boundary on branch net. Defaults to False.	False
trunk_weight_norm	bool	Whether to apply weight norm on parameter(s) for trunk net. Defaults to False.	False
use_bias	bool	Whether to add bias on predicted G(u)(y). Defaults to True.	True

Examples:

>>> import ppsci
>>> model = ppsci.arch.ChipDeepONets(
...     ('u',),
...     ('bc',),
...     ('bc_data',),
...     ("x",'y'),
...     ("T",),
...     324,
...     1,
...     76,
...     400,
...     9,
...     9,
...     6,
...     256,
...     256,
...     128,
...     branch_activation="swish",
...     BC_activation="swish",
...     trunk_activation="swish",
...     use_bias=True,
... )

Source code in ppsci/arch/chip_deeponets.py

class ChipDeepONets(base.Arch):
    """Multi-branch physics-informed deep operator neural network. The network consists of three branch networks: random heat source, boundary function, and boundary type, as well as a trunk network.

    Args:
        branch_input_keys (Tuple[str, ...]): Name of input data for internal heat source on branch nets.
        BCtype_input_keys (Tuple[str, ...]): Name of input data for boundary types on branch nets.
        BC_input_keys (Tuple[str, ...]): Name of input data for boundary on branch nets.
        trunk_input_keys (Tuple[str, ...]): Name of input data for trunk net.
        output_keys (Tuple[str, ...]): Output name of predicted temperature.
        num_loc (int): Number of sampled input data for internal heat source.
        bctype_loc (int): Number of sampled input data for boundary types.
        BC_num_loc (int): Number of sampled input data for boundary.
        num_features (int): Number of features extracted from trunk net, same for all branch nets.
        branch_num_layers (int): Number of hidden layers of internal heat source on branch nets.
        BC_num_layers (int): Number of hidden layers of boundary on branch nets.
        trunk_num_layers (int): Number of hidden layers of trunk net.
        branch_hidden_size (Union[int, Tuple[int, ...]]): Number of hidden size of internal heat source on branch nets.
            An integer for all layers, or list of integer specify each layer's size.
        BC_hidden_size (Union[int, Tuple[int, ...]]): Number of hidden size of boundary on branch nets.
            An integer for all layers, or list of integer specify each layer's size.
        trunk_hidden_size (Union[int, Tuple[int, ...]]): Number of hidden size of trunk net.
            An integer for all layers, or list of integer specify each layer's size.
        branch_skip_connection (bool, optional): Whether to use skip connection for internal heat source on branch net. Defaults to False.
        BC_skip_connection (bool, optional): Whether to use skip connection for boundary on branch net. Defaults to False.
        trunk_skip_connection (bool, optional): Whether to use skip connection for trunk net. Defaults to False.
        branch_activation (str, optional): Name of activation function for internal heat source on branch net. Defaults to "tanh".
        BC_activation (str, optional): Name of activation function for boundary on branch net. Defaults to "tanh".
        trunk_activation (str, optional): Name of activation function for trunk net. Defaults to "tanh".
        branch_weight_norm (bool, optional): Whether to apply weight norm on parameter(s) for internal heat source on branch net. Defaults to False.
        BC_weight_norm (bool, optional): Whether to apply weight norm on parameter(s) for boundary on branch net. Defaults to False.
        trunk_weight_norm (bool, optional): Whether to apply weight norm on parameter(s) for trunk net. Defaults to False.
        use_bias (bool, optional): Whether to add bias on predicted G(u)(y). Defaults to True.

    Examples:
        >>> import ppsci
        >>> model = ppsci.arch.ChipDeepONets(
        ...     ('u',),
        ...     ('bc',),
        ...     ('bc_data',),
        ...     ("x",'y'),
        ...     ("T",),
        ...     324,
        ...     1,
        ...     76,
        ...     400,
        ...     9,
        ...     9,
        ...     6,
        ...     256,
        ...     256,
        ...     128,
        ...     branch_activation="swish",
        ...     BC_activation="swish",
        ...     trunk_activation="swish",
        ...     use_bias=True,
        ... )
    """

    def __init__(
        self,
        branch_input_keys: Tuple[str, ...],
        BCtype_input_keys: Tuple[str, ...],
        BC_input_keys: Tuple[str, ...],
        trunk_input_keys: Tuple[str, ...],
        output_keys: Tuple[str, ...],
        num_loc: int,
        bctype_loc: int,
        BC_num_loc: int,
        num_features: int,
        branch_num_layers: int,
        BC_num_layers: int,
        trunk_num_layers: int,
        branch_hidden_size: Union[int, Tuple[int, ...]],
        BC_hidden_size: Union[int, Tuple[int, ...]],
        trunk_hidden_size: Union[int, Tuple[int, ...]],
        branch_skip_connection: bool = False,
        BC_skip_connection: bool = False,
        trunk_skip_connection: bool = False,
        branch_activation: str = "tanh",
        BC_activation: str = "tanh",
        trunk_activation: str = "tanh",
        branch_weight_norm: bool = False,
        BC_weight_norm: bool = False,
        trunk_weight_norm: bool = False,
        use_bias: bool = True,
    ):
        super().__init__()
        self.trunk_input_keys = trunk_input_keys
        self.branch_input_keys = branch_input_keys
        self.BCtype_input_keys = BCtype_input_keys
        self.BC_input_keys = BC_input_keys
        self.input_keys = (
            self.trunk_input_keys
            + self.branch_input_keys
            + self.BC_input_keys
            + self.BCtype_input_keys
        )
        self.output_keys = output_keys

        self.branch_net = mlp.MLP(
            self.branch_input_keys,
            ("b",),
            branch_num_layers,
            branch_hidden_size,
            branch_activation,
            branch_skip_connection,
            branch_weight_norm,
            input_dim=num_loc,
            output_dim=num_features,
        )

        self.BCtype_net = mlp.MLP(
            self.BCtype_input_keys,
            ("bctype",),
            BC_num_layers,
            BC_hidden_size,
            BC_activation,
            BC_skip_connection,
            BC_weight_norm,
            input_dim=bctype_loc,
            output_dim=num_features,
        )

        self.BC_net = mlp.MLP(
            self.BC_input_keys,
            ("bc",),
            BC_num_layers,
            BC_hidden_size,
            BC_activation,
            BC_skip_connection,
            BC_weight_norm,
            input_dim=BC_num_loc,
            output_dim=num_features,
        )

        self.trunk_net = mlp.MLP(
            self.trunk_input_keys,
            ("t",),
            trunk_num_layers,
            trunk_hidden_size,
            trunk_activation,
            trunk_skip_connection,
            trunk_weight_norm,
            input_dim=len(self.trunk_input_keys),
            output_dim=num_features,
        )
        self.trunk_act = act_mod.get_activation(trunk_activation)
        self.bc_act = act_mod.get_activation(BC_activation)
        self.branch_act = act_mod.get_activation(branch_activation)

        self.use_bias = use_bias
        if use_bias:
            # register bias to parameter for updating in optimizer and storage
            self.b = self.create_parameter(
                shape=(1,),
                attr=nn.initializer.Constant(0.0),
            )

    def forward(self, x):

        if self._input_transform is not None:
            x = self._input_transform(x)

        # Branch net to encode the input function
        u_features = self.branch_net(x)[self.branch_net.output_keys[0]]
        bc_features = self.BC_net(x)[self.BC_net.output_keys[0]]
        bctype_features = self.BCtype_net(x)[self.BCtype_net.output_keys[0]]
        # Trunk net to encode the domain of the output function
        y_features = self.trunk_net(x)[self.trunk_net.output_keys[0]]
        y_features = self.trunk_act(y_features)
        # Dot product
        G_u = paddle.sum(
            u_features * y_features * bc_features * bctype_features,
            axis=1,
            keepdim=True,
        )
        # Add bias
        if self.use_bias:
            G_u += self.b

        result_dict = {
            self.output_keys[0]: G_u,
        }
        if self._output_transform is not None:
            result_dict = self._output_transform(x, result_dict)

        return result_dict

Climateformer

Bases: Arch

Climateformer is a class that represents a Spatial-Temporal Transformer model designed for climate prediction with multiple meteorological variables.

Parameters:

Name	Type	Description	Default
input_keys	Tuple[str, ...]	A tuple of input keys.	required
output_keys	Tuple[str, ...]	A tuple of output keys.	required
shape_in	Tuple[int, ...]	The shape of the input data (T, C, H, W), where T is the number of time steps, C is the number of channels, H and W are the spatial dimensions.	required
hid_S	int	The number of hidden channels in the spatial encoder.	64
hid_T	int	The number of hidden units in the temporal encoder.	256
N_S	int	The number of spatial transformer layers.	4
N_T	int	The number of temporal transformer layers.	4
incep_ker	Tuple[int, ...]	The kernel sizes used in the inception block.	(3, 5, 7, 11)
groups	int	The number of groups for grouped convolutions.	8
num_classes	int	The number of predicted meteorological variables.	12

Examples:

>>> import paddle
>>> import ppsci
>>> model = ppsci.arch.Climateformer(
...     input_keys=("input",),
...     output_keys=("output",),
...     shape_in=(6, 12, 192, 256),
...     hid_S=64,
...     hid_T=256,
...     N_S=4,
...     N_T=4,
...     incep_ker=(3, 5, 7, 11),
...     groups=8,
...     num_classes=4,
... )
>>> input_dict = {"input": paddle.rand([8, 6, 12, 192, 256])}
>>> output_dict = model(input_dict)
>>> print(output_dict["output"].shape)
[8, 6, 12, 192, 256]

Source code in ppsci/arch/climateformer.py

class Climateformer(base.Arch):
    """
    Climateformer is a class that represents a Spatial-Temporal Transformer model designed for climate prediction with multiple meteorological variables.

    Args:
        input_keys (Tuple[str, ...]): A tuple of input keys.
        output_keys (Tuple[str, ...]): A tuple of output keys.
        shape_in (Tuple[int, ...]): The shape of the input data (T, C, H, W), where
            T is the number of time steps, C is the number of channels,
            H and W are the spatial dimensions.
        hid_S (int): The number of hidden channels in the spatial encoder.
        hid_T (int): The number of hidden units in the temporal encoder.
        N_S (int): The number of spatial transformer layers.
        N_T (int): The number of temporal transformer layers.
        incep_ker (Tuple[int, ...]): The kernel sizes used in the inception block.
        groups (int): The number of groups for grouped convolutions.
        num_classes (int): The number of predicted meteorological variables.

    Examples:
        >>> import paddle
        >>> import ppsci
        >>> model = ppsci.arch.Climateformer(
        ...     input_keys=("input",),
        ...     output_keys=("output",),
        ...     shape_in=(6, 12, 192, 256),
        ...     hid_S=64,
        ...     hid_T=256,
        ...     N_S=4,
        ...     N_T=4,
        ...     incep_ker=(3, 5, 7, 11),
        ...     groups=8,
        ...     num_classes=4,
        ... )
        >>> input_dict = {"input": paddle.rand([8, 6, 12, 192, 256])}
        >>> output_dict = model(input_dict)
        >>> print(output_dict["output"].shape)
        [8, 6, 12, 192, 256]
    """

    def __init__(
        self,
        input_keys: Tuple[str, ...],
        output_keys: Tuple[str, ...],
        shape_in: Tuple[int, ...],
        hid_S: int = 64,
        hid_T: int = 256,
        N_S: int = 4,
        N_T: int = 4,
        incep_ker: Tuple[int, ...] = (3, 5, 7, 11),
        groups: int = 8,
        num_classes: int = 12,
    ):
        super().__init__()
        self.input_keys = input_keys
        self.output_keys = output_keys
        self.num_classes = num_classes

        T, C, H, W = shape_in
        self.enc = Encoder(C, hid_S, N_S)
        self.hid1 = MidXnet(T * hid_S, hid_T // 2, N_T, incep_ker, groups)
        self.dec = Decoder(T * hid_S, T * self.num_classes, N_S)

    def forward(self, x):
        if self._input_transform is not None:
            x = self._input_transform(x)

        x = self.concat_to_tensor(x, self.input_keys)

        B, T, C, H, W = x.shape
        x = x.reshape([B * T, C, H, W])

        # encoded
        embed = self.enc(x)
        _, C_4, H_4, W_4 = embed[-1].shape

        # translator
        z = embed[-1].reshape([B, T, C_4, H_4, W_4])
        hid = self.hid1(z)
        hid = hid.transpose(perm=[0, 2, 1]).reshape([B, -1, H_4, W_4])

        # decoded
        y = self.dec(hid, embed[0])
        y = y.reshape([B, T, self.num_classes, H, W])

        y = self.split_to_dict(y, self.output_keys)
        if self._output_transform is not None:
            y = self._output_transform(x, y)

        return y  # {self.output_keys[0]: Y}

CrystalGraphConvNet

Bases: Arch

Create a crystal graph convolutional neural network for predicting total material properties.

Parameters:

Name	Type	Description	Default
orig_atom_fea_len	int	Number of atom features in the input.	required
nbr_fea_len	int	Number of bond features.	required
atom_fea_len	int	Number of hidden atom features in the convolutional layers.	required
n_conv	int	Number of convolutional layers.	required
h_fea_len	int	Number of hidden features after pooling.	required
n_h	int	Number of hidden layers after pooling.	required

Examples:

>>> import paddle
>>> import ppsci
>>> model = ppsci.arch.CrystalGraphConvNet(
...     orig_atom_fea_len=92,
...     nbr_fea_len=41,
...     atom_fea_len=64,
...     n_conv=3,
...     h_fea_len=128,
...     n_h=1,
... )
>>> input_dict = {
...     "i": [
...         paddle.rand(shape=[45, 92]), paddle.rand(shape=[45, 12, 41]),
...         paddle.randint(high=45, shape=[45, 12]),
...         [
...             paddle.randint(high=32, shape=[32]), paddle.randint(high=8, shape=[8]),
...             paddle.randint(high=2, shape=[2]), paddle.randint(high=3, shape=[3])
...         ]
...     ]
... }
>>> output_dict = model(input_dict)
>>> print(output_dict["out"].shape)
[4, 1]

Source code in ppsci/arch/crystalgraphconvnet.py

class CrystalGraphConvNet(base.Arch):
    """
    Create a crystal graph convolutional neural network for predicting total
    material properties.

    Args:
        orig_atom_fea_len (int): Number of atom features in the input.
        nbr_fea_len (int): Number of bond features.
        atom_fea_len (int): Number of hidden atom features in the convolutional layers.
        n_conv (int): Number of convolutional layers.
        h_fea_len (int): Number of hidden features after pooling.
        n_h (int): Number of hidden layers after pooling.

    Examples:
         >>> import paddle
         >>> import ppsci
         >>> model = ppsci.arch.CrystalGraphConvNet(
         ...     orig_atom_fea_len=92,
         ...     nbr_fea_len=41,
         ...     atom_fea_len=64,
         ...     n_conv=3,
         ...     h_fea_len=128,
         ...     n_h=1,
         ... )
         >>> input_dict = {
         ...     "i": [
         ...         paddle.rand(shape=[45, 92]), paddle.rand(shape=[45, 12, 41]),
         ...         paddle.randint(high=45, shape=[45, 12]),
         ...         [
         ...             paddle.randint(high=32, shape=[32]), paddle.randint(high=8, shape=[8]),
         ...             paddle.randint(high=2, shape=[2]), paddle.randint(high=3, shape=[3])
         ...         ]
         ...     ]
         ... }
         >>> output_dict = model(input_dict)
         >>> print(output_dict["out"].shape)
         [4, 1]
    """

    def __init__(
        self,
        orig_atom_fea_len: int,
        nbr_fea_len: int,
        atom_fea_len: int,
        n_conv: int,
        h_fea_len: int,
        n_h: int,
    ):

        super().__init__()
        self.embedding = nn.Linear(orig_atom_fea_len, atom_fea_len)
        self.convs = nn.LayerList(
            [
                ConvLayer(atom_fea_len=atom_fea_len, nbr_fea_len=nbr_fea_len)
                for _ in range(n_conv)
            ]
        )
        self.conv_to_fc = nn.Linear(atom_fea_len, h_fea_len)
        self.conv_to_fc_softplus = nn.Softplus()
        if n_h > 1:
            self.fcs = nn.LayerList(
                [nn.Linear(h_fea_len, h_fea_len) for _ in range(n_h - 1)]
            )
            self.softpluses = nn.LayerList([nn.Softplus() for _ in range(n_h - 1)])

        self.fc_out = nn.Linear(h_fea_len, 1)

    def forward(self, input) -> paddle.Tensor:
        """
        Forward pass.

        N: Total number of atoms in the batch.
        M: Max number of neighbors.
        N0: Total number of crystals in the batch.

        Args:
            input (list): List of input, which includes the following elements:
                - atom_fea (paddle.Tensor): Shape (N, orig_atom_fea_len). Atom features from atom type.
                - nbr_fea (paddle.Tensor): Shape (N, M, nbr_fea_len). Bond features of each atom's M neighbors.
                - nbr_fea_idx (paddle.Tensor): Shape (N, M). Indices of M neighbors of each atom.
                - crystal_atom_idx (list): List of paddle.Tensor of length N0. Mapping from the crystal idx to atom idx.

        Returns:
            paddle.Tensor: Shape (N,). Atom hidden features after convolution.
        """
        atom_fea, nbr_fea, nbr_fea_idx, crystal_atom_idx = input["i"]
        atom_fea = self.embedding(atom_fea)
        for conv_func in self.convs:
            atom_fea = conv_func(atom_fea, nbr_fea, nbr_fea_idx)
        crys_fea = self.pooling(atom_fea, crystal_atom_idx)
        crys_fea = self.conv_to_fc(self.conv_to_fc_softplus(crys_fea))
        crys_fea = self.conv_to_fc_softplus(crys_fea)
        if hasattr(self, "fcs") and hasattr(self, "softpluses"):
            for fc, softplus in zip(self.fcs, self.softpluses):
                crys_fea = softplus(fc(crys_fea))
        out = self.fc_out(crys_fea)
        out_dict = {"out": out}
        return out_dict

    def pooling(self, atom_fea, crystal_atom_idx):
        """
        Pooling the atom features to crystal features

        N: Total number of atoms in the batch
        N0: Total number of crystals in the batch

        Args:
            atom_fea (paddle.Tensor): Shape (N, atom_fea_len). Atom feature vectors of the batch.
            crystal_atom_idx (List[paddle.Tensor]): Length N0. Mapping from the crystal idx to atom idx
        """
        assert (
            sum([len(idx_map) for idx_map in crystal_atom_idx])
            == atom_fea.data.shape[0]
        )
        summed_fea = [
            paddle.mean(atom_fea[idx_map], axis=0, keepdim=True)
            for idx_map in crystal_atom_idx
        ]
        return paddle.concat(summed_fea, axis=0)

forward(input)

Forward pass.

N: Total number of atoms in the batch. M: Max number of neighbors. N0: Total number of crystals in the batch.

Parameters:

Name	Type	Description	Default
input	list	List of input, which includes the following elements: - atom_fea (paddle.Tensor): Shape (N, orig_atom_fea_len). Atom features from atom type. - nbr_fea (paddle.Tensor): Shape (N, M, nbr_fea_len). Bond features of each atom's M neighbors. - nbr_fea_idx (paddle.Tensor): Shape (N, M). Indices of M neighbors of each atom. - crystal_atom_idx (list): List of paddle.Tensor of length N0. Mapping from the crystal idx to atom idx.	required

Returns:

Type	Description
Tensor	paddle.Tensor: Shape (N,). Atom hidden features after convolution.

Source code in ppsci/arch/crystalgraphconvnet.py

def forward(self, input) -> paddle.Tensor:
    """
    Forward pass.

    N: Total number of atoms in the batch.
    M: Max number of neighbors.
    N0: Total number of crystals in the batch.

    Args:
        input (list): List of input, which includes the following elements:
            - atom_fea (paddle.Tensor): Shape (N, orig_atom_fea_len). Atom features from atom type.
            - nbr_fea (paddle.Tensor): Shape (N, M, nbr_fea_len). Bond features of each atom's M neighbors.
            - nbr_fea_idx (paddle.Tensor): Shape (N, M). Indices of M neighbors of each atom.
            - crystal_atom_idx (list): List of paddle.Tensor of length N0. Mapping from the crystal idx to atom idx.

    Returns:
        paddle.Tensor: Shape (N,). Atom hidden features after convolution.
    """
    atom_fea, nbr_fea, nbr_fea_idx, crystal_atom_idx = input["i"]
    atom_fea = self.embedding(atom_fea)
    for conv_func in self.convs:
        atom_fea = conv_func(atom_fea, nbr_fea, nbr_fea_idx)
    crys_fea = self.pooling(atom_fea, crystal_atom_idx)
    crys_fea = self.conv_to_fc(self.conv_to_fc_softplus(crys_fea))
    crys_fea = self.conv_to_fc_softplus(crys_fea)
    if hasattr(self, "fcs") and hasattr(self, "softpluses"):
        for fc, softplus in zip(self.fcs, self.softpluses):
            crys_fea = softplus(fc(crys_fea))
    out = self.fc_out(crys_fea)
    out_dict = {"out": out}
    return out_dict

pooling(atom_fea, crystal_atom_idx)

Pooling the atom features to crystal features

N: Total number of atoms in the batch N0: Total number of crystals in the batch

Parameters:

Name	Type	Description	Default
atom_fea	Tensor	Shape (N, atom_fea_len). Atom feature vectors of the batch.	required
crystal_atom_idx	List[Tensor]	Length N0. Mapping from the crystal idx to atom idx	required

Source code in ppsci/arch/crystalgraphconvnet.py

def pooling(self, atom_fea, crystal_atom_idx):
    """
    Pooling the atom features to crystal features

    N: Total number of atoms in the batch
    N0: Total number of crystals in the batch

    Args:
        atom_fea (paddle.Tensor): Shape (N, atom_fea_len). Atom feature vectors of the batch.
        crystal_atom_idx (List[paddle.Tensor]): Length N0. Mapping from the crystal idx to atom idx
    """
    assert (
        sum([len(idx_map) for idx_map in crystal_atom_idx])
        == atom_fea.data.shape[0]
    )
    summed_fea = [
        paddle.mean(atom_fea[idx_map], axis=0, keepdim=True)
        for idx_map in crystal_atom_idx
    ]
    return paddle.concat(summed_fea, axis=0)

CuboidTransformer

Bases: Arch

Cuboid Transformer for spatiotemporal forecasting

We adopt the Non-autoregressive encoder-decoder architecture. The decoder takes the multi-scale memory output from the encoder.

The initial downsampling / upsampling layers will be Downsampling: [K x Conv2D --> PatchMerge] Upsampling: [Nearest Interpolation-based Upsample --> K x Conv2D]

x --> downsample (optional) ---> (+pos_embed) ---> enc --> mem_l initial_z (+pos_embed) ---> FC | | |------------| | | y <--- upsample (optional) <--- dec <----------

Parameters:

Name	Type	Description	Default
input_keys	Tuple[str, ...]	Name of input keys, such as ("input",).	required
output_keys	Tuple[str, ...]	Name of output keys, such as ("output",).	required
input_shape	Tuple[int, ...]	The shape of the input data.	required
target_shape	Tuple[int, ...]	The shape of the target data.	required
base_units	int	The base units. Defaults to 128.	128
block_units	int	The block units. Defaults to None.	None
scale_alpha	float	We scale up the channels based on the formula: - round_to(base_units * max(downsample_scale) ** units_alpha, 4). Defaults to 1.0.	1.0
num_heads	int	The number of heads. Defaults to 4.	4
attn_drop	float	The attention dropout. Defaults to 0.0.	0.0
proj_drop	float	The projection dropout. Defaults to 0.0.	0.0
ffn_drop	float	The ffn dropout. Defaults to 0.0.	0.0
downsample	int	The rate of downsample. Defaults to 2.	2
downsample_type	str	The type of downsample. Defaults to "patch_merge".	'patch_merge'
upsample_type	str	The rate of upsample. Defaults to "upsample".	'upsample'
upsample_kernel_size	int	The kernel size of upsample. Defaults to 3.	3
enc_depth	list	The depth of encoder. Defaults to [4, 4, 4].	[4, 4, 4]
enc_attn_patterns	str	The pattern of encoder attention. Defaults to None.	None
enc_cuboid_size	list	The cuboid size of encoder. Defaults to [(4, 4, 4), (4, 4, 4)].	[(4, 4, 4), (4, 4, 4)]
enc_cuboid_strategy	list	The cuboid strategy of encoder. Defaults to [("l", "l", "l"), ("d", "d", "d")].	[('l', 'l', 'l'), ('d', 'd', 'd')]
enc_shift_size	list	The shift size of encoder. Defaults to [(0, 0, 0), (0, 0, 0)].	[(0, 0, 0), (0, 0, 0)]
enc_use_inter_ffn	bool	Whether to use intermediate FFN for encoder. Defaults to True.	True
dec_depth	list	The depth of decoder. Defaults to [2, 2].	[2, 2]
dec_cross_start	int	The cross start of decoder. Defaults to 0.	0
dec_self_attn_patterns	str	The partterns of decoder. Defaults to None.	None
dec_self_cuboid_size	list	The cuboid size of decoder. Defaults to [(4, 4, 4), (4, 4, 4)].	[(4, 4, 4), (4, 4, 4)]
dec_self_cuboid_strategy	list	The strategy of decoder. Defaults to [("l", "l", "l"), ("d", "d", "d")].	[('l', 'l', 'l'), ('d', 'd', 'd')]
dec_self_shift_size	list	The shift size of decoder. Defaults to [(1, 1, 1), (0, 0, 0)].	[(1, 1, 1), (0, 0, 0)]
dec_cross_attn_patterns	_type_	The cross attention patterns of decoder. Defaults to None.	None
dec_cross_cuboid_hw	list	The cuboid_hw of decoder. Defaults to [(4, 4), (4, 4)].	[(4, 4), (4, 4)]
dec_cross_cuboid_strategy	list	The cuboid strategy of decoder. Defaults to [("l", "l", "l"), ("d", "l", "l")].	[('l', 'l', 'l'), ('d', 'l', 'l')]
dec_cross_shift_hw	list	The shift_hw of decoder. Defaults to [(0, 0), (0, 0)].	[(0, 0), (0, 0)]
dec_cross_n_temporal	list	The cross_n_temporal of decoder. Defaults to [1, 2].	[1, 2]
dec_cross_last_n_frames	int	The cross_last_n_frames of decoder. Defaults to None.	None
dec_use_inter_ffn	bool	Whether to use intermediate FFN for decoder. Defaults to True.	True
dec_hierarchical_pos_embed	bool	Whether to use hierarchical pos_embed for decoder. Defaults to False.	False
num_global_vectors	int	The num of global vectors. Defaults to 4.	4
use_dec_self_global	bool	Whether to use global vector for decoder. Defaults to True.	True
dec_self_update_global	bool	Whether to update global vector for decoder. Defaults to True.	True
use_dec_cross_global	bool	Whether to use cross global vector for decoder. Defaults to True.	True
use_global_vector_ffn	bool	Whether to use global vector FFN. Defaults to True.	True
use_global_self_attn	bool	Whether to use global attentions. Defaults to False.	False
separate_global_qkv	bool	Whether to separate global qkv. Defaults to False.	False
global_dim_ratio	int	The ratio of global dim. Defaults to 1.	1
self_pattern	str	The pattern. Defaults to "axial".	'axial'
cross_self_pattern	str	The self cross pattern. Defaults to "axial".	'axial'
cross_pattern	str	The cross pattern. Defaults to "cross_1x1".	'cross_1x1'
z_init_method	str	How the initial input to the decoder is initialized. Defaults to "nearest_interp".	'nearest_interp'
initial_downsample_type	str	The downsample type of initial. Defaults to "conv".	'conv'
initial_downsample_activation	str	The downsample activation of initial. Defaults to "leaky".	'leaky'
initial_downsample_scale	int	The downsample scale of initial. Defaults to 1.	1
initial_downsample_conv_layers	int	The conv layer of downsample of initial. Defaults to 2.	2
final_upsample_conv_layers	int	The conv layer of final upsample. Defaults to 2.	2
initial_downsample_stack_conv_num_layers	int	The num of stack conv layer of initial downsample. Defaults to 1.	1
initial_downsample_stack_conv_dim_list	list	The dim list of stack conv of initial downsample. Defaults to None.	None
initial_downsample_stack_conv_downscale_list	list	The downscale list of stack conv of initial downsample. Defaults to [1].	[1]
initial_downsample_stack_conv_num_conv_list	list	The num of stack conv list of initial downsample. Defaults to [2].	[2]
ffn_activation	str	The activation of FFN. Defaults to "leaky".	'leaky'
gated_ffn	bool	Whether to use gate FFN. Defaults to False.	False
norm_layer	str	The type of normilize. Defaults to "layer_norm".	'layer_norm'
padding_type	str	The type of padding. Defaults to "ignore".	'ignore'
pos_embed_type	str	The type of pos embedding. Defaults to "t+hw".	't+hw'
checkpoint_level	bool	Whether to use checkpoint. Defaults to True.	True
use_relative_pos	bool	Whether to use relative pose. Defaults to True.	True
self_attn_use_final_proj	bool	Whether to use final projection. Defaults to True.	True
dec_use_first_self_attn	bool	Whether to use first self attention for decoder. Defaults to False.	False
attn_linear_init_mode	str	The mode of attention linear init. Defaults to "0".	'0'
ffn_linear_init_mode	str	The mode of FFN linear init. Defaults to "0".	'0'
conv_init_mode	str	The mode of conv init. Defaults to "0".	'0'
down_up_linear_init_mode	str	The mode of downsample and upsample linear init. Defaults to "0".	'0'
norm_init_mode	str	The mode of normalization init. Defaults to "0".	'0'

Source code in ppsci/arch/cuboid_transformer.py

class CuboidTransformer(base.Arch):
    """Cuboid Transformer for spatiotemporal forecasting

    We adopt the Non-autoregressive encoder-decoder architecture.
    The decoder takes the multi-scale memory output from the encoder.

    The initial downsampling / upsampling layers will be
    Downsampling: [K x Conv2D --> PatchMerge]
    Upsampling: [Nearest Interpolation-based Upsample --> K x Conv2D]

    x --> downsample (optional) ---> (+pos_embed) ---> enc --> mem_l         initial_z (+pos_embed) ---> FC
                                                     |            |
                                                     |------------|
                                                           |
                                                           |
             y <--- upsample (optional) <--- dec <----------

    Args:
        input_keys (Tuple[str, ...]): Name of input keys, such as ("input",).
        output_keys (Tuple[str, ...]): Name of output keys, such as ("output",).
        input_shape (Tuple[int, ...]): The shape of the input data.
        target_shape (Tuple[int, ...]): The shape of the target data.
        base_units (int, optional): The base units. Defaults to 128.
        block_units (int, optional): The block units. Defaults to None.
        scale_alpha (float, optional): We scale up the channels based on the formula:
            - round_to(base_units * max(downsample_scale) ** units_alpha, 4). Defaults to 1.0.
        num_heads (int, optional): The number of heads. Defaults to 4.
        attn_drop (float, optional): The attention dropout. Defaults to 0.0.
        proj_drop (float, optional): The projection dropout. Defaults to 0.0.
        ffn_drop (float, optional): The ffn dropout. Defaults to 0.0.
        downsample (int, optional): The rate of downsample. Defaults to 2.
        downsample_type (str, optional): The type of downsample. Defaults to "patch_merge".
        upsample_type (str, optional): The rate of upsample. Defaults to "upsample".
        upsample_kernel_size (int, optional): The kernel size of upsample. Defaults to 3.
        enc_depth (list, optional): The depth of encoder. Defaults to [4, 4, 4].
        enc_attn_patterns (str, optional): The pattern of encoder attention. Defaults to None.
        enc_cuboid_size (list, optional): The cuboid size of encoder. Defaults to [(4, 4, 4), (4, 4, 4)].
        enc_cuboid_strategy (list, optional): The cuboid strategy of encoder. Defaults to [("l", "l", "l"), ("d", "d", "d")].
        enc_shift_size (list, optional): The shift size of encoder. Defaults to [(0, 0, 0), (0, 0, 0)].
        enc_use_inter_ffn (bool, optional): Whether to use intermediate FFN for encoder. Defaults to True.
        dec_depth (list, optional): The depth of decoder. Defaults to [2, 2].
        dec_cross_start (int, optional): The cross start of decoder. Defaults to 0.
        dec_self_attn_patterns (str, optional): The partterns of decoder. Defaults to None.
        dec_self_cuboid_size (list, optional): The cuboid size of decoder. Defaults to [(4, 4, 4), (4, 4, 4)].
        dec_self_cuboid_strategy (list, optional): The strategy of decoder. Defaults to [("l", "l", "l"), ("d", "d", "d")].
        dec_self_shift_size (list, optional): The shift size of decoder. Defaults to [(1, 1, 1), (0, 0, 0)].
        dec_cross_attn_patterns (_type_, optional): The cross attention patterns of decoder. Defaults to None.
        dec_cross_cuboid_hw (list, optional): The cuboid_hw of decoder. Defaults to [(4, 4), (4, 4)].
        dec_cross_cuboid_strategy (list, optional): The cuboid strategy of decoder. Defaults to [("l", "l", "l"), ("d", "l", "l")].
        dec_cross_shift_hw (list, optional): The shift_hw of decoder. Defaults to [(0, 0), (0, 0)].
        dec_cross_n_temporal (list, optional): The cross_n_temporal of decoder. Defaults to [1, 2].
        dec_cross_last_n_frames (int, optional): The cross_last_n_frames of decoder. Defaults to None.
        dec_use_inter_ffn (bool, optional): Whether to use intermediate FFN for decoder. Defaults to True.
        dec_hierarchical_pos_embed (bool, optional): Whether to use hierarchical pos_embed for decoder. Defaults to False.
        num_global_vectors (int, optional): The num of global vectors. Defaults to 4.
        use_dec_self_global (bool, optional): Whether to use global vector for decoder. Defaults to True.
        dec_self_update_global (bool, optional): Whether to update global vector for decoder. Defaults to True.
        use_dec_cross_global (bool, optional): Whether to use cross global vector for decoder. Defaults to True.
        use_global_vector_ffn (bool, optional): Whether to use global vector FFN. Defaults to True.
        use_global_self_attn (bool, optional): Whether to use global attentions. Defaults to False.
        separate_global_qkv (bool, optional): Whether to separate global qkv. Defaults to False.
        global_dim_ratio (int, optional): The ratio of global dim. Defaults to 1.
        self_pattern (str, optional): The pattern. Defaults to "axial".
        cross_self_pattern (str, optional): The self cross pattern. Defaults to "axial".
        cross_pattern (str, optional): The cross pattern. Defaults to "cross_1x1".
        z_init_method (str, optional): How the initial input to the decoder is initialized. Defaults to "nearest_interp".
        initial_downsample_type (str, optional): The downsample type of initial. Defaults to "conv".
        initial_downsample_activation (str, optional): The downsample activation of initial. Defaults to "leaky".
        initial_downsample_scale (int, optional): The downsample scale of initial. Defaults to 1.
        initial_downsample_conv_layers (int, optional): The conv layer of downsample of initial. Defaults to 2.
        final_upsample_conv_layers (int, optional): The conv layer of final upsample. Defaults to 2.
        initial_downsample_stack_conv_num_layers (int, optional): The num of stack conv layer of initial downsample. Defaults to 1.
        initial_downsample_stack_conv_dim_list (list, optional): The dim list of stack conv of initial downsample. Defaults to None.
        initial_downsample_stack_conv_downscale_list (list, optional): The downscale list of stack conv of initial downsample. Defaults to [1].
        initial_downsample_stack_conv_num_conv_list (list, optional): The num of stack conv list of initial downsample. Defaults to [2].
        ffn_activation (str, optional): The activation of FFN. Defaults to "leaky".
        gated_ffn (bool, optional): Whether to use gate FFN. Defaults to False.
        norm_layer (str, optional): The type of normilize. Defaults to "layer_norm".
        padding_type (str, optional): The type of padding. Defaults to "ignore".
        pos_embed_type (str, optional): The type of pos embedding. Defaults to "t+hw".
        checkpoint_level (bool, optional): Whether to use checkpoint. Defaults to True.
        use_relative_pos (bool, optional): Whether to use relative pose. Defaults to True.
        self_attn_use_final_proj (bool, optional): Whether to use final projection. Defaults to True.
        dec_use_first_self_attn (bool, optional): Whether to use first self attention for decoder. Defaults to False.
        attn_linear_init_mode (str, optional): The mode of attention linear init. Defaults to "0".
        ffn_linear_init_mode (str, optional): The mode of FFN linear init. Defaults to "0".
        conv_init_mode (str, optional): The mode of conv init. Defaults to "0".
        down_up_linear_init_mode (str, optional): The mode of downsample and upsample linear init. Defaults to "0".
        norm_init_mode (str, optional): The mode of normalization init. Defaults to "0".
    """

    def __init__(
        self,
        input_keys: Tuple[str, ...],
        output_keys: Tuple[str, ...],
        input_shape: Tuple[int, ...],
        target_shape: Tuple[int, ...],
        base_units: int = 128,
        block_units: int = None,
        scale_alpha: float = 1.0,
        num_heads: int = 4,
        attn_drop: float = 0.0,
        proj_drop: float = 0.0,
        ffn_drop: float = 0.0,
        downsample: int = 2,
        downsample_type: str = "patch_merge",
        upsample_type: str = "upsample",
        upsample_kernel_size: int = 3,
        enc_depth: Tuple[int, ...] = [4, 4, 4],
        enc_attn_patterns: str = None,
        enc_cuboid_size: Tuple[Tuple[int, ...], ...] = [(4, 4, 4), (4, 4, 4)],
        enc_cuboid_strategy: Tuple[Tuple[str, ...], ...] = [
            ("l", "l", "l"),
            ("d", "d", "d"),
        ],
        enc_shift_size: Tuple[Tuple[int, ...], ...] = [(0, 0, 0), (0, 0, 0)],
        enc_use_inter_ffn: bool = True,
        dec_depth: Tuple[int, ...] = [2, 2],
        dec_cross_start: int = 0,
        dec_self_attn_patterns: str = None,
        dec_self_cuboid_size: Tuple[Tuple[int, ...], ...] = [(4, 4, 4), (4, 4, 4)],
        dec_self_cuboid_strategy: Tuple[Tuple[str, ...], ...] = [
            ("l", "l", "l"),
            ("d", "d", "d"),
        ],
        dec_self_shift_size: Tuple[Tuple[int, ...], ...] = [(1, 1, 1), (0, 0, 0)],
        dec_cross_attn_patterns: str = None,
        dec_cross_cuboid_hw: Tuple[Tuple[int, ...], ...] = [(4, 4), (4, 4)],
        dec_cross_cuboid_strategy: Tuple[Tuple[str, ...], ...] = [
            ("l", "l", "l"),
            ("d", "l", "l"),
        ],
        dec_cross_shift_hw: Tuple[Tuple[int, ...], ...] = [(0, 0), (0, 0)],
        dec_cross_n_temporal: Tuple[int, ...] = [1, 2],
        dec_cross_last_n_frames: int = None,
        dec_use_inter_ffn: bool = True,
        dec_hierarchical_pos_embed: bool = False,
        num_global_vectors: int = 4,
        use_dec_self_global: bool = True,
        dec_self_update_global: bool = True,
        use_dec_cross_global: bool = True,
        use_global_vector_ffn: bool = True,
        use_global_self_attn: bool = False,
        separate_global_qkv: bool = False,
        global_dim_ratio: int = 1,
        self_pattern: str = "axial",
        cross_self_pattern: str = "axial",
        cross_pattern: str = "cross_1x1",
        z_init_method: str = "nearest_interp",
        initial_downsample_type: str = "conv",
        initial_downsample_activation: str = "leaky",
        initial_downsample_scale: int = 1,
        initial_downsample_conv_layers: int = 2,
        final_upsample_conv_layers: int = 2,
        initial_downsample_stack_conv_num_layers: int = 1,
        initial_downsample_stack_conv_dim_list: Tuple[int, ...] = None,
        initial_downsample_stack_conv_downscale_list: Tuple[int, ...] = [1],
        initial_downsample_stack_conv_num_conv_list: Tuple[int, ...] = [2],
        ffn_activation: str = "leaky",
        gated_ffn: bool = False,
        norm_layer: str = "layer_norm",
        padding_type: str = "ignore",
        pos_embed_type: str = "t+hw",
        checkpoint_level: bool = True,
        use_relative_pos: bool = True,
        self_attn_use_final_proj: bool = True,
        dec_use_first_self_attn: bool = False,
        attn_linear_init_mode: str = "0",
        ffn_linear_init_mode: str = "0",
        conv_init_mode: str = "0",
        down_up_linear_init_mode: str = "0",
        norm_init_mode: str = "0",
    ):
        super().__init__()
        self.input_keys = input_keys
        self.output_keys = output_keys
        self.attn_linear_init_mode = attn_linear_init_mode
        self.ffn_linear_init_mode = ffn_linear_init_mode
        self.conv_init_mode = conv_init_mode
        self.down_up_linear_init_mode = down_up_linear_init_mode
        self.norm_init_mode = norm_init_mode
        assert len(enc_depth) == len(dec_depth)
        self.base_units = base_units
        self.num_global_vectors = num_global_vectors

        num_blocks = len(enc_depth)
        if isinstance(self_pattern, str):
            enc_attn_patterns = [self_pattern] * num_blocks

        if isinstance(cross_self_pattern, str):
            dec_self_attn_patterns = [cross_self_pattern] * num_blocks

        if isinstance(cross_pattern, str):
            dec_cross_attn_patterns = [cross_pattern] * num_blocks

        if global_dim_ratio != 1:
            assert (
                separate_global_qkv is True
            ), "Setting global_dim_ratio != 1 requires separate_global_qkv == True."
        self.global_dim_ratio = global_dim_ratio
        self.z_init_method = z_init_method
        assert self.z_init_method in ["zeros", "nearest_interp", "last", "mean"]
        self.input_shape = input_shape
        self.target_shape = target_shape
        T_in, H_in, W_in, C_in = input_shape
        T_out, H_out, W_out, C_out = target_shape
        assert H_in == H_out and W_in == W_out
        if self.num_global_vectors > 0:
            init_data = paddle.zeros(
                (self.num_global_vectors, global_dim_ratio * base_units)
            )
            self.init_global_vectors = paddle.create_parameter(
                shape=init_data.shape,
                dtype=init_data.dtype,
                default_initializer=nn.initializer.Constant(0.0),
            )

            self.init_global_vectors.stop_gradient = not True
        new_input_shape = self.get_initial_encoder_final_decoder(
            initial_downsample_scale=initial_downsample_scale,
            initial_downsample_type=initial_downsample_type,
            activation=initial_downsample_activation,
            initial_downsample_conv_layers=initial_downsample_conv_layers,
            final_upsample_conv_layers=final_upsample_conv_layers,
            padding_type=padding_type,
            initial_downsample_stack_conv_num_layers=initial_downsample_stack_conv_num_layers,
            initial_downsample_stack_conv_dim_list=initial_downsample_stack_conv_dim_list,
            initial_downsample_stack_conv_downscale_list=initial_downsample_stack_conv_downscale_list,
            initial_downsample_stack_conv_num_conv_list=initial_downsample_stack_conv_num_conv_list,
        )
        T_in, H_in, W_in, _ = new_input_shape
        self.encoder = cuboid_encoder.CuboidTransformerEncoder(
            input_shape=(T_in, H_in, W_in, base_units),
            base_units=base_units,
            block_units=block_units,
            scale_alpha=scale_alpha,
            depth=enc_depth,
            downsample=downsample,
            downsample_type=downsample_type,
            block_attn_patterns=enc_attn_patterns,
            block_cuboid_size=enc_cuboid_size,
            block_strategy=enc_cuboid_strategy,
            block_shift_size=enc_shift_size,
            num_heads=num_heads,
            attn_drop=attn_drop,
            proj_drop=proj_drop,
            ffn_drop=ffn_drop,
            gated_ffn=gated_ffn,
            ffn_activation=ffn_activation,
            norm_layer=norm_layer,
            use_inter_ffn=enc_use_inter_ffn,
            padding_type=padding_type,
            use_global_vector=num_global_vectors > 0,
            use_global_vector_ffn=use_global_vector_ffn,
            use_global_self_attn=use_global_self_attn,
            separate_global_qkv=separate_global_qkv,
            global_dim_ratio=global_dim_ratio,
            checkpoint_level=checkpoint_level,
            use_relative_pos=use_relative_pos,
            self_attn_use_final_proj=self_attn_use_final_proj,
            attn_linear_init_mode=attn_linear_init_mode,
            ffn_linear_init_mode=ffn_linear_init_mode,
            conv_init_mode=conv_init_mode,
            down_linear_init_mode=down_up_linear_init_mode,
            norm_init_mode=norm_init_mode,
        )
        self.enc_pos_embed = cuboid_decoder.PosEmbed(
            embed_dim=base_units, typ=pos_embed_type, maxH=H_in, maxW=W_in, maxT=T_in
        )
        mem_shapes = self.encoder.get_mem_shapes()
        self.z_proj = nn.Linear(
            in_features=mem_shapes[-1][-1], out_features=mem_shapes[-1][-1]
        )
        self.dec_pos_embed = cuboid_decoder.PosEmbed(
            embed_dim=mem_shapes[-1][-1],
            typ=pos_embed_type,
            maxT=T_out,
            maxH=mem_shapes[-1][1],
            maxW=mem_shapes[-1][2],
        )
        self.decoder = cuboid_decoder.CuboidTransformerDecoder(
            target_temporal_length=T_out,
            mem_shapes=mem_shapes,
            cross_start=dec_cross_start,
            depth=dec_depth,
            upsample_type=upsample_type,
            block_self_attn_patterns=dec_self_attn_patterns,
            block_self_cuboid_size=dec_self_cuboid_size,
            block_self_shift_size=dec_self_shift_size,
            block_self_cuboid_strategy=dec_self_cuboid_strategy,
            block_cross_attn_patterns=dec_cross_attn_patterns,
            block_cross_cuboid_hw=dec_cross_cuboid_hw,
            block_cross_shift_hw=dec_cross_shift_hw,
            block_cross_cuboid_strategy=dec_cross_cuboid_strategy,
            block_cross_n_temporal=dec_cross_n_temporal,
            cross_last_n_frames=dec_cross_last_n_frames,
            num_heads=num_heads,
            attn_drop=attn_drop,
            proj_drop=proj_drop,
            ffn_drop=ffn_drop,
            upsample_kernel_size=upsample_kernel_size,
            ffn_activation=ffn_activation,
            gated_ffn=gated_ffn,
            norm_layer=norm_layer,
            use_inter_ffn=dec_use_inter_ffn,
            max_temporal_relative=T_in + T_out,
            padding_type=padding_type,
            hierarchical_pos_embed=dec_hierarchical_pos_embed,
            pos_embed_type=pos_embed_type,
            use_self_global=num_global_vectors > 0 and use_dec_self_global,
            self_update_global=dec_self_update_global,
            use_cross_global=num_global_vectors > 0 and use_dec_cross_global,
            use_global_vector_ffn=use_global_vector_ffn,
            use_global_self_attn=use_global_self_attn,
            separate_global_qkv=separate_global_qkv,
            global_dim_ratio=global_dim_ratio,
            checkpoint_level=checkpoint_level,
            use_relative_pos=use_relative_pos,
            self_attn_use_final_proj=self_attn_use_final_proj,
            use_first_self_attn=dec_use_first_self_attn,
            attn_linear_init_mode=attn_linear_init_mode,
            ffn_linear_init_mode=ffn_linear_init_mode,
            conv_init_mode=conv_init_mode,
            up_linear_init_mode=down_up_linear_init_mode,
            norm_init_mode=norm_init_mode,
        )
        self.reset_parameters()

    def get_initial_encoder_final_decoder(
        self,
        initial_downsample_type,
        activation,
        initial_downsample_scale,
        initial_downsample_conv_layers,
        final_upsample_conv_layers,
        padding_type,
        initial_downsample_stack_conv_num_layers,
        initial_downsample_stack_conv_dim_list,
        initial_downsample_stack_conv_downscale_list,
        initial_downsample_stack_conv_num_conv_list,
    ):
        T_in, H_in, W_in, C_in = self.input_shape
        T_out, H_out, W_out, C_out = self.target_shape
        self.initial_downsample_type = initial_downsample_type
        if self.initial_downsample_type == "conv":
            if isinstance(initial_downsample_scale, int):
                initial_downsample_scale = (
                    1,
                    initial_downsample_scale,
                    initial_downsample_scale,
                )
            elif len(initial_downsample_scale) == 2:
                initial_downsample_scale = 1, *initial_downsample_scale
            elif len(initial_downsample_scale) == 3:
                initial_downsample_scale = tuple(initial_downsample_scale)
            else:
                raise NotImplementedError(
                    f"initial_downsample_scale {initial_downsample_scale} format not supported!"
                )
            self.initial_encoder = InitialEncoder(
                dim=C_in,
                out_dim=self.base_units,
                downsample_scale=initial_downsample_scale,
                num_conv_layers=initial_downsample_conv_layers,
                padding_type=padding_type,
                activation=activation,
                conv_init_mode=self.conv_init_mode,
                linear_init_mode=self.down_up_linear_init_mode,
                norm_init_mode=self.norm_init_mode,
            )

            self.final_decoder = FinalDecoder(
                dim=self.base_units,
                target_thw=(T_out, H_out, W_out),
                num_conv_layers=final_upsample_conv_layers,
                activation=activation,
                conv_init_mode=self.conv_init_mode,
                linear_init_mode=self.down_up_linear_init_mode,
                norm_init_mode=self.norm_init_mode,
            )
            new_input_shape = self.initial_encoder.patch_merge.get_out_shape(
                self.input_shape
            )
            self.dec_final_proj = nn.Linear(
                in_features=self.base_units, out_features=C_out
            )
        elif self.initial_downsample_type == "stack_conv":
            if initial_downsample_stack_conv_dim_list is None:
                initial_downsample_stack_conv_dim_list = [
                    self.base_units
                ] * initial_downsample_stack_conv_num_layers
            self.initial_encoder = InitialStackPatchMergingEncoder(
                num_merge=initial_downsample_stack_conv_num_layers,
                in_dim=C_in,
                out_dim_list=initial_downsample_stack_conv_dim_list,
                downsample_scale_list=initial_downsample_stack_conv_downscale_list,
                num_conv_per_merge_list=initial_downsample_stack_conv_num_conv_list,
                padding_type=padding_type,
                activation=activation,
                conv_init_mode=self.conv_init_mode,
                linear_init_mode=self.down_up_linear_init_mode,
                norm_init_mode=self.norm_init_mode,
            )
            initial_encoder_out_shape_list = self.initial_encoder.get_out_shape_list(
                self.target_shape
            )
            (
                dec_target_shape_list,
                dec_in_dim,
            ) = FinalStackUpsamplingDecoder.get_init_params(
                enc_input_shape=self.target_shape,
                enc_out_shape_list=initial_encoder_out_shape_list,
                large_channel=True,
            )
            self.final_decoder = FinalStackUpsamplingDecoder(
                target_shape_list=dec_target_shape_list,
                in_dim=dec_in_dim,
                num_conv_per_up_list=initial_downsample_stack_conv_num_conv_list[::-1],
                activation=activation,
                conv_init_mode=self.conv_init_mode,
                linear_init_mode=self.down_up_linear_init_mode,
                norm_init_mode=self.norm_init_mode,
            )
            self.dec_final_proj = nn.Linear(
                in_features=dec_target_shape_list[-1][-1], out_features=C_out
            )
            new_input_shape = self.initial_encoder.get_out_shape_list(self.input_shape)[
                -1
            ]
        else:
            raise NotImplementedError(f"{self.initial_downsample_type} is invalid.")
        self.input_shape_after_initial_downsample = new_input_shape
        T_in, H_in, W_in, _ = new_input_shape
        return new_input_shape

    def reset_parameters(self):
        if self.num_global_vectors > 0:
            self.init_global_vectors = initializer.trunc_normal_(
                self.init_global_vectors, std=0.02
            )
        if hasattr(self.initial_encoder, "reset_parameters"):
            self.initial_encoder.reset_parameters()
        else:
            cuboid_utils.apply_initialization(
                self.initial_encoder,
                conv_mode=self.conv_init_mode,
                linear_mode=self.down_up_linear_init_mode,
                norm_mode=self.norm_init_mode,
            )
        if hasattr(self.final_decoder, "reset_parameters"):
            self.final_decoder.reset_parameters()
        else:
            cuboid_utils.apply_initialization(
                self.final_decoder,
                conv_mode=self.conv_init_mode,
                linear_mode=self.down_up_linear_init_mode,
                norm_mode=self.norm_init_mode,
            )
        cuboid_utils.apply_initialization(
            self.dec_final_proj, linear_mode=self.down_up_linear_init_mode
        )
        self.encoder.reset_parameters()
        self.enc_pos_embed.reset_parameters()
        self.decoder.reset_parameters()
        self.dec_pos_embed.reset_parameters()
        cuboid_utils.apply_initialization(self.z_proj, linear_mode="0")

    def get_initial_z(self, final_mem, T_out):
        B = final_mem.shape[0]
        if self.z_init_method == "zeros":
            z_shape = list((1, T_out)) + final_mem.shape[2:]
            initial_z = paddle.zeros(shape=z_shape, dtype=final_mem.dtype)
            initial_z = self.z_proj(self.dec_pos_embed(initial_z)).expand(
                shape=[B, -1, -1, -1, -1]
            )
        elif self.z_init_method == "nearest_interp":
            initial_z = nn.functional.interpolate(
                x=final_mem.transpose(perm=[0, 4, 1, 2, 3]),
                size=(T_out, final_mem.shape[2], final_mem.shape[3]),
            ).transpose(perm=[0, 2, 3, 4, 1])
            initial_z = self.z_proj(initial_z)
        elif self.z_init_method == "last":
            initial_z = paddle.broadcast_to(
                x=final_mem[:, -1:, :, :, :], shape=(B, T_out) + final_mem.shape[2:]
            )
            initial_z = self.z_proj(initial_z)
        elif self.z_init_method == "mean":
            initial_z = paddle.broadcast_to(
                x=final_mem.mean(axis=1, keepdims=True),
                shape=(B, T_out) + final_mem.shape[2:],
            )
            initial_z = self.z_proj(initial_z)
        else:
            raise NotImplementedError
        return initial_z

    def forward(self, x: "paddle.Tensor", verbose: bool = False) -> "paddle.Tensor":
        """
        Args:
            x (paddle.Tensor): Tensor with shape (B, T, H, W, C).
            verbose (bool): If True, print intermediate shapes.

        Returns:
            out (paddle.Tensor): The output Shape (B, T_out, H, W, C_out)
        """

        x = self.concat_to_tensor(x, self.input_keys)
        flag_ndim = x.ndim
        if flag_ndim == 6:
            x = x.reshape([-1, *x.shape[2:]])
        B, _, _, _, _ = x.shape

        T_out = self.target_shape[0]
        x = self.initial_encoder(x)
        x = self.enc_pos_embed(x)

        if self.num_global_vectors > 0:
            init_global_vectors = self.init_global_vectors.expand(
                shape=[
                    B,
                    self.num_global_vectors,
                    self.global_dim_ratio * self.base_units,
                ]
            )
            mem_l, mem_global_vector_l = self.encoder(x, init_global_vectors)
        else:
            mem_l = self.encoder(x)

        if verbose:
            for i, mem in enumerate(mem_l):
                print(f"mem[{i}].shape = {mem.shape}")
        initial_z = self.get_initial_z(final_mem=mem_l[-1], T_out=T_out)

        if self.num_global_vectors > 0:
            dec_out = self.decoder(initial_z, mem_l, mem_global_vector_l)
        else:
            dec_out = self.decoder(initial_z, mem_l)

        dec_out = self.final_decoder(dec_out)

        out = self.dec_final_proj(dec_out)
        if flag_ndim == 6:
            out = out.reshape([-1, *out.shape])
        return {key: out for key in self.output_keys}

forward(x, verbose=False)

Parameters:

Name	Type	Description	Default
x	Tensor	Tensor with shape (B, T, H, W, C).	required
verbose	bool	If True, print intermediate shapes.	False

Returns:

Name	Type	Description
out	Tensor	The output Shape (B, T_out, H, W, C_out)

Source code in ppsci/arch/cuboid_transformer.py

def forward(self, x: "paddle.Tensor", verbose: bool = False) -> "paddle.Tensor":
    """
    Args:
        x (paddle.Tensor): Tensor with shape (B, T, H, W, C).
        verbose (bool): If True, print intermediate shapes.

    Returns:
        out (paddle.Tensor): The output Shape (B, T_out, H, W, C_out)
    """

    x = self.concat_to_tensor(x, self.input_keys)
    flag_ndim = x.ndim
    if flag_ndim == 6:
        x = x.reshape([-1, *x.shape[2:]])
    B, _, _, _, _ = x.shape

    T_out = self.target_shape[0]
    x = self.initial_encoder(x)
    x = self.enc_pos_embed(x)

    if self.num_global_vectors > 0:
        init_global_vectors = self.init_global_vectors.expand(
            shape=[
                B,
                self.num_global_vectors,
                self.global_dim_ratio * self.base_units,
            ]
        )
        mem_l, mem_global_vector_l = self.encoder(x, init_global_vectors)
    else:
        mem_l = self.encoder(x)

    if verbose:
        for i, mem in enumerate(mem_l):
            print(f"mem[{i}].shape = {mem.shape}")
    initial_z = self.get_initial_z(final_mem=mem_l[-1], T_out=T_out)

    if self.num_global_vectors > 0:
        dec_out = self.decoder(initial_z, mem_l, mem_global_vector_l)
    else:
        dec_out = self.decoder(initial_z, mem_l)

    dec_out = self.final_decoder(dec_out)

    out = self.dec_final_proj(dec_out)
    if flag_ndim == 6:
        out = out.reshape([-1, *out.shape])
    return {key: out for key in self.output_keys}

CVit1D

Bases: Arch

1D Convolutional Vision Transformer (CVit1D) class.

Bridging Operator Learning and Conditioned Neural Fields: A Unifying Perspective

Parameters:

Name	Type	Description	Default
input_keys	Sequence[str]	Keys identifying the input tensors.	required
output_keys	Sequence[str]	Keys identifying the output tensors.	required
spatial_dims	int	The spatial dimensions of the input data.	required
in_dim	int	The dimensionality of the input data.	required
coords_dim	int	The dimensionality of the positional encoding.	required
patch_size	Sequence[int]	Size of the patches. Defaults to (4,).	(4,)
grid_size	Sequence[int]	Size of the grid. Defaults to (200,).	(200,)
latent_dim	int	Dimensionality of the latent space. Defaults to 256.	256
emb_dim	int	Dimensionality of the embedding space. Defaults to 256.	256
depth	int	Number of transformer encoder layers. Defaults to 3.	3
num_heads	int	Number of attention heads. Defaults to 8.	8
dec_emb_dim	int	Dimensionality of the decoder embedding space. Defaults to 256.	256
dec_num_heads	int	Number of decoder attention heads. Defaults to 8.	8
dec_depth	int	Number of decoder transformer layers. Defaults to 1.	1
num_mlp_layers	int	Number of layers in the MLP. Defaults to 1.	1
mlp_ratio	int	Ratio for determining the size of the MLP's hidden layer. Defaults to 1.	1
out_dim	int	Dimensionality of the output data. Defaults to 1.	1
layer_norm_eps	float	Epsilon for layer normalization. Defaults to 1e-5.	1e-05
embedding_type	str	Type of embedding to use ("grid" or other options). Defaults to "grid".	'grid'

Examples:

>>> import ppsci
>>> b, l, c = 2, 32, 1
>>> l_query = 42
>>> c_in = 1
>>> c_out = 1
>>> model = ppsci.arch.CVit1D(
...     input_keys=["u", "y"],
...     output_keys=["s"],
...     in_dim=c_in,
...     coords_dim=1,
...     spatial_dims=l,
...     patch_size=[4],
...     grid_size=[l],
...     latent_dim=32,
...     emb_dim=32,
...     depth=3,
...     num_heads=8,
...     dec_emb_dim=32,
...     dec_num_heads=8,
...     dec_depth=1,
...     num_mlp_layers=1,
...     mlp_ratio=1,
...     out_dim=c_out,
...     layer_norm_eps=1e-5,
...     embedding_type="grid",
... )
>>> x = paddle.randn([b, l, c_in])
>>> coords = paddle.randn([l_query, 1])
>>> out = model({"u": x, "y": coords})["s"]
>>> print(out.shape) # output shape should be [b, l_query, c_out]
[2, 42, 1]

Source code in ppsci/arch/cvit.py

class CVit1D(base.Arch):
    """
    1D Convolutional Vision Transformer (CVit1D) class.

    [Bridging Operator Learning and Conditioned Neural Fields: A Unifying Perspective](https://arxiv.org/abs/2405.13998)

    Args:
        input_keys (Sequence[str]): Keys identifying the input tensors.
        output_keys (Sequence[str]): Keys identifying the output tensors.
        spatial_dims (int): The spatial dimensions of the input data.
        in_dim (int): The dimensionality of the input data.
        coords_dim (int): The dimensionality of the positional encoding.
        patch_size (Sequence[int], optional): Size of the patches. Defaults to (4,).
        grid_size (Sequence[int], optional): Size of the grid. Defaults to (200,).
        latent_dim (int, optional): Dimensionality of the latent space. Defaults to 256.
        emb_dim (int, optional): Dimensionality of the embedding space. Defaults to 256.
        depth (int, optional): Number of transformer encoder layers. Defaults to 3.
        num_heads (int, optional): Number of attention heads. Defaults to 8.
        dec_emb_dim (int, optional): Dimensionality of the decoder embedding space. Defaults to 256.
        dec_num_heads (int, optional): Number of decoder attention heads. Defaults to 8.
        dec_depth (int, optional): Number of decoder transformer layers. Defaults to 1.
        num_mlp_layers (int, optional): Number of layers in the MLP. Defaults to 1.
        mlp_ratio (int, optional): Ratio for determining the size of the MLP's hidden layer. Defaults to 1.
        out_dim (int, optional): Dimensionality of the output data. Defaults to 1.
        layer_norm_eps (float, optional): Epsilon for layer normalization. Defaults to 1e-5.
        embedding_type (str, optional): Type of embedding to use ("grid" or other options). Defaults to "grid".

    Examples:
        >>> import ppsci
        >>> b, l, c = 2, 32, 1
        >>> l_query = 42
        >>> c_in = 1
        >>> c_out = 1
        >>> model = ppsci.arch.CVit1D(
        ...     input_keys=["u", "y"],
        ...     output_keys=["s"],
        ...     in_dim=c_in,
        ...     coords_dim=1,
        ...     spatial_dims=l,
        ...     patch_size=[4],
        ...     grid_size=[l],
        ...     latent_dim=32,
        ...     emb_dim=32,
        ...     depth=3,
        ...     num_heads=8,
        ...     dec_emb_dim=32,
        ...     dec_num_heads=8,
        ...     dec_depth=1,
        ...     num_mlp_layers=1,
        ...     mlp_ratio=1,
        ...     out_dim=c_out,
        ...     layer_norm_eps=1e-5,
        ...     embedding_type="grid",
        ... )
        >>> x = paddle.randn([b, l, c_in])
        >>> coords = paddle.randn([l_query, 1])
        >>> out = model({"u": x, "y": coords})["s"]
        >>> print(out.shape) # output shape should be [b, l_query, c_out]
        [2, 42, 1]
    """

    def __init__(
        self,
        input_keys: Sequence[str],
        output_keys: Sequence[str],
        spatial_dims: int,
        in_dim: int,
        coords_dim: int,
        patch_size: Sequence[int] = (4,),
        grid_size: Sequence[int] = (200,),
        latent_dim: int = 256,
        emb_dim: int = 256,
        depth: int = 3,
        num_heads: int = 8,
        dec_emb_dim: int = 256,
        dec_num_heads: int = 8,
        dec_depth: int = 1,
        num_mlp_layers: int = 1,
        mlp_ratio: int = 1,
        out_dim: int = 1,
        layer_norm_eps: float = 1e-5,
        embedding_type: str = "grid",
    ):
        if not importlib.util.find_spec("einops"):
            raise ModuleNotFoundError(
                "Please install `einops` by running 'pip install einops'."
            )
        super().__init__()
        self.input_keys = input_keys
        self.output_keys = output_keys
        self.spatial_dims = spatial_dims
        self.in_dim = in_dim
        self.coords_dim = coords_dim
        self.patch_size = patch_size
        self.grid_size = grid_size
        self.latent_dim = latent_dim
        self.emb_dim = emb_dim
        self.depth = depth
        self.num_heads = num_heads
        self.dec_emb_dim = dec_emb_dim
        self.dec_num_heads = dec_num_heads
        self.dec_depth = dec_depth
        self.num_mlp_layers = num_mlp_layers
        self.mlp_ratio = mlp_ratio
        self.out_dim = out_dim
        self.layer_norm_eps = layer_norm_eps
        self.embedding_type = embedding_type

        if self.embedding_type == "grid":
            # Create grid and latents
            n_x = self.grid_size[0]
            self.grid = paddle.linspace(0, 1, n_x)
            self.latents = self.create_parameter(
                [n_x, self.latent_dim],
                default_initializer=nn.initializer.Normal(std=1e-2),
            )
            self.fc = nn.Linear(self.latent_dim, self.dec_emb_dim)
            self.norm = nn.LayerNorm(self.dec_emb_dim, self.layer_norm_eps)
        elif self.embedding_type == "mlp":
            self.mlp = MlpBlock(self.coords_dim, self.dec_emb_dim, self.dec_emb_dim)
            self.norm = nn.LayerNorm(self.dec_emb_dim, self.layer_norm_eps)

        self.encoder = Encoder1D(
            self.in_dim,
            self.spatial_dims,
            self.patch_size,
            self.emb_dim,
            self.depth,
            self.num_heads,
            self.mlp_ratio,
            self.layer_norm_eps,
        )
        self.enc_norm = nn.LayerNorm(self.emb_dim, self.layer_norm_eps)
        self.fc1 = nn.Linear(self.emb_dim, self.dec_emb_dim)
        self.cross_attn_blocks = nn.LayerList(
            [
                CrossAttnBlock(
                    self.dec_num_heads,
                    self.dec_emb_dim,
                    self.mlp_ratio,
                    self.layer_norm_eps,
                    self.dec_emb_dim,
                    self.dec_emb_dim,
                )
                for _ in range(self.dec_depth)
            ]
        )
        self.block_norm = nn.LayerNorm(self.dec_emb_dim, self.layer_norm_eps)
        self.final_mlp = Mlp(
            self.num_mlp_layers,
            self.dec_emb_dim,
            self.out_dim,
            layer_norm_eps=self.layer_norm_eps,
        )

    def forward_tensor(self, x, coords):
        b, h, c = x.shape

        # process query coordinates
        if self.embedding_type == "grid":
            d2 = (coords - self.grid.unsqueeze(0)) ** 2
            w = paddle.exp(-1e5 * d2) / paddle.exp(-1e5 * d2).sum(axis=1, keepdim=True)
            coords = paddle.einsum("ic,pi->pc", self.latents, w)
            coords = self.fc(coords)
            coords = self.norm(coords)
        elif self.embedding_type == "mlp":
            coords = self.mlp(coords)
            coords = self.norm(coords)

        coords = einops.repeat(coords, "n d -> b n d", b=b)

        # process input function(encoder)
        x = self.encoder(x)
        x = self.enc_norm(x)
        x = self.fc1(x)

        # decoder
        for i, block in enumerate(self.cross_attn_blocks):
            coords = block(coords, x)

        # mlp
        x = self.block_norm(coords)
        x = self.final_mlp(x)

        return x

    def forward(self, x_dict):
        if self._input_transform is not None:
            x = self._input_transform(x_dict)

        x, coords = x_dict[self.input_keys[0]], x_dict[self.input_keys[1]]
        if coords.ndim >= 3:
            coords = coords[0]  # [b, n, c] -> [n, c]

        y = self.forward_tensor(x, coords)

        y_dict = {self.output_keys[0]: y}
        if self._output_transform is not None:
            y_dict = self._output_transform(x_dict, y_dict)

        return y_dict

CylinderEmbedding

Bases: Arch

Embedding Koopman model for the Cylinder system.

Parameters:

Name	Type	Description	Default
input_keys	Tuple[str, ...]	Input keys, such as ("states", "visc").	required
output_keys	Tuple[str, ...]	Output keys, such as ("pred_states", "recover_states").	required
mean	Optional[Tuple[float, ...]]	Mean of training dataset. Defaults to None.	None
std	Optional[Tuple[float, ...]]	Standard Deviation of training dataset. Defaults to None.	None
embed_size	int	Number of embedding size. Defaults to 128.	128
encoder_channels	Optional[Tuple[int, ...]]	Number of channels in encoder network. Defaults to None.	None
decoder_channels	Optional[Tuple[int, ...]]	Number of channels in decoder network. Defaults to None.	None
drop	float	Probability of dropout the units. Defaults to 0.0.	0.0

Examples:

>>> import paddle
>>> import ppsci
>>> model = ppsci.arch.CylinderEmbedding(("states", "visc"), ("pred_states", "recover_states"))
>>> states_shape = [32, 10, 3, 64, 128]
>>> visc_shape = [32, 1]
>>> input_dict = {"states" : paddle.rand(states_shape),
...               "visc" : paddle.rand(visc_shape)}
>>> out_dict = model(input_dict)
>>> print(out_dict["pred_states"].shape)
[32, 9, 3, 64, 128]
>>> print(out_dict["recover_states"].shape)
[32, 10, 3, 64, 128]

Source code in ppsci/arch/embedding_koopman.py

class CylinderEmbedding(base.Arch):
    """Embedding Koopman model for the Cylinder system.

    Args:
        input_keys (Tuple[str, ...]): Input keys, such as ("states", "visc").
        output_keys (Tuple[str, ...]): Output keys, such as ("pred_states", "recover_states").
        mean (Optional[Tuple[float, ...]]): Mean of training dataset. Defaults to None.
        std (Optional[Tuple[float, ...]]): Standard Deviation of training dataset. Defaults to None.
        embed_size (int, optional): Number of embedding size. Defaults to 128.
        encoder_channels (Optional[Tuple[int, ...]]): Number of channels in encoder network. Defaults to None.
        decoder_channels (Optional[Tuple[int, ...]]): Number of channels in decoder network. Defaults to None.
        drop (float, optional):  Probability of dropout the units. Defaults to 0.0.

    Examples:
        >>> import paddle
        >>> import ppsci
        >>> model = ppsci.arch.CylinderEmbedding(("states", "visc"), ("pred_states", "recover_states"))
        >>> states_shape = [32, 10, 3, 64, 128]
        >>> visc_shape = [32, 1]
        >>> input_dict = {"states" : paddle.rand(states_shape),
        ...               "visc" : paddle.rand(visc_shape)}
        >>> out_dict = model(input_dict)
        >>> print(out_dict["pred_states"].shape)
        [32, 9, 3, 64, 128]
        >>> print(out_dict["recover_states"].shape)
        [32, 10, 3, 64, 128]
    """

    def __init__(
        self,
        input_keys: Tuple[str, ...],
        output_keys: Tuple[str, ...],
        mean: Optional[Tuple[float, ...]] = None,
        std: Optional[Tuple[float, ...]] = None,
        embed_size: int = 128,
        encoder_channels: Optional[Tuple[int, ...]] = None,
        decoder_channels: Optional[Tuple[int, ...]] = None,
        drop: float = 0.0,
    ):
        super().__init__()
        self.input_keys = input_keys
        self.output_keys = output_keys
        self.embed_size = embed_size

        X, Y = np.meshgrid(
            np.linspace(-2, 14, 128, dtype=paddle.get_default_dtype()),
            np.linspace(-4, 4, 64, dtype=paddle.get_default_dtype()),
        )
        self.mask = paddle.to_tensor(np.sqrt(X**2 + Y**2)).unsqueeze(0).unsqueeze(0)

        encoder_channels = (
            [4, 16, 32, 64, 128] if encoder_channels is None else encoder_channels
        )
        decoder_channels = (
            [embed_size // 32, 128, 64, 32, 16]
            if decoder_channels is None
            else decoder_channels
        )
        self.encoder_net = self.build_encoder(embed_size, encoder_channels, drop)
        self.k_diag_net, self.k_ut_net, self.k_lt_net = self.build_koopman_operator(
            embed_size
        )
        self.decoder_net = self.build_decoder(decoder_channels)

        xidx = []
        yidx = []
        for i in range(1, 5):
            yidx.append(np.arange(i, embed_size))
            xidx.append(np.arange(0, embed_size - i))
        self.xidx = paddle.to_tensor(np.concatenate(xidx), dtype="int64")
        self.yidx = paddle.to_tensor(np.concatenate(yidx), dtype="int64")

        mean = [0.0, 0.0, 0.0, 0.0] if mean is None else mean
        std = [1.0, 1.0, 1.0, 1.0] if std is None else std
        self.register_buffer("mean", paddle.to_tensor(mean).reshape([1, 4, 1, 1]))
        self.register_buffer("std", paddle.to_tensor(std).reshape([1, 4, 1, 1]))

        self.apply(self._init_weights)

    def _init_weights(self, m):
        if isinstance(m, nn.Linear):
            k = 1 / m.weight.shape[0]
            uniform = Uniform(-(k**0.5), k**0.5)
            uniform(m.weight)
            if m.bias is not None:
                uniform(m.bias)
        elif isinstance(m, nn.LayerNorm):
            zeros_(m.bias)
            ones_(m.weight)
        elif isinstance(m, nn.Conv2D):
            k = 1 / (m.weight.shape[1] * m.weight.shape[2] * m.weight.shape[3])
            uniform = Uniform(-(k**0.5), k**0.5)
            uniform(m.weight)
            if m.bias is not None:
                uniform(m.bias)

    def _build_conv_relu_list(
        self, in_channels: Tuple[int, ...], out_channels: Tuple[int, ...]
    ):
        net_list = [
            nn.Conv2D(
                in_channels,
                out_channels,
                kernel_size=(3, 3),
                stride=2,
                padding=1,
                padding_mode="replicate",
            ),
            nn.ReLU(),
        ]
        return net_list

    def build_encoder(
        self, embed_size: int, channels: Tuple[int, ...], drop: float = 0.0
    ):
        net = []
        for i in range(1, len(channels)):
            net.extend(self._build_conv_relu_list(channels[i - 1], channels[i]))
        net.append(
            nn.Conv2D(
                channels[-1],
                embed_size // 32,
                kernel_size=(3, 3),
                padding=1,
                padding_mode="replicate",
            )
        )
        net.append(
            nn.LayerNorm(
                (4, 4, 8),
            )
        )
        net.append(nn.Dropout(drop))
        net = nn.Sequential(*net)
        return net

    def _build_upsample_conv_relu(
        self, in_channels: Tuple[int, ...], out_channels: Tuple[int, ...]
    ):
        net_list = [
            nn.Upsample(scale_factor=2, mode="bilinear", align_corners=True),
            nn.Conv2D(
                in_channels,
                out_channels,
                kernel_size=(3, 3),
                stride=1,
                padding=1,
                padding_mode="replicate",
            ),
            nn.ReLU(),
        ]
        return net_list

    def build_decoder(self, channels: Tuple[int, ...]):
        net = []
        for i in range(1, len(channels)):
            net.extend(self._build_upsample_conv_relu(channels[i - 1], channels[i]))
        net.append(
            nn.Conv2D(
                channels[-1],
                3,
                kernel_size=(3, 3),
                stride=1,
                padding=1,
                padding_mode="replicate",
            ),
        )
        net = nn.Sequential(*net)
        return net

    def build_koopman_operator(self, embed_size: int):
        # Learned Koopman operator parameters
        k_diag_net = nn.Sequential(
            nn.Linear(1, 50), nn.ReLU(), nn.Linear(50, embed_size)
        )

        k_ut_net = nn.Sequential(
            nn.Linear(1, 50), nn.ReLU(), nn.Linear(50, 4 * embed_size - 10)
        )
        k_lt_net = nn.Sequential(
            nn.Linear(1, 50), nn.ReLU(), nn.Linear(50, 4 * embed_size - 10)
        )
        return k_diag_net, k_ut_net, k_lt_net

    def encoder(self, x: paddle.Tensor, viscosity: paddle.Tensor):
        B, T, C, H, W = x.shape
        x = x.reshape((B * T, C, H, W))
        viscosity = viscosity.repeat_interleave(T, axis=1).reshape((B * T, 1))
        x = paddle.concat(
            [x, viscosity.unsqueeze(-1).unsqueeze(-1) * paddle.ones_like(x[:, :1])],
            axis=1,
        )
        x = self._normalize(x)
        g = self.encoder_net(x)
        g = g.reshape([B, T, -1])
        return g

    def decoder(self, g: paddle.Tensor):
        B, T, _ = g.shape
        x = self.decoder_net(g.reshape([-1, self.embed_size // 32, 4, 8]))
        x = self._unnormalize(x)
        mask0 = (
            self.mask.repeat_interleave(x.shape[1], axis=1).repeat_interleave(
                x.shape[0], axis=0
            )
            < 1
        )
        x[mask0] = 0
        _, C, H, W = x.shape
        x = x.reshape([B, T, C, H, W])
        return x

    def get_koopman_matrix(self, g: paddle.Tensor, visc: paddle.Tensor):
        # # Koopman operator
        kMatrix = paddle.zeros([g.shape[0], self.embed_size, self.embed_size])
        kMatrix.stop_gradient = False
        # Populate the off diagonal terms
        kMatrixUT_data = self.k_ut_net(100 * visc)
        kMatrixLT_data = self.k_lt_net(100 * visc)

        kMatrix = kMatrix.transpose([1, 2, 0])
        kMatrixUT_data_t = kMatrixUT_data.transpose([1, 0])
        kMatrixLT_data_t = kMatrixLT_data.transpose([1, 0])
        kMatrix[self.xidx, self.yidx] = kMatrixUT_data_t
        kMatrix[self.yidx, self.xidx] = kMatrixLT_data_t

        # Populate the diagonal
        ind = np.diag_indices(kMatrix.shape[1])
        ind = paddle.to_tensor(ind, dtype="int64")

        kMatrixDiag = self.k_diag_net(100 * visc)
        kMatrixDiag_t = kMatrixDiag.transpose([1, 0])
        kMatrix[ind[0], ind[1]] = kMatrixDiag_t
        return kMatrix.transpose([2, 0, 1])

    def koopman_operation(self, embed_data: paddle.Tensor, k_matrix: paddle.Tensor):
        embed_pred_data = paddle.bmm(
            k_matrix, embed_data.transpose([0, 2, 1])
        ).transpose([0, 2, 1])
        return embed_pred_data

    def _normalize(self, x: paddle.Tensor):
        x = (x - self.mean) / self.std
        return x

    def _unnormalize(self, x: paddle.Tensor):
        return self.std[:, :3] * x + self.mean[:, :3]

    def forward_tensor(self, states, visc):
        # states.shape=(B, T, C, H, W)
        embed_data = self.encoder(states, visc)
        recover_data = self.decoder(embed_data)

        k_matrix = self.get_koopman_matrix(embed_data, visc)
        embed_pred_data = self.koopman_operation(embed_data, k_matrix)
        pred_data = self.decoder(embed_pred_data)

        return (pred_data[:, :-1], recover_data, k_matrix)

    @staticmethod
    def split_to_dict(data_tensors: Tuple[paddle.Tensor, ...], keys: Tuple[str, ...]):
        return {key: data_tensors[i] for i, key in enumerate(keys)}

    def forward(self, x):

        if self._input_transform is not None:
            x = self._input_transform(x)

        y = self.forward_tensor(**x)
        y = self.split_to_dict(y, self.output_keys)

        if self._output_transform is not None:
            y = self._output_transform(x, y)
        return y

DeepONet

Bases: Arch

Deep operator network.

Lu et al. Learning nonlinear operators via DeepONet based on the universal approximation theorem of operators. Nat Mach Intell, 2021.

Parameters:

Name	Type	Description	Default
u_key	str	Name of function data for input function u(x).	required
y_key	str	Name of location data for input function G(u).	required
G_key	str	Output name of predicted G(u)(y).	required
num_loc	int	Number of sampled u(x), i.e. m in paper.	required
num_features	int	Number of features extracted from u(x), same for y.	required
branch_num_layers	int	Number of hidden layers of branch net.	required
trunk_num_layers	int	Number of hidden layers of trunk net.	required
branch_hidden_size	Union[int, Tuple[int, ...]]	Number of hidden size of branch net. An integer for all layers, or list of integer specify each layer's size.	required
trunk_hidden_size	Union[int, Tuple[int, ...]]	Number of hidden size of trunk net. An integer for all layers, or list of integer specify each layer's size.	required
branch_skip_connection	bool	Whether to use skip connection for branch net. Defaults to False.	False
trunk_skip_connection	bool	Whether to use skip connection for trunk net. Defaults to False.	False
branch_activation	str	Name of activation function. Defaults to "tanh".	'tanh'
trunk_activation	str	Name of activation function. Defaults to "tanh".	'tanh'
branch_weight_norm	bool	Whether to apply weight norm on parameter(s) for branch net. Defaults to False.	False
trunk_weight_norm	bool	Whether to apply weight norm on parameter(s) for trunk net. Defaults to False.	False
use_bias	bool	Whether to add bias on predicted G(u)(y). Defaults to True.	True

Examples:

>>> import paddle
>>> import ppsci
>>> model = ppsci.arch.DeepONet(
...     "u", "y", "G",
...     100, 40,
...     1, 1,
...     40, 40,
...     branch_activation="relu", trunk_activation="relu",
...     use_bias=True,
... )
>>> input_dict = {"u": paddle.rand([200, 100]),
...               "y": paddle.rand([200, 1])}
>>> output_dict = model(input_dict)
>>> print(output_dict["G"].shape)
[200, 1]

Source code in ppsci/arch/deeponet.py

class DeepONet(base.Arch):
    """Deep operator network.

    [Lu et al. Learning nonlinear operators via DeepONet based on the universal approximation theorem of operators. Nat Mach Intell, 2021.](https://doi.org/10.1038/s42256-021-00302-5)

    Args:
        u_key (str): Name of function data for input function u(x).
        y_key (str): Name of location data for input function G(u).
        G_key (str): Output name of predicted G(u)(y).
        num_loc (int): Number of sampled u(x), i.e. `m` in paper.
        num_features (int): Number of features extracted from u(x), same for y.
        branch_num_layers (int): Number of hidden layers of branch net.
        trunk_num_layers (int): Number of hidden layers of trunk net.
        branch_hidden_size (Union[int, Tuple[int, ...]]): Number of hidden size of branch net.
            An integer for all layers, or list of integer specify each layer's size.
        trunk_hidden_size (Union[int, Tuple[int, ...]]): Number of hidden size of trunk net.
            An integer for all layers, or list of integer specify each layer's size.
        branch_skip_connection (bool, optional): Whether to use skip connection for branch net. Defaults to False.
        trunk_skip_connection (bool, optional): Whether to use skip connection for trunk net. Defaults to False.
        branch_activation (str, optional): Name of activation function. Defaults to "tanh".
        trunk_activation (str, optional): Name of activation function. Defaults to "tanh".
        branch_weight_norm (bool, optional): Whether to apply weight norm on parameter(s) for branch net. Defaults to False.
        trunk_weight_norm (bool, optional): Whether to apply weight norm on parameter(s) for trunk net. Defaults to False.
        use_bias (bool, optional): Whether to add bias on predicted G(u)(y). Defaults to True.

    Examples:
        >>> import paddle
        >>> import ppsci
        >>> model = ppsci.arch.DeepONet(
        ...     "u", "y", "G",
        ...     100, 40,
        ...     1, 1,
        ...     40, 40,
        ...     branch_activation="relu", trunk_activation="relu",
        ...     use_bias=True,
        ... )
        >>> input_dict = {"u": paddle.rand([200, 100]),
        ...               "y": paddle.rand([200, 1])}
        >>> output_dict = model(input_dict)
        >>> print(output_dict["G"].shape)
        [200, 1]
    """

    def __init__(
        self,
        u_key: str,
        y_key: str,
        G_key: str,
        num_loc: int,
        num_features: int,
        branch_num_layers: int,
        trunk_num_layers: int,
        branch_hidden_size: Union[int, Tuple[int, ...]],
        trunk_hidden_size: Union[int, Tuple[int, ...]],
        branch_skip_connection: bool = False,
        trunk_skip_connection: bool = False,
        branch_activation: str = "tanh",
        trunk_activation: str = "tanh",
        branch_weight_norm: bool = False,
        trunk_weight_norm: bool = False,
        use_bias: bool = True,
    ):
        super().__init__()
        self.u_key = u_key
        self.y_key = y_key
        self.input_keys = (u_key, y_key)
        self.output_keys = (G_key,)

        self.branch_net = mlp.MLP(
            (self.u_key,),
            ("b",),
            branch_num_layers,
            branch_hidden_size,
            branch_activation,
            branch_skip_connection,
            branch_weight_norm,
            input_dim=num_loc,
            output_dim=num_features,
        )

        self.trunk_net = mlp.MLP(
            (self.y_key,),
            ("t",),
            trunk_num_layers,
            trunk_hidden_size,
            trunk_activation,
            trunk_skip_connection,
            trunk_weight_norm,
            input_dim=1,
            output_dim=num_features,
        )
        self.trunk_act = act_mod.get_activation(trunk_activation)

        self.use_bias = use_bias
        if use_bias:
            # register bias to parameter for updating in optimizer and storage
            self.b = self.create_parameter(
                shape=(1,),
                attr=nn.initializer.Constant(0.0),
            )

    def forward(self, x):
        if self._input_transform is not None:
            x = self._input_transform(x)

        # Branch net to encode the input function
        u_features = self.branch_net(x)[self.branch_net.output_keys[0]]

        # Trunk net to encode the domain of the output function
        y_features = self.trunk_net(x)
        y_features = self.trunk_act(y_features[self.trunk_net.output_keys[0]])

        # Dot product
        G_u = paddle.einsum("bi,bi->b", u_features, y_features)  # [batch_size, ]
        G_u = paddle.reshape(G_u, [-1, 1])  # reshape [batch_size, ] to [batch_size, 1]

        # Add bias
        if self.use_bias:
            G_u += self.b

        result_dict = {
            self.output_keys[0]: G_u,
        }
        if self._output_transform is not None:
            result_dict = self._output_transform(x, result_dict)

        return result_dict

DeepPhyLSTM

Bases: Arch

DeepPhyLSTM init function.

Parameters:

Name	Type	Description	Default
input_size	int	The input size.	required
output_size	int	The output size.	required
hidden_size	int	The hidden size. Defaults to 100.	100
model_type	int	The model type, value is 2 or 3, 2 indicates having two sub-models, 3 indicates having three submodels. Defaults to 2.	2

Examples:

>>> import paddle
>>> import ppsci
>>> # model_type is `2`
>>> model = ppsci.arch.DeepPhyLSTM(
...     input_size=16,
...     output_size=1,
...     hidden_size=100,
...     model_type=2)
>>> out = model(
...     {"ag":paddle.rand([64, 16, 16]),
...     "ag_c":paddle.rand([64, 16, 16]),
...     "phi":paddle.rand([1, 16, 16])})
>>> for k, v in out.items():
...     print(f"{k} {v.dtype} {v.shape}")
eta_pred paddle.float32 [64, 16, 1]
eta_dot_pred paddle.float32 [64, 16, 1]
g_pred paddle.float32 [64, 16, 1]
eta_t_pred_c paddle.float32 [64, 16, 1]
eta_dot_pred_c paddle.float32 [64, 16, 1]
lift_pred_c paddle.float32 [64, 16, 1]
>>> # model_type is `3`
>>> model = ppsci.arch.DeepPhyLSTM(
...     input_size=16,
...     output_size=1,
...     hidden_size=100,
...     model_type=3)
>>> out = model(
...     {"ag":paddle.rand([64, 16, 1]),
...     "ag_c":paddle.rand([64, 16, 1]),
...     "phi":paddle.rand([1, 16, 16])})
>>> for k, v in out.items():
...     print(f"{k} {v.dtype} {v.shape}")
eta_pred paddle.float32 [64, 16, 1]
eta_dot_pred paddle.float32 [64, 16, 1]
g_pred paddle.float32 [64, 16, 1]
eta_t_pred_c paddle.float32 [64, 16, 1]
eta_dot_pred_c paddle.float32 [64, 16, 1]
lift_pred_c paddle.float32 [64, 16, 1]
g_t_pred_c paddle.float32 [64, 16, 1]
g_dot_pred_c paddle.float32 [64, 16, 1]

Source code in ppsci/arch/phylstm.py

class DeepPhyLSTM(base.Arch):
    """DeepPhyLSTM init function.

    Args:
        input_size (int): The input size.
        output_size (int): The output size.
        hidden_size (int, optional): The hidden size. Defaults to 100.
        model_type (int, optional): The model type, value is 2 or 3, 2 indicates having two sub-models, 3 indicates having three submodels. Defaults to 2.

    Examples:
        >>> import paddle
        >>> import ppsci
        >>> # model_type is `2`
        >>> model = ppsci.arch.DeepPhyLSTM(
        ...     input_size=16,
        ...     output_size=1,
        ...     hidden_size=100,
        ...     model_type=2)
        >>> out = model(
        ...     {"ag":paddle.rand([64, 16, 16]),
        ...     "ag_c":paddle.rand([64, 16, 16]),
        ...     "phi":paddle.rand([1, 16, 16])})
        >>> for k, v in out.items():
        ...     print(f"{k} {v.dtype} {v.shape}")
        eta_pred paddle.float32 [64, 16, 1]
        eta_dot_pred paddle.float32 [64, 16, 1]
        g_pred paddle.float32 [64, 16, 1]
        eta_t_pred_c paddle.float32 [64, 16, 1]
        eta_dot_pred_c paddle.float32 [64, 16, 1]
        lift_pred_c paddle.float32 [64, 16, 1]
        >>> # model_type is `3`
        >>> model = ppsci.arch.DeepPhyLSTM(
        ...     input_size=16,
        ...     output_size=1,
        ...     hidden_size=100,
        ...     model_type=3)
        >>> out = model(
        ...     {"ag":paddle.rand([64, 16, 1]),
        ...     "ag_c":paddle.rand([64, 16, 1]),
        ...     "phi":paddle.rand([1, 16, 16])})
        >>> for k, v in out.items():
        ...     print(f"{k} {v.dtype} {v.shape}")
        eta_pred paddle.float32 [64, 16, 1]
        eta_dot_pred paddle.float32 [64, 16, 1]
        g_pred paddle.float32 [64, 16, 1]
        eta_t_pred_c paddle.float32 [64, 16, 1]
        eta_dot_pred_c paddle.float32 [64, 16, 1]
        lift_pred_c paddle.float32 [64, 16, 1]
        g_t_pred_c paddle.float32 [64, 16, 1]
        g_dot_pred_c paddle.float32 [64, 16, 1]
    """

    def __init__(self, input_size, output_size, hidden_size=100, model_type=2):
        super().__init__()
        self.input_size = input_size
        self.output_size = output_size
        self.hidden_size = hidden_size
        self.model_type = model_type

        if self.model_type == 2:
            self.lstm_model = nn.Sequential(
                nn.LSTM(input_size, hidden_size),
                nn.ReLU(),
                nn.LSTM(hidden_size, hidden_size),
                nn.ReLU(),
                nn.LSTM(hidden_size, hidden_size),
                nn.ReLU(),
                nn.Linear(hidden_size, hidden_size),
                nn.Linear(hidden_size, 3 * output_size),
            )

            self.lstm_model_f = nn.Sequential(
                nn.LSTM(3 * output_size, hidden_size),
                nn.ReLU(),
                nn.LSTM(hidden_size, hidden_size),
                nn.ReLU(),
                nn.LSTM(hidden_size, hidden_size),
                nn.ReLU(),
                nn.Linear(hidden_size, hidden_size),
                nn.Linear(hidden_size, output_size),
            )
        elif self.model_type == 3:
            self.lstm_model = nn.Sequential(
                nn.LSTM(1, hidden_size),
                nn.ReLU(),
                nn.LSTM(hidden_size, hidden_size),
                nn.ReLU(),
                nn.LSTM(hidden_size, hidden_size),
                nn.ReLU(),
                nn.Linear(hidden_size, 3 * output_size),
            )

            self.lstm_model_f = nn.Sequential(
                nn.LSTM(3 * output_size, hidden_size),
                nn.ReLU(),
                nn.LSTM(hidden_size, hidden_size),
                nn.ReLU(),
                nn.LSTM(hidden_size, hidden_size),
                nn.ReLU(),
                nn.Linear(hidden_size, output_size),
            )

            self.lstm_model_g = nn.Sequential(
                nn.LSTM(2 * output_size, hidden_size),
                nn.ReLU(),
                nn.LSTM(hidden_size, hidden_size),
                nn.ReLU(),
                nn.LSTM(hidden_size, hidden_size),
                nn.ReLU(),
                nn.Linear(hidden_size, output_size),
            )
        else:
            raise ValueError(f"model_type should be 2 or 3, but got {model_type}")

    def forward(self, x):
        if self._input_transform is not None:
            x = self._input_transform(x)

        if self.model_type == 2:
            result_dict = self._forward_type_2(x)
        elif self.model_type == 3:
            result_dict = self._forward_type_3(x)
        if self._output_transform is not None:
            result_dict = self._output_transform(x, result_dict)
        return result_dict

    def _forward_type_2(self, x):
        output = x["ag"]
        for layer in self.lstm_model:
            output = layer(output)
            if isinstance(output, tuple):
                output = output[0]

        eta_pred = output[:, :, 0 : self.output_size]
        eta_dot_pred = output[:, :, self.output_size : 2 * self.output_size]
        g_pred = output[:, :, 2 * self.output_size :]

        # for ag_c
        output_c = x["ag_c"]
        for layer in self.lstm_model:
            output_c = layer(output_c)
            if isinstance(output_c, tuple):
                output_c = output_c[0]

        eta_pred_c = output_c[:, :, 0 : self.output_size]
        eta_dot_pred_c = output_c[:, :, self.output_size : 2 * self.output_size]
        g_pred_c = output_c[:, :, 2 * self.output_size :]
        eta_t_pred_c = paddle.matmul(x["phi"], eta_pred_c)
        eta_tt_pred_c = paddle.matmul(x["phi"], eta_dot_pred_c)
        eta_dot1_pred_c = eta_dot_pred_c[:, :, 0:1]
        tmp = paddle.concat([eta_pred_c, eta_dot1_pred_c, g_pred_c], 2)
        f = tmp
        for layer in self.lstm_model_f:
            f = layer(f)
            if isinstance(f, tuple):
                f = f[0]

        lift_pred_c = eta_tt_pred_c + f

        return {
            "eta_pred": eta_pred,
            "eta_dot_pred": eta_dot_pred,
            "g_pred": g_pred,
            "eta_t_pred_c": eta_t_pred_c,
            "eta_dot_pred_c": eta_dot_pred_c,
            "lift_pred_c": lift_pred_c,
        }

    def _forward_type_3(self, x):
        # physics informed neural networks
        output = x["ag"]
        for layer in self.lstm_model:
            output = layer(output)
            if isinstance(output, tuple):
                output = output[0]

        eta_pred = output[:, :, 0 : self.output_size]
        eta_dot_pred = output[:, :, self.output_size : 2 * self.output_size]
        g_pred = output[:, :, 2 * self.output_size :]

        output_c = x["ag_c"]
        for layer in self.lstm_model:
            output_c = layer(output_c)
            if isinstance(output_c, tuple):
                output_c = output_c[0]

        eta_pred_c = output_c[:, :, 0 : self.output_size]
        eta_dot_pred_c = output_c[:, :, self.output_size : 2 * self.output_size]
        g_pred_c = output_c[:, :, 2 * self.output_size :]

        eta_t_pred_c = paddle.matmul(x["phi"], eta_pred_c)
        eta_tt_pred_c = paddle.matmul(x["phi"], eta_dot_pred_c)
        g_t_pred_c = paddle.matmul(x["phi"], g_pred_c)

        f = paddle.concat([eta_pred_c, eta_dot_pred_c, g_pred_c], 2)
        for layer in self.lstm_model_f:
            f = layer(f)
            if isinstance(f, tuple):
                f = f[0]

        lift_pred_c = eta_tt_pred_c + f

        eta_dot1_pred_c = eta_dot_pred_c[:, :, 0:1]
        g_dot_pred_c = paddle.concat([eta_dot1_pred_c, g_pred_c], 2)
        for layer in self.lstm_model_g:
            g_dot_pred_c = layer(g_dot_pred_c)
            if isinstance(g_dot_pred_c, tuple):
                g_dot_pred_c = g_dot_pred_c[0]

        return {
            "eta_pred": eta_pred,
            "eta_dot_pred": eta_dot_pred,
            "g_pred": g_pred,
            "eta_t_pred_c": eta_t_pred_c,
            "eta_dot_pred_c": eta_dot_pred_c,
            "lift_pred_c": lift_pred_c,
            "g_t_pred_c": g_t_pred_c,
            "g_dot_pred_c": g_dot_pred_c,
        }

DGMR

Bases: Arch

Deep Generative Model of Radar. Nowcasting GAN is an attempt to recreate DeepMind's Skillful Nowcasting GAN from https://arxiv.org/abs/2104.00954. but slightly modified for multiple satellite channels

Parameters:

Name	Type	Description	Default
input_keys	Tuple[str, ...]	Name of input keys, such as ("input",).	required
output_keys	Tuple[str, ...]	Name of output keys, such as ("output",).	required
forecast_steps	int	Number of steps to predict in the future	18
input_channels	int	Number of input channels per image	1
gen_lr	float	Learning rate for the generator	5e-05
disc_lr	float	Learning rate for the discriminators, shared for both temporal and spatial discriminator	0.0002
conv_type	str	Type of 2d convolution to use, see satflow/models/utils.py for options	'standard'
beta1	float	Beta1 for Adam optimizer	0.0
beta2	float	Beta2 for Adam optimizer	0.999
num_samples	int	Number of samples of the latent space to sample for training/validation	6
grid_lambda	float	Lambda for the grid regularization loss	20.0
output_shape	int	Shape of the output predictions, generally should be same as the input shape	256
generation_steps	int	Number of generation steps to use in forward pass, in paper is 6 and the best is chosen for the loss this results in huge amounts of GPU memory though, so less might work better for training.	6
context_channels	int	Number of output channels for the lowest block of conditioning stack	384
latent_channels	int	Number of channels that the latent space should be reshaped to, input dimension into ConvGRU, also affects the number of channels for other linked inputs/outputs	768

Examples:

>>> import ppsci
>>> import paddle
>>> model = ppsci.arch.DGMR(("input", ), ("output", ))
>>> input_dict = {"input": paddle.randn((1, 4, 1, 256, 256))}
>>> output_dict = model(input_dict)
>>> print(output_dict["output"].shape)
[1, 18, 1, 256, 256]

Source code in ppsci/arch/dgmr.py

class DGMR(base.Arch):
    """Deep Generative Model of Radar.
        Nowcasting GAN is an attempt to recreate DeepMind's Skillful Nowcasting GAN from https://arxiv.org/abs/2104.00954.
        but slightly modified for multiple satellite channels

    Args:
        input_keys (Tuple[str, ...]): Name of input keys, such as ("input",).
        output_keys (Tuple[str, ...]): Name of output keys, such as ("output",).
        forecast_steps (int, optional): Number of steps to predict in the future
        input_channels (int, optional): Number of input channels per image
        gen_lr (float, optional): Learning rate for the generator
        disc_lr (float, optional): Learning rate for the discriminators, shared for both temporal and spatial discriminator
        conv_type (str, optional): Type of 2d convolution to use, see satflow/models/utils.py for options
        beta1 (float, optional): Beta1 for Adam optimizer
        beta2 (float, optional): Beta2 for Adam optimizer
        num_samples (int, optional): Number of samples of the latent space to sample for training/validation
        grid_lambda (float, optional): Lambda for the grid regularization loss
        output_shape (int, optional): Shape of the output predictions, generally should be same as the input shape
        generation_steps (int, optional): Number of generation steps to use in forward pass, in paper is 6 and the best is chosen for the loss
            this results in huge amounts of GPU memory though, so less might work better for training.
        context_channels (int, optional): Number of output channels for the lowest block of conditioning stack
        latent_channels (int, optional): Number of channels that the latent space should be reshaped to,
            input dimension into ConvGRU, also affects the number of channels for other linked inputs/outputs

    Examples:
        >>> import ppsci
        >>> import paddle
        >>> model = ppsci.arch.DGMR(("input", ), ("output", ))
        >>> input_dict = {"input": paddle.randn((1, 4, 1, 256, 256))}
        >>> output_dict = model(input_dict) # doctest: +SKIP
        >>> print(output_dict["output"].shape) # doctest: +SKIP
        [1, 18, 1, 256, 256]
    """

    def __init__(
        self,
        input_keys: Tuple[str, ...],
        output_keys: Tuple[str, ...],
        forecast_steps: int = 18,
        input_channels: int = 1,
        output_shape: int = 256,
        gen_lr: float = 5e-05,
        disc_lr: float = 0.0002,
        conv_type: str = "standard",
        num_samples: int = 6,
        grid_lambda: float = 20.0,
        beta1: float = 0.0,
        beta2: float = 0.999,
        latent_channels: int = 768,
        context_channels: int = 384,
        generation_steps: int = 6,
    ):
        super().__init__()
        self.input_keys = input_keys
        self.output_keys = output_keys
        self.gen_lr = gen_lr
        self.disc_lr = disc_lr
        self.beta1 = beta1
        self.beta2 = beta2
        self.grid_lambda = grid_lambda
        self.num_samples = num_samples
        self.latent_channels = latent_channels
        self.context_channels = context_channels
        self.input_channels = input_channels
        self.generation_steps = generation_steps
        self.conditioning_stack = ContextConditioningStack(
            input_channels=input_channels,
            conv_type=conv_type,
            output_channels=self.context_channels,
        )
        self.latent_stack = LatentConditioningStack(
            shape=(8 * self.input_channels, output_shape // 32, output_shape // 32),
            output_channels=self.latent_channels,
        )
        self.sampler = Sampler(
            forecast_steps=forecast_steps,
            latent_channels=self.latent_channels,
            context_channels=self.context_channels,
        )
        self.generator = Generator(
            self.conditioning_stack, self.latent_stack, self.sampler
        )
        self.discriminator = Discriminator(input_channels)
        self.global_iteration = 0
        self.automatic_optimization = False

    def split_to_dict(
        self, data_tensors: Tuple[paddle.Tensor, ...], keys: Tuple[str, ...]
    ):
        return {key: data_tensors[i] for i, key in enumerate(keys)}

    def forward(self, x):
        if self._input_transform is not None:
            x = self._input_transform(x)
        x_tensor = self.concat_to_tensor(x, self.input_keys)
        y = [self.generator(x_tensor)]
        y = self.split_to_dict(y, self.output_keys)

        if self._output_transform is not None:
            y = self._output_transform(x, y)
        return y

Discriminator

Bases: Arch

Discriminator Net of GAN.

Parameters:

Name	Type	Description	Default
input_keys	Tuple[str, ...]	Name of input keys, such as ("input1", "input2").	required
output_keys	Tuple[str, ...]	Name of output keys, such as ("output1", "output2").	required
in_channel	int	Number of input channels of the first conv layer.	required
out_channels	Tuple[int, ...]	Number of output channels of all conv layers, such as (out_conv0, out_conv1, out_conv2).	required
fc_channel	int	Number of input features of linear layer. Number of output features of the layer is set to 1 in this Net to construct a fully_connected layer.	required
kernel_sizes	Tuple[int, ...]	Number of kernel_size of all conv layers, such as (kernel_size_conv0, kernel_size_conv1, kernel_size_conv2).	required
strides	Tuple[int, ...]	Number of stride of all conv layers, such as (stride_conv0, stride_conv1, stride_conv2).	required
use_bns	Tuple[bool, ...]	Whether to use the batch_norm layer after each conv layer.	required
acts	Tuple[str, ...]	Whether to use the activation layer after each conv layer. If so, witch activation to use, such as (act_conv0, act_conv1, act_conv2).	required

Examples:

>>> import ppsci
>>> in_channel = 2
>>> in_channel_tempo = 3
>>> out_channels = (32, 64, 128, 256)
>>> fc_channel = 65536
>>> kernel_sizes = ((4, 4), (4, 4), (4, 4), (4, 4))
>>> strides = (2, 2, 2, 1)
>>> use_bns = (False, True, True, True)
>>> acts = ("leaky_relu", "leaky_relu", "leaky_relu", "leaky_relu", None)
>>> output_keys_disc = ("out_1", "out_2", "out_3", "out_4", "out_5", "out_6", "out_7", "out_8", "out_9", "out_10")
>>> model = ppsci.arch.Discriminator(("in_1","in_2"), output_keys_disc, in_channel, out_channels, fc_channel, kernel_sizes, strides, use_bns, acts)
>>> input_data = [paddle.to_tensor(paddle.randn([1, in_channel, 128, 128])),paddle.to_tensor(paddle.randn([1, in_channel, 128, 128]))]
>>> input_dict = {"in_1": input_data[0],"in_2": input_data[1]}
>>> out_dict = model(input_dict)
>>> for k, v in out_dict.items():
...     print(k, v.shape)
out_1 [1, 32, 64, 64]
out_2 [1, 64, 32, 32]
out_3 [1, 128, 16, 16]
out_4 [1, 256, 16, 16]
out_5 [1, 1]
out_6 [1, 32, 64, 64]
out_7 [1, 64, 32, 32]
out_8 [1, 128, 16, 16]
out_9 [1, 256, 16, 16]
out_10 [1, 1]

Source code in ppsci/arch/gan.py

class Discriminator(base.Arch):
    """Discriminator Net of GAN.

    Args:
        input_keys (Tuple[str, ...]): Name of input keys, such as ("input1", "input2").
        output_keys (Tuple[str, ...]): Name of output keys, such as ("output1", "output2").
        in_channel (int):  Number of input channels of the first conv layer.
        out_channels (Tuple[int, ...]): Number of output channels of all conv layers,
            such as (out_conv0, out_conv1, out_conv2).
        fc_channel (int):  Number of input features of linear layer. Number of output features of the layer
            is set to 1 in this Net to construct a fully_connected layer.
        kernel_sizes (Tuple[int, ...]): Number of kernel_size of all conv layers,
            such as (kernel_size_conv0, kernel_size_conv1, kernel_size_conv2).
        strides (Tuple[int, ...]): Number of stride of all conv layers,
            such as (stride_conv0, stride_conv1, stride_conv2).
        use_bns (Tuple[bool, ...]): Whether to use the batch_norm layer after each conv layer.
        acts (Tuple[str, ...]): Whether to use the activation layer after each conv layer. If so, witch activation to use,
            such as (act_conv0, act_conv1, act_conv2).

    Examples:
        >>> import ppsci
        >>> in_channel = 2
        >>> in_channel_tempo = 3
        >>> out_channels = (32, 64, 128, 256)
        >>> fc_channel = 65536
        >>> kernel_sizes = ((4, 4), (4, 4), (4, 4), (4, 4))
        >>> strides = (2, 2, 2, 1)
        >>> use_bns = (False, True, True, True)
        >>> acts = ("leaky_relu", "leaky_relu", "leaky_relu", "leaky_relu", None)
        >>> output_keys_disc = ("out_1", "out_2", "out_3", "out_4", "out_5", "out_6", "out_7", "out_8", "out_9", "out_10")
        >>> model = ppsci.arch.Discriminator(("in_1","in_2"), output_keys_disc, in_channel, out_channels, fc_channel, kernel_sizes, strides, use_bns, acts)
        >>> input_data = [paddle.to_tensor(paddle.randn([1, in_channel, 128, 128])),paddle.to_tensor(paddle.randn([1, in_channel, 128, 128]))]
        >>> input_dict = {"in_1": input_data[0],"in_2": input_data[1]}
        >>> out_dict = model(input_dict)
        >>> for k, v in out_dict.items():
        ...     print(k, v.shape)
        out_1 [1, 32, 64, 64]
        out_2 [1, 64, 32, 32]
        out_3 [1, 128, 16, 16]
        out_4 [1, 256, 16, 16]
        out_5 [1, 1]
        out_6 [1, 32, 64, 64]
        out_7 [1, 64, 32, 32]
        out_8 [1, 128, 16, 16]
        out_9 [1, 256, 16, 16]
        out_10 [1, 1]
    """

    def __init__(
        self,
        input_keys: Tuple[str, ...],
        output_keys: Tuple[str, ...],
        in_channel: int,
        out_channels: Tuple[int, ...],
        fc_channel: int,
        kernel_sizes: Tuple[int, ...],
        strides: Tuple[int, ...],
        use_bns: Tuple[bool, ...],
        acts: Tuple[str, ...],
    ):
        super().__init__()
        self.input_keys = input_keys
        self.output_keys = output_keys
        self.in_channel = in_channel
        self.out_channels = out_channels
        self.fc_channel = fc_channel
        self.kernel_sizes = kernel_sizes
        self.strides = strides
        self.use_bns = use_bns
        self.acts = acts

        self.init_layers()

    def init_layers(self):
        layers_list = []
        for i in range(len(self.out_channels)):
            in_channel = self.in_channel if i == 0 else self.out_channels[i - 1]
            layers_list.append(
                Conv2DBlock(
                    in_channel=in_channel,
                    out_channel=self.out_channels[i],
                    kernel_size=self.kernel_sizes[i],
                    stride=self.strides[i],
                    use_bn=self.use_bns[i],
                    act=self.acts[i],
                    mean=0.0,
                    std=0.04,
                    value=0.1,
                )
            )

        layers_list.append(
            FCBlock(self.fc_channel, self.acts[4], mean=0.0, std=0.04, value=0.1)
        )
        self.layers = nn.LayerList(layers_list)

    def forward_tensor(self, x):
        y = x
        y_list = []
        for layer in self.layers:
            y = layer(y)
            y_list.append(y)
        return y_list  # y_conv1, y_conv2, y_conv3, y_conv4, y_fc(y_out)

    def forward(self, x):
        if self._input_transform is not None:
            x = self._input_transform(x)

        y_list = []
        # y1_conv1, y1_conv2, y1_conv3, y1_conv4, y1_fc, y2_conv1, y2_conv2, y2_conv3, y2_conv4, y2_fc
        for k in x:
            y_list.extend(self.forward_tensor(x[k]))

        y = self.split_to_dict(y_list, self.output_keys)

        if self._output_transform is not None:
            y = self._output_transform(x, y)

        return y

    @staticmethod
    def split_to_dict(
        data_list: List[paddle.Tensor], keys: Tuple[str, ...]
    ) -> Dict[str, paddle.Tensor]:
        """Overwrite of split_to_dict() method belongs to Class base.Arch.

        Reason for overwriting is there is no concat_to_tensor() method called in "tempoGAN" example.
        That is because input in "tempoGAN" example is not in a regular format, but a format like:
        {
            "input1": paddle.concat([in1, in2], axis=1),
            "input2": paddle.concat([in1, in3], axis=1),
        }

        Args:
            data_list (List[paddle.Tensor]): The data to be split. It should be a list of tensor(s), but not a paddle.Tensor.
            keys (Tuple[str, ...]): Keys of outputs.

        Returns:
            Dict[str, paddle.Tensor]: Dict with split data.
        """
        if len(keys) == 1:
            return {keys[0]: data_list[0]}
        return {key: data_list[i] for i, key in enumerate(keys)}

split_to_dict(data_list, keys) staticmethod

Overwrite of split_to_dict() method belongs to Class base.Arch.

Reason for overwriting is there is no concat_to_tensor() method called in "tempoGAN" example. That is because input in "tempoGAN" example is not in a regular format, but a format like: { "input1": paddle.concat([in1, in2], axis=1), "input2": paddle.concat([in1, in3], axis=1), }

Parameters:

Name	Type	Description	Default
data_list	List[Tensor]	The data to be split. It should be a list of tensor(s), but not a paddle.Tensor.	required
keys	Tuple[str, ...]	Keys of outputs.	required

Returns:

Type	Description
Dict[str, Tensor]	Dict[str, paddle.Tensor]: Dict with split data.

Source code in ppsci/arch/gan.py

@staticmethod
def split_to_dict(
    data_list: List[paddle.Tensor], keys: Tuple[str, ...]
) -> Dict[str, paddle.Tensor]:
    """Overwrite of split_to_dict() method belongs to Class base.Arch.

    Reason for overwriting is there is no concat_to_tensor() method called in "tempoGAN" example.
    That is because input in "tempoGAN" example is not in a regular format, but a format like:
    {
        "input1": paddle.concat([in1, in2], axis=1),
        "input2": paddle.concat([in1, in3], axis=1),
    }

    Args:
        data_list (List[paddle.Tensor]): The data to be split. It should be a list of tensor(s), but not a paddle.Tensor.
        keys (Tuple[str, ...]): Keys of outputs.

    Returns:
        Dict[str, paddle.Tensor]: Dict with split data.
    """
    if len(keys) == 1:
        return {keys[0]: data_list[0]}
    return {key: data_list[i] for i, key in enumerate(keys)}

ExtFormerMoECuboid

Bases: Arch

Cuboid Transformer for spatiotemporal forecasting

We adopt the Non-autoregressive encoder-decoder architecture. The decoder takes the multi-scale memory output from the encoder.

The initial downsampling / upsampling layers will be Downsampling: [K x Conv2D --> PatchMerge] Upsampling: [Nearest Interpolation-based Upsample --> K x Conv2D]

x --> downsample (optional) ---> (+pos_embed) ---> enc --> mem_l initial_z (+pos_embed) ---> FC | | |------------| | | y <--- upsample (optional) <--- dec <----------

Parameters:

Name	Type	Description	Default
input_keys	Tuple[str, ...]	Name of input keys, such as ("input",).	required
output_keys	Tuple[str, ...]	Name of output keys, such as ("output",).	required
input_shape	Tuple[int, ...]	The shape of the input data.	required
target_shape	Tuple[int, ...]	The shape of the target data.	required
base_units	int	The base units. Defaults to 128.	128
block_units	int	The block units. Defaults to None.	None
scale_alpha	float	We scale up the channels based on the formula: - round_to(base_units * max(downsample_scale) ** units_alpha, 4). Defaults to 1.0.	1.0
num_heads	int	The number of heads. Defaults to 4.	4
attn_drop	float	The attention dropout. Defaults to 0.0.	0.0
proj_drop	float	The projection dropout. Defaults to 0.0.	0.0
ffn_drop	float	The ffn dropout. Defaults to 0.0.	0.0
downsample	int	The rate of downsample. Defaults to 2.	2
downsample_type	str	The type of downsample. Defaults to "patch_merge".	'patch_merge'
upsample_type	str	The rate of upsample. Defaults to "upsample".	'upsample'
upsample_kernel_size	int	The kernel size of upsample. Defaults to 3.	3
enc_depth	list	The depth of encoder. Defaults to [4, 4, 4].	[4, 4, 4]
enc_attn_patterns	str	The pattern of encoder attention. Defaults to None.	None
enc_cuboid_size	list	The cuboid size of encoder. Defaults to [(4, 4, 4), (4, 4, 4)].	[(4, 4, 4), (4, 4, 4)]
enc_cuboid_strategy	list	The cuboid strategy of encoder. Defaults to [("l", "l", "l"), ("d", "d", "d")].	[('l', 'l', 'l'), ('d', 'd', 'd')]
enc_shift_size	list	The shift size of encoder. Defaults to [(0, 0, 0), (0, 0, 0)].	[(0, 0, 0), (0, 0, 0)]
enc_use_inter_ffn	bool	Whether to use intermediate FFN for encoder. Defaults to True.	True
dec_depth	list	The depth of decoder. Defaults to [2, 2].	[2, 2]
dec_cross_start	int	The cross start of decoder. Defaults to 0.	0
dec_self_attn_patterns	str	The partterns of decoder. Defaults to None.	None
dec_self_cuboid_size	list	The cuboid size of decoder. Defaults to [(4, 4, 4), (4, 4, 4)].	[(4, 4, 4), (4, 4, 4)]
dec_self_cuboid_strategy	list	The strategy of decoder. Defaults to [("l", "l", "l"), ("d", "d", "d")].	[('l', 'l', 'l'), ('d', 'd', 'd')]
dec_self_shift_size	list	The shift size of decoder. Defaults to [(1, 1, 1), (0, 0, 0)].	[(1, 1, 1), (0, 0, 0)]
dec_cross_attn_patterns	_type_	The cross attention patterns of decoder. Defaults to None.	None
dec_cross_cuboid_hw	list	The cuboid_hw of decoder. Defaults to [(4, 4), (4, 4)].	[(4, 4), (4, 4)]
dec_cross_cuboid_strategy	list	The cuboid strategy of decoder. Defaults to [("l", "l", "l"), ("d", "l", "l")].	[('l', 'l', 'l'), ('d', 'l', 'l')]
dec_cross_shift_hw	list	The shift_hw of decoder. Defaults to [(0, 0), (0, 0)].	[(0, 0), (0, 0)]
dec_cross_n_temporal	list	The cross_n_temporal of decoder. Defaults to [1, 2].	[1, 2]
dec_cross_last_n_frames	int	The cross_last_n_frames of decoder. Defaults to None.	None
dec_use_inter_ffn	bool	Whether to use intermediate FFN for decoder. Defaults to True.	True
dec_hierarchical_pos_embed	bool	Whether to use hierarchical pos_embed for decoder. Defaults to False.	False
num_global_vectors	int	The num of global vectors. Defaults to 4.	4
use_dec_self_global	bool	Whether to use global vector for decoder. Defaults to True.	True
dec_self_update_global	bool	Whether to update global vector for decoder. Defaults to True.	True
use_dec_cross_global	bool	Whether to use cross global vector for decoder. Defaults to True.	True
use_global_vector_ffn	bool	Whether to use global vector FFN. Defaults to True.	True
use_global_self_attn	bool	Whether to use global attentions. Defaults to False.	False
separate_global_qkv	bool	Whether to separate global qkv. Defaults to False.	False
global_dim_ratio	int	The ratio of global dim. Defaults to 1.	1
self_pattern	str	The pattern. Defaults to "axial".	'axial'
cross_self_pattern	str	The self cross pattern. Defaults to "axial".	'axial'
cross_pattern	str	The cross pattern. Defaults to "cross_1x1".	'cross_1x1'
z_init_method	str	How the initial input to the decoder is initialized. Defaults to "nearest_interp".	'nearest_interp'
initial_downsample_type	str	The downsample type of initial. Defaults to "conv".	'conv'
initial_downsample_activation	str	The downsample activation of initial. Defaults to "leaky".	'leaky'
initial_downsample_scale	int	The downsample scale of initial. Defaults to 1.	1
initial_downsample_conv_layers	int	The conv layer of downsample of initial. Defaults to 2.	2
final_upsample_conv_layers	int	The conv layer of final upsample. Defaults to 2.	2
initial_downsample_stack_conv_num_layers	int	The num of stack conv layer of initial downsample. Defaults to 1.	1
initial_downsample_stack_conv_dim_list	list	The dim list of stack conv of initial downsample. Defaults to None.	None
initial_downsample_stack_conv_downscale_list	list	The downscale list of stack conv of initial downsample. Defaults to [1].	[1]
initial_downsample_stack_conv_num_conv_list	list	The num of stack conv list of initial downsample. Defaults to [2].	[2]
ffn_activation	str	The activation of FFN. Defaults to "leaky".	'leaky'
gated_ffn	bool	Whether to use gate FFN. Defaults to False.	False
norm_layer	str	The type of normilize. Defaults to "layer_norm".	'layer_norm'
padding_type	str	The type of padding. Defaults to "ignore".	'ignore'
pos_embed_type	str	The type of pos embedding. Defaults to "t+hw".	't+hw'
checkpoint_level	bool	Whether to use checkpoint. Defaults to True.	True
use_relative_pos	bool	Whether to use relative pose. Defaults to True.	True
self_attn_use_final_proj	bool	Whether to use final projection. Defaults to True.	True
dec_use_first_self_attn	bool	Whether to use first self attention for decoder. Defaults to False.	False
attn_linear_init_mode	str	The mode of attention linear init. Defaults to "0".	'0'
ffn_linear_init_mode	str	The mode of FFN linear init. Defaults to "0".	'0'
conv_init_mode	str	The mode of conv init. Defaults to "0".	'0'
down_up_linear_init_mode	str	The mode of downsample and upsample linear init. Defaults to "0".	'0'
norm_init_mode	str	The mode of normalization init. Defaults to "0".	'0'

Source code in ppsci/arch/extformer_moe_cuboid.py

class ExtFormerMoECuboid(base.Arch):
    """Cuboid Transformer for spatiotemporal forecasting

    We adopt the Non-autoregressive encoder-decoder architecture.
    The decoder takes the multi-scale memory output from the encoder.

    The initial downsampling / upsampling layers will be
    Downsampling: [K x Conv2D --> PatchMerge]
    Upsampling: [Nearest Interpolation-based Upsample --> K x Conv2D]

    x --> downsample (optional) ---> (+pos_embed) ---> enc --> mem_l         initial_z (+pos_embed) ---> FC
                                                     |            |
                                                     |------------|
                                                           |
                                                           |
             y <--- upsample (optional) <--- dec <----------

    Args:
        input_keys (Tuple[str, ...]): Name of input keys, such as ("input",).
        output_keys (Tuple[str, ...]): Name of output keys, such as ("output",).
        input_shape (Tuple[int, ...]): The shape of the input data.
        target_shape (Tuple[int, ...]): The shape of the target data.
        base_units (int, optional): The base units. Defaults to 128.
        block_units (int, optional): The block units. Defaults to None.
        scale_alpha (float, optional): We scale up the channels based on the formula:
            - round_to(base_units * max(downsample_scale) ** units_alpha, 4). Defaults to 1.0.
        num_heads (int, optional): The number of heads. Defaults to 4.
        attn_drop (float, optional): The attention dropout. Defaults to 0.0.
        proj_drop (float, optional): The projection dropout. Defaults to 0.0.
        ffn_drop (float, optional): The ffn dropout. Defaults to 0.0.
        downsample (int, optional): The rate of downsample. Defaults to 2.
        downsample_type (str, optional): The type of downsample. Defaults to "patch_merge".
        upsample_type (str, optional): The rate of upsample. Defaults to "upsample".
        upsample_kernel_size (int, optional): The kernel size of upsample. Defaults to 3.
        enc_depth (list, optional): The depth of encoder. Defaults to [4, 4, 4].
        enc_attn_patterns (str, optional): The pattern of encoder attention. Defaults to None.
        enc_cuboid_size (list, optional): The cuboid size of encoder. Defaults to [(4, 4, 4), (4, 4, 4)].
        enc_cuboid_strategy (list, optional): The cuboid strategy of encoder. Defaults to [("l", "l", "l"), ("d", "d", "d")].
        enc_shift_size (list, optional): The shift size of encoder. Defaults to [(0, 0, 0), (0, 0, 0)].
        enc_use_inter_ffn (bool, optional): Whether to use intermediate FFN for encoder. Defaults to True.
        dec_depth (list, optional): The depth of decoder. Defaults to [2, 2].
        dec_cross_start (int, optional): The cross start of decoder. Defaults to 0.
        dec_self_attn_patterns (str, optional): The partterns of decoder. Defaults to None.
        dec_self_cuboid_size (list, optional): The cuboid size of decoder. Defaults to [(4, 4, 4), (4, 4, 4)].
        dec_self_cuboid_strategy (list, optional): The strategy of decoder. Defaults to [("l", "l", "l"), ("d", "d", "d")].
        dec_self_shift_size (list, optional): The shift size of decoder. Defaults to [(1, 1, 1), (0, 0, 0)].
        dec_cross_attn_patterns (_type_, optional): The cross attention patterns of decoder. Defaults to None.
        dec_cross_cuboid_hw (list, optional): The cuboid_hw of decoder. Defaults to [(4, 4), (4, 4)].
        dec_cross_cuboid_strategy (list, optional): The cuboid strategy of decoder. Defaults to [("l", "l", "l"), ("d", "l", "l")].
        dec_cross_shift_hw (list, optional): The shift_hw of decoder. Defaults to [(0, 0), (0, 0)].
        dec_cross_n_temporal (list, optional): The cross_n_temporal of decoder. Defaults to [1, 2].
        dec_cross_last_n_frames (int, optional): The cross_last_n_frames of decoder. Defaults to None.
        dec_use_inter_ffn (bool, optional): Whether to use intermediate FFN for decoder. Defaults to True.
        dec_hierarchical_pos_embed (bool, optional): Whether to use hierarchical pos_embed for decoder. Defaults to False.
        num_global_vectors (int, optional): The num of global vectors. Defaults to 4.
        use_dec_self_global (bool, optional): Whether to use global vector for decoder. Defaults to True.
        dec_self_update_global (bool, optional): Whether to update global vector for decoder. Defaults to True.
        use_dec_cross_global (bool, optional): Whether to use cross global vector for decoder. Defaults to True.
        use_global_vector_ffn (bool, optional): Whether to use global vector FFN. Defaults to True.
        use_global_self_attn (bool, optional): Whether to use global attentions. Defaults to False.
        separate_global_qkv (bool, optional): Whether to separate global qkv. Defaults to False.
        global_dim_ratio (int, optional): The ratio of global dim. Defaults to 1.
        self_pattern (str, optional): The pattern. Defaults to "axial".
        cross_self_pattern (str, optional): The self cross pattern. Defaults to "axial".
        cross_pattern (str, optional): The cross pattern. Defaults to "cross_1x1".
        z_init_method (str, optional): How the initial input to the decoder is initialized. Defaults to "nearest_interp".
        initial_downsample_type (str, optional): The downsample type of initial. Defaults to "conv".
        initial_downsample_activation (str, optional): The downsample activation of initial. Defaults to "leaky".
        initial_downsample_scale (int, optional): The downsample scale of initial. Defaults to 1.
        initial_downsample_conv_layers (int, optional): The conv layer of downsample of initial. Defaults to 2.
        final_upsample_conv_layers (int, optional): The conv layer of final upsample. Defaults to 2.
        initial_downsample_stack_conv_num_layers (int, optional): The num of stack conv layer of initial downsample. Defaults to 1.
        initial_downsample_stack_conv_dim_list (list, optional): The dim list of stack conv of initial downsample. Defaults to None.
        initial_downsample_stack_conv_downscale_list (list, optional): The downscale list of stack conv of initial downsample. Defaults to [1].
        initial_downsample_stack_conv_num_conv_list (list, optional): The num of stack conv list of initial downsample. Defaults to [2].
        ffn_activation (str, optional): The activation of FFN. Defaults to "leaky".
        gated_ffn (bool, optional): Whether to use gate FFN. Defaults to False.
        norm_layer (str, optional): The type of normilize. Defaults to "layer_norm".
        padding_type (str, optional): The type of padding. Defaults to "ignore".
        pos_embed_type (str, optional): The type of pos embedding. Defaults to "t+hw".
        checkpoint_level (bool, optional): Whether to use checkpoint. Defaults to True.
        use_relative_pos (bool, optional): Whether to use relative pose. Defaults to True.
        self_attn_use_final_proj (bool, optional): Whether to use final projection. Defaults to True.
        dec_use_first_self_attn (bool, optional): Whether to use first self attention for decoder. Defaults to False.
        attn_linear_init_mode (str, optional): The mode of attention linear init. Defaults to "0".
        ffn_linear_init_mode (str, optional): The mode of FFN linear init. Defaults to "0".
        conv_init_mode (str, optional): The mode of conv init. Defaults to "0".
        down_up_linear_init_mode (str, optional): The mode of downsample and upsample linear init. Defaults to "0".
        norm_init_mode (str, optional): The mode of normalization init. Defaults to "0".
    """

    def __init__(
        self,
        input_keys: Tuple[str, ...],
        output_keys: Tuple[str, ...],
        input_shape: Tuple[int, ...],
        target_shape: Tuple[int, ...],
        base_units: int = 128,
        block_units: int = None,
        scale_alpha: float = 1.0,
        num_heads: int = 4,
        attn_drop: float = 0.0,
        proj_drop: float = 0.0,
        ffn_drop: float = 0.0,
        downsample: int = 2,
        downsample_type: str = "patch_merge",
        upsample_type: str = "upsample",
        upsample_kernel_size: int = 3,
        enc_depth: Tuple[int, ...] = [4, 4, 4],
        enc_attn_patterns: str = None,
        enc_cuboid_size: Tuple[Tuple[int, ...], ...] = [(4, 4, 4), (4, 4, 4)],
        enc_cuboid_strategy: Tuple[Tuple[str, ...], ...] = [
            ("l", "l", "l"),
            ("d", "d", "d"),
        ],
        enc_shift_size: Tuple[Tuple[int, ...], ...] = [(0, 0, 0), (0, 0, 0)],
        enc_use_inter_ffn: bool = True,
        dec_depth: Tuple[int, ...] = [2, 2],
        dec_cross_start: int = 0,
        dec_self_attn_patterns: str = None,
        dec_self_cuboid_size: Tuple[Tuple[int, ...], ...] = [(4, 4, 4), (4, 4, 4)],
        dec_self_cuboid_strategy: Tuple[Tuple[str, ...], ...] = [
            ("l", "l", "l"),
            ("d", "d", "d"),
        ],
        dec_self_shift_size: Tuple[Tuple[int, ...], ...] = [(1, 1, 1), (0, 0, 0)],
        dec_cross_attn_patterns: str = None,
        dec_cross_cuboid_hw: Tuple[Tuple[int, ...], ...] = [(4, 4), (4, 4)],
        dec_cross_cuboid_strategy: Tuple[Tuple[str, ...], ...] = [
            ("l", "l", "l"),
            ("d", "l", "l"),
        ],
        dec_cross_shift_hw: Tuple[Tuple[int, ...], ...] = [(0, 0), (0, 0)],
        dec_cross_n_temporal: Tuple[int, ...] = [1, 2],
        dec_cross_last_n_frames: int = None,
        dec_use_inter_ffn: bool = True,
        dec_hierarchical_pos_embed: bool = False,
        num_global_vectors: int = 4,
        use_dec_self_global: bool = True,
        dec_self_update_global: bool = True,
        use_dec_cross_global: bool = True,
        use_global_vector_ffn: bool = True,
        use_global_self_attn: bool = False,
        separate_global_qkv: bool = False,
        global_dim_ratio: int = 1,
        self_pattern: str = "axial",
        cross_self_pattern: str = "axial",
        cross_pattern: str = "cross_1x1",
        z_init_method: str = "nearest_interp",
        initial_downsample_type: str = "conv",
        initial_downsample_activation: str = "leaky",
        initial_downsample_scale: int = 1,
        initial_downsample_conv_layers: int = 2,
        final_upsample_conv_layers: int = 2,
        initial_downsample_stack_conv_num_layers: int = 1,
        initial_downsample_stack_conv_dim_list: Tuple[int, ...] = None,
        initial_downsample_stack_conv_downscale_list: Tuple[int, ...] = [1],
        initial_downsample_stack_conv_num_conv_list: Tuple[int, ...] = [2],
        ffn_activation: str = "leaky",
        gated_ffn: bool = False,
        norm_layer: str = "layer_norm",
        padding_type: str = "ignore",
        pos_embed_type: str = "t+hw",
        checkpoint_level: bool = True,
        use_relative_pos: bool = True,
        self_attn_use_final_proj: bool = True,
        dec_use_first_self_attn: bool = False,
        attn_linear_init_mode: str = "0",
        ffn_linear_init_mode: str = "0",
        conv_init_mode: str = "0",
        down_up_linear_init_mode: str = "0",
        norm_init_mode: str = "0",
        moe_config: dict = None,
        rnc_config: dict = None,
    ):
        super().__init__()
        self.input_keys = input_keys
        self.output_keys = output_keys
        self.attn_linear_init_mode = attn_linear_init_mode
        self.ffn_linear_init_mode = ffn_linear_init_mode
        self.conv_init_mode = conv_init_mode
        self.down_up_linear_init_mode = down_up_linear_init_mode
        self.norm_init_mode = norm_init_mode
        assert len(enc_depth) == len(dec_depth)
        self.base_units = base_units
        self.num_global_vectors = num_global_vectors
        self.moe_config = moe_config
        self.rnc_config = rnc_config
        self.checkpoint_level = checkpoint_level

        num_blocks = len(enc_depth)
        if isinstance(self_pattern, str):
            enc_attn_patterns = [self_pattern] * num_blocks
        if isinstance(cross_self_pattern, str):
            dec_self_attn_patterns = [cross_self_pattern] * num_blocks
        if isinstance(cross_pattern, str):
            dec_cross_attn_patterns = [cross_pattern] * num_blocks
        if global_dim_ratio != 1:
            assert (
                separate_global_qkv is True
            ), "Setting global_dim_ratio != 1 requires separate_global_qkv == True."
        self.global_dim_ratio = global_dim_ratio
        self.z_init_method = z_init_method
        assert self.z_init_method in ["zeros", "nearest_interp", "last", "mean"]
        self.input_shape = input_shape
        self.target_shape = target_shape
        T_in, H_in, W_in, C_in = input_shape
        T_out, H_out, W_out, C_out = target_shape
        assert H_in == H_out and W_in == W_out
        if self.num_global_vectors > 0:
            init_data = paddle.zeros(
                (self.num_global_vectors, global_dim_ratio * base_units)
            )
            self.init_global_vectors = paddle.create_parameter(
                shape=init_data.shape,
                dtype=init_data.dtype,
                default_initializer=nn.initializer.Constant(0.0),
            )

            self.init_global_vectors.stop_gradient = not True
        new_input_shape = self.get_initial_encoder_final_decoder(
            initial_downsample_scale=initial_downsample_scale,
            initial_downsample_type=initial_downsample_type,
            activation=initial_downsample_activation,
            initial_downsample_conv_layers=initial_downsample_conv_layers,
            final_upsample_conv_layers=final_upsample_conv_layers,
            padding_type=padding_type,
            initial_downsample_stack_conv_num_layers=initial_downsample_stack_conv_num_layers,
            initial_downsample_stack_conv_dim_list=initial_downsample_stack_conv_dim_list,
            initial_downsample_stack_conv_downscale_list=initial_downsample_stack_conv_downscale_list,
            initial_downsample_stack_conv_num_conv_list=initial_downsample_stack_conv_num_conv_list,
        )
        T_in, H_in, W_in, _ = new_input_shape
        self.encoder = cuboid_encoder.CuboidTransformerEncoder(
            input_shape=(T_in, H_in, W_in, base_units),
            base_units=base_units,
            block_units=block_units,
            scale_alpha=scale_alpha,
            depth=enc_depth,
            downsample=downsample,
            downsample_type=downsample_type,
            block_attn_patterns=enc_attn_patterns,
            block_cuboid_size=enc_cuboid_size,
            block_strategy=enc_cuboid_strategy,
            block_shift_size=enc_shift_size,
            num_heads=num_heads,
            attn_drop=attn_drop,
            proj_drop=proj_drop,
            ffn_drop=ffn_drop,
            gated_ffn=gated_ffn,
            ffn_activation=ffn_activation,
            norm_layer=norm_layer,
            use_inter_ffn=enc_use_inter_ffn,
            padding_type=padding_type,
            use_global_vector=num_global_vectors > 0,
            use_global_vector_ffn=use_global_vector_ffn,
            use_global_self_attn=use_global_self_attn,
            separate_global_qkv=separate_global_qkv,
            global_dim_ratio=global_dim_ratio,
            checkpoint_level=checkpoint_level,
            use_relative_pos=use_relative_pos,
            self_attn_use_final_proj=self_attn_use_final_proj,
            attn_linear_init_mode=attn_linear_init_mode,
            ffn_linear_init_mode=ffn_linear_init_mode,
            conv_init_mode=conv_init_mode,
            down_linear_init_mode=down_up_linear_init_mode,
            norm_init_mode=norm_init_mode,
            moe_config=moe_config,
        )
        self.enc_pos_embed = cuboid_decoder.PosEmbed(
            embed_dim=base_units, typ=pos_embed_type, maxH=H_in, maxW=W_in, maxT=T_in
        )
        mem_shapes = self.encoder.get_mem_shapes()
        self.z_proj = nn.Linear(
            in_features=mem_shapes[-1][-1], out_features=mem_shapes[-1][-1]
        )
        self.dec_pos_embed = cuboid_decoder.PosEmbed(
            embed_dim=mem_shapes[-1][-1],
            typ=pos_embed_type,
            maxT=T_out,
            maxH=mem_shapes[-1][1],
            maxW=mem_shapes[-1][2],
        )
        self.decoder = cuboid_decoder.CuboidTransformerDecoder(
            target_temporal_length=T_out,
            mem_shapes=mem_shapes,
            cross_start=dec_cross_start,
            depth=dec_depth,
            upsample_type=upsample_type,
            block_self_attn_patterns=dec_self_attn_patterns,
            block_self_cuboid_size=dec_self_cuboid_size,
            block_self_shift_size=dec_self_shift_size,
            block_self_cuboid_strategy=dec_self_cuboid_strategy,
            block_cross_attn_patterns=dec_cross_attn_patterns,
            block_cross_cuboid_hw=dec_cross_cuboid_hw,
            block_cross_shift_hw=dec_cross_shift_hw,
            block_cross_cuboid_strategy=dec_cross_cuboid_strategy,
            block_cross_n_temporal=dec_cross_n_temporal,
            cross_last_n_frames=dec_cross_last_n_frames,
            num_heads=num_heads,
            attn_drop=attn_drop,
            proj_drop=proj_drop,
            ffn_drop=ffn_drop,
            upsample_kernel_size=upsample_kernel_size,
            ffn_activation=ffn_activation,
            gated_ffn=gated_ffn,
            norm_layer=norm_layer,
            use_inter_ffn=dec_use_inter_ffn,
            max_temporal_relative=T_in + T_out,
            padding_type=padding_type,
            hierarchical_pos_embed=dec_hierarchical_pos_embed,
            pos_embed_type=pos_embed_type,
            use_self_global=num_global_vectors > 0 and use_dec_self_global,
            self_update_global=dec_self_update_global,
            use_cross_global=num_global_vectors > 0 and use_dec_cross_global,
            use_global_vector_ffn=use_global_vector_ffn,
            use_global_self_attn=use_global_self_attn,
            separate_global_qkv=separate_global_qkv,
            global_dim_ratio=global_dim_ratio,
            checkpoint_level=checkpoint_level,
            use_relative_pos=use_relative_pos,
            self_attn_use_final_proj=self_attn_use_final_proj,
            use_first_self_attn=dec_use_first_self_attn,
            attn_linear_init_mode=attn_linear_init_mode,
            ffn_linear_init_mode=ffn_linear_init_mode,
            conv_init_mode=conv_init_mode,
            up_linear_init_mode=down_up_linear_init_mode,
            norm_init_mode=norm_init_mode,
            moe_config=moe_config,
        )

        if rnc_config["use_rnc"]:
            self.rnc_cri = extformer_moe_utils.RnCLoss(rnc_config)

        self.reset_parameters()

    def get_initial_encoder_final_decoder(
        self,
        initial_downsample_type,
        activation,
        initial_downsample_scale,
        initial_downsample_conv_layers,
        final_upsample_conv_layers,
        padding_type,
        initial_downsample_stack_conv_num_layers,
        initial_downsample_stack_conv_dim_list,
        initial_downsample_stack_conv_downscale_list,
        initial_downsample_stack_conv_num_conv_list,
    ):
        T_in, H_in, W_in, C_in = self.input_shape
        T_out, H_out, W_out, C_out = self.target_shape
        self.initial_downsample_type = initial_downsample_type
        if self.initial_downsample_type == "conv":
            if isinstance(initial_downsample_scale, int):
                initial_downsample_scale = (
                    1,
                    initial_downsample_scale,
                    initial_downsample_scale,
                )
            elif len(initial_downsample_scale) == 2:
                initial_downsample_scale = 1, *initial_downsample_scale
            elif len(initial_downsample_scale) == 3:
                initial_downsample_scale = tuple(initial_downsample_scale)
            else:
                raise NotImplementedError(
                    f"initial_downsample_scale {initial_downsample_scale} format not supported!"
                )
            self.initial_encoder = InitialEncoder(
                dim=C_in,
                out_dim=self.base_units,
                downsample_scale=initial_downsample_scale,
                num_conv_layers=initial_downsample_conv_layers,
                padding_type=padding_type,
                activation=activation,
                conv_init_mode=self.conv_init_mode,
                linear_init_mode=self.down_up_linear_init_mode,
                norm_init_mode=self.norm_init_mode,
            )

            self.final_decoder = FinalDecoder(
                dim=self.base_units,
                target_thw=(T_out, H_out, W_out),
                num_conv_layers=final_upsample_conv_layers,
                activation=activation,
                conv_init_mode=self.conv_init_mode,
                linear_init_mode=self.down_up_linear_init_mode,
                norm_init_mode=self.norm_init_mode,
            )
            new_input_shape = self.initial_encoder.patch_merge.get_out_shape(
                self.input_shape
            )
            self.dec_final_proj = nn.Linear(
                in_features=self.base_units, out_features=C_out
            )
        elif self.initial_downsample_type == "stack_conv":
            if initial_downsample_stack_conv_dim_list is None:
                initial_downsample_stack_conv_dim_list = [
                    self.base_units
                ] * initial_downsample_stack_conv_num_layers
            self.initial_encoder = InitialStackPatchMergingEncoder(
                num_merge=initial_downsample_stack_conv_num_layers,
                in_dim=C_in,
                out_dim_list=initial_downsample_stack_conv_dim_list,
                downsample_scale_list=initial_downsample_stack_conv_downscale_list,
                num_conv_per_merge_list=initial_downsample_stack_conv_num_conv_list,
                padding_type=padding_type,
                activation=activation,
                conv_init_mode=self.conv_init_mode,
                linear_init_mode=self.down_up_linear_init_mode,
                norm_init_mode=self.norm_init_mode,
            )
            initial_encoder_out_shape_list = self.initial_encoder.get_out_shape_list(
                self.target_shape
            )
            (
                dec_target_shape_list,
                dec_in_dim,
            ) = FinalStackUpsamplingDecoder.get_init_params(
                enc_input_shape=self.target_shape,
                enc_out_shape_list=initial_encoder_out_shape_list,
                large_channel=True,
            )
            self.final_decoder = FinalStackUpsamplingDecoder(
                target_shape_list=dec_target_shape_list,
                in_dim=dec_in_dim,
                num_conv_per_up_list=initial_downsample_stack_conv_num_conv_list[::-1],
                activation=activation,
                conv_init_mode=self.conv_init_mode,
                linear_init_mode=self.down_up_linear_init_mode,
                norm_init_mode=self.norm_init_mode,
            )
            self.dec_final_proj = nn.Linear(
                in_features=dec_target_shape_list[-1][-1], out_features=C_out
            )
            new_input_shape = self.initial_encoder.get_out_shape_list(self.input_shape)[
                -1
            ]
        else:
            raise NotImplementedError(f"{self.initial_downsample_type} is invalid.")
        self.input_shape_after_initial_downsample = new_input_shape
        T_in, H_in, W_in, _ = new_input_shape
        return new_input_shape

    def reset_parameters(self):
        if self.num_global_vectors > 0:
            self.init_global_vectors = initializer.trunc_normal_(
                self.init_global_vectors, std=0.02
            )
        if hasattr(self.initial_encoder, "reset_parameters"):
            self.initial_encoder.reset_parameters()
        else:
            cuboid_utils.apply_initialization(
                self.initial_encoder,
                conv_mode=self.conv_init_mode,
                linear_mode=self.down_up_linear_init_mode,
                norm_mode=self.norm_init_mode,
            )
        if hasattr(self.final_decoder, "reset_parameters"):
            self.final_decoder.reset_parameters()
        else:
            cuboid_utils.apply_initialization(
                self.final_decoder,
                conv_mode=self.conv_init_mode,
                linear_mode=self.down_up_linear_init_mode,
                norm_mode=self.norm_init_mode,
            )
        cuboid_utils.apply_initialization(
            self.dec_final_proj, linear_mode=self.down_up_linear_init_mode
        )
        self.encoder.reset_parameters()
        self.enc_pos_embed.reset_parameters()
        self.decoder.reset_parameters()
        self.dec_pos_embed.reset_parameters()
        cuboid_utils.apply_initialization(self.z_proj, linear_mode="0")

    def get_initial_z(self, final_mem, T_out):
        B = final_mem.shape[0]
        if self.z_init_method == "zeros":
            z_shape = list((1, T_out)) + final_mem.shape[2:]
            initial_z = paddle.zeros(shape=z_shape, dtype=final_mem.dtype)
            initial_z = self.z_proj(self.dec_pos_embed(initial_z)).expand(
                shape=[B, -1, -1, -1, -1]
            )
        elif self.z_init_method == "nearest_interp":
            initial_z = nn.functional.interpolate(
                x=final_mem.transpose(perm=[0, 4, 1, 2, 3]),
                size=(T_out, final_mem.shape[2], final_mem.shape[3]),
            ).transpose(perm=[0, 2, 3, 4, 1])
            initial_z = self.z_proj(initial_z)
        elif self.z_init_method == "last":
            initial_z = paddle.broadcast_to(
                x=final_mem[:, -1:, :, :, :], shape=(B, T_out) + final_mem.shape[2:]
            )
            initial_z = self.z_proj(initial_z)
        elif self.z_init_method == "mean":
            initial_z = paddle.broadcast_to(
                x=final_mem.mean(axis=1, keepdims=True),
                shape=(B, T_out) + final_mem.shape[2:],
            )
            initial_z = self.z_proj(initial_z)
        else:
            raise NotImplementedError
        return initial_z

    def forward(self, x: paddle.Tensor, verbose: bool = False) -> paddle.Tensor:
        """
        Args:
            x (paddle.Tensor): Tensor with shape (B, T, H, W, C).
            verbose (bool): if True, print intermediate shapes.

        Returns:
            out (paddle.Tensor): The output Shape (B, T_out, H, W, C_out)
        """

        labels = x["sst_target"]
        x = self.concat_to_tensor(x, self.input_keys)
        flag_ndim = x.ndim
        if flag_ndim == 6:
            x = x.reshape([-1, *x.shape[2:]])
        B, _, _, _, _ = x.shape

        T_out = self.target_shape[0]
        x = self.initial_encoder(x)
        x = self.enc_pos_embed(x)

        if self.num_global_vectors > 0:
            init_global_vectors = self.init_global_vectors.expand(
                shape=[
                    B,
                    self.num_global_vectors,
                    self.global_dim_ratio * self.base_units,
                ]
            )
            mem_l, mem_global_vector_l = self.encoder(x, init_global_vectors)
        else:
            mem_l = self.encoder(x)

        if verbose:
            for i, mem in enumerate(mem_l):
                print(f"mem[{i}].shape = {mem.shape}")
        initial_z = self.get_initial_z(final_mem=mem_l[-1], T_out=T_out)

        if self.num_global_vectors > 0:
            dec_out = self.decoder(initial_z, mem_l, mem_global_vector_l)
        else:
            dec_out = self.decoder(initial_z, mem_l)

        dec_out = self.final_decoder(dec_out)
        out = self.dec_final_proj(dec_out)

        if flag_ndim == 6:
            out = out.reshape([-1, *out.shape])

        out_dict = {key: out for key in self.output_keys[:2]}

        # moe loss
        if self.training:
            aux_losses = extformer_moe_utils.aggregate_aux_losses(self)
            if len(aux_losses) > 0:
                aux_loss = paddle.concat(aux_losses).mean()
            else:
                aux_loss = None
        else:
            aux_loss = None
        assert "aux_loss" in self.output_keys
        out_dict["aux_loss"] = aux_loss

        # rnc
        if self.training and self.rnc_config["use_rnc"]:
            rank_loss = self.rnc_cri(dec_out, labels)
            rank_loss = rank_loss.unsqueeze(0)
        else:
            rank_loss = None
        assert "rank_loss" in self.output_keys
        out_dict["rank_loss"] = rank_loss

        return out_dict

forward(x, verbose=False)

Parameters:

Name	Type	Description	Default
x	Tensor	Tensor with shape (B, T, H, W, C).	required
verbose	bool	if True, print intermediate shapes.	False

Returns:

Name	Type	Description
out	Tensor	The output Shape (B, T_out, H, W, C_out)

Source code in ppsci/arch/extformer_moe_cuboid.py

def forward(self, x: paddle.Tensor, verbose: bool = False) -> paddle.Tensor:
    """
    Args:
        x (paddle.Tensor): Tensor with shape (B, T, H, W, C).
        verbose (bool): if True, print intermediate shapes.

    Returns:
        out (paddle.Tensor): The output Shape (B, T_out, H, W, C_out)
    """

    labels = x["sst_target"]
    x = self.concat_to_tensor(x, self.input_keys)
    flag_ndim = x.ndim
    if flag_ndim == 6:
        x = x.reshape([-1, *x.shape[2:]])
    B, _, _, _, _ = x.shape

    T_out = self.target_shape[0]
    x = self.initial_encoder(x)
    x = self.enc_pos_embed(x)

    if self.num_global_vectors > 0:
        init_global_vectors = self.init_global_vectors.expand(
            shape=[
                B,
                self.num_global_vectors,
                self.global_dim_ratio * self.base_units,
            ]
        )
        mem_l, mem_global_vector_l = self.encoder(x, init_global_vectors)
    else:
        mem_l = self.encoder(x)

    if verbose:
        for i, mem in enumerate(mem_l):
            print(f"mem[{i}].shape = {mem.shape}")
    initial_z = self.get_initial_z(final_mem=mem_l[-1], T_out=T_out)

    if self.num_global_vectors > 0:
        dec_out = self.decoder(initial_z, mem_l, mem_global_vector_l)
    else:
        dec_out = self.decoder(initial_z, mem_l)

    dec_out = self.final_decoder(dec_out)
    out = self.dec_final_proj(dec_out)

    if flag_ndim == 6:
        out = out.reshape([-1, *out.shape])

    out_dict = {key: out for key in self.output_keys[:2]}

    # moe loss
    if self.training:
        aux_losses = extformer_moe_utils.aggregate_aux_losses(self)
        if len(aux_losses) > 0:
            aux_loss = paddle.concat(aux_losses).mean()
        else:
            aux_loss = None
    else:
        aux_loss = None
    assert "aux_loss" in self.output_keys
    out_dict["aux_loss"] = aux_loss

    # rnc
    if self.training and self.rnc_config["use_rnc"]:
        rank_loss = self.rnc_cri(dec_out, labels)
        rank_loss = rank_loss.unsqueeze(0)
    else:
        rank_loss = None
    assert "rank_loss" in self.output_keys
    out_dict["rank_loss"] = rank_loss

    return out_dict

FNO1d

Bases: Layer

The overall network. It contains 4 layers of the Fourier layer. 1. Lift the input to the desire channel dimension by self.fc0 . 2. 4 layers of the integral operators u' = (W + K)(u). W defined by self.w; K defined by self.conv . 3. Project from the channel space to the output space by self.fc1 and self.fc2 .

Parameters:

Name	Type	Description	Default
input_key	Tuple[str, ...]	Key to get the input tensor from the dict. Defaults to ("intput",).	('input',)
output_key	Tuple[str, ...]	Key to save the output tensor into the dict. Defaults to ("output",).	('output',)
modes	(int, optional)	Number of Fourier modes to compute, it should be the same as that in fft part of the code below. Defaults to 64.	64
width	(int, optional)	Number of channels in each Fourier layer. Defaults to 64.	64
padding	(int, optional)	How many zeros to pad to the input Tensor. Defaults to 100.	100
input_channel	(int, optional)	Number of channels of the input tensor. Defaults to 2.	2
output_np	(int, optional)	Number of points to sample the solution. Defaults to 2001.	2001

Examples:

>>> model = ppsci.arch.FNO1d()
>>> input_data = paddle.randn([100, 2001, 2])
>>> input_dict = {"input": input_data}
>>> out_dict = model(input_dict)
>>> for k, v in out_dict.items():
...     print(k, v.shape)
output [100, 1]

Source code in ppsci/arch/geofno.py

class FNO1d(nn.Layer):
    """The overall network. It contains 4 layers of the Fourier layer.
    1. Lift the input to the desire channel dimension by self.fc0 .
    2. 4 layers of the integral operators u' = (W + K)(u).
         W defined by self.w; K defined by self.conv .
    3. Project from the channel space to the output space by self.fc1 and self.fc2 .

    Args:
        input_key (Tuple[str, ...], optional): Key to get the input tensor from the dict. Defaults to ("intput",).
        output_key (Tuple[str, ...], optional): Key to save the output tensor into the dict. Defaults to ("output",).
        modes (int, optional, optional): Number of Fourier modes to compute, it should be the same as
            that in fft part of the code below. Defaults to 64.
        width (int, optional, optional): Number of channels in each Fourier layer. Defaults to 64.
        padding (int, optional, optional): How many zeros to pad to the input Tensor. Defaults to 100.
        input_channel (int, optional, optional): Number of channels of the input tensor. Defaults to 2.
        output_np (int, optional, optional): Number of points to sample the solution. Defaults to 2001.

    Examples:
        >>> model = ppsci.arch.FNO1d()
        >>> input_data = paddle.randn([100, 2001, 2])
        >>> input_dict = {"input": input_data}
        >>> out_dict = model(input_dict)
        >>> for k, v in out_dict.items():
        ...     print(k, v.shape)
        output [100, 1]
    """

    def __init__(
        self,
        input_key=("input",),
        output_key=("output",),
        modes=64,
        width=64,
        padding=100,
        input_channel=2,
        output_np=2001,
    ):
        super().__init__()
        self.input_keys = input_key
        self.output_keys = output_key

        self.output_np = output_np
        self.modes1 = modes
        self.width = width
        self.padding = padding
        self.fc0 = nn.Linear(input_channel, self.width)

        self.conv0 = SpectralConv1d(self.width, self.width, self.modes1)
        self.conv1 = SpectralConv1d(self.width, self.width, self.modes1)
        self.conv2 = SpectralConv1d(self.width, self.width, self.modes1)
        self.conv3 = SpectralConv1d(self.width, self.width, self.modes1)
        self.conv4 = SpectralConv1d(self.width, self.width, self.modes1)

        self.w0 = nn.Conv1D(self.width, self.width, 1)
        self.w1 = nn.Conv1D(self.width, self.width, 1)
        self.w2 = nn.Conv1D(self.width, self.width, 1)
        self.w3 = nn.Conv1D(self.width, self.width, 1)

        self.fc1 = nn.Linear(self.width, 128)
        self.fc2 = nn.Linear(128, 1)

    def _functional_pad(self, x, pad, mode="constant", value=0.0, data_format="NCL"):
        if len(x.shape) * 2 == len(pad) and mode == "constant":
            pad = (
                paddle.to_tensor(pad, dtype="float32")
                .reshape((-1, 2))
                .flip([0])
                .flatten()
                .tolist()
            )
        return F.pad(x, pad, mode, value, data_format)

    def forward(self, x):
        x = x[self.input_keys[0]]
        # Dict
        x = self.fc0(x)
        x = paddle.transpose(x, perm=[0, 2, 1])
        # pad the domain if input is non-periodic
        x = self._functional_pad(x, [0, self.padding])

        x1 = self.conv0(x)
        x2 = self.w0(x)
        x = x1 + x2
        x = F.gelu(x=x, approximate=False)

        x1 = self.conv1(x)
        x2 = self.w1(x)
        x = x1 + x2
        x = F.gelu(x, approximate=False)

        x1 = self.conv2(x)
        x2 = self.w2(x)
        x = x1 + x2
        x = F.gelu(x, approximate=False)

        x1 = self.conv3(x)
        x2 = self.w3(x)
        x = x1 + x2
        x = F.gelu(x, approximate=False)

        x = x[..., : -self.padding]
        x1 = self.conv4(x, self.output_np)
        x2 = F.interpolate(x, size=[self.output_np], mode="linear", align_corners=True)
        x = x1 + x2
        # Change the x-dimension to (batch, channel, 2001)
        x = x.transpose(perm=[0, 2, 1])
        x = self.fc1(x)
        x = F.gelu(x, approximate=False)
        x = self.fc2(x)

        return {self.output_keys[0]: x}

Generator

Bases: Arch

Generator Net of GAN. Attention, the net using a kind of variant of ResBlock which is unique to "tempoGAN" example but not an open source network.

Parameters:

Name	Type	Description	Default
input_keys	Tuple[str, ...]	Name of input keys, such as ("input1", "input2").	required
output_keys	Tuple[str, ...]	Name of output keys, such as ("output1", "output2").	required
in_channel	int	Number of input channels of the first conv layer.	required
out_channels_tuple	Tuple[Tuple[int, ...], ...]	Number of output channels of all conv layers, such as [[out_res0_conv0, out_res0_conv1], [out_res1_conv0, out_res1_conv1]]	required
kernel_sizes_tuple	Tuple[Tuple[int, ...], ...]	Number of kernel_size of all conv layers, such as [[kernel_size_res0_conv0, kernel_size_res0_conv1], [kernel_size_res1_conv0, kernel_size_res1_conv1]]	required
strides_tuple	Tuple[Tuple[int, ...], ...]	Number of stride of all conv layers, such as [[stride_res0_conv0, stride_res0_conv1], [stride_res1_conv0, stride_res1_conv1]]	required
use_bns_tuple	Tuple[Tuple[bool, ...], ...]	Whether to use the batch_norm layer after each conv layer.	required
acts_tuple	Tuple[Tuple[str, ...], ...]	Whether to use the activation layer after each conv layer. If so, witch activation to use, such as [[act_res0_conv0, act_res0_conv1], [act_res1_conv0, act_res1_conv1]]	required

Examples:

>>> import ppsci
>>> in_channel = 1
>>> rb_channel0 = (2, 8, 8)
>>> rb_channel1 = (128, 128, 128)
>>> rb_channel2 = (32, 8, 8)
>>> rb_channel3 = (2, 1, 1)
>>> out_channels_tuple = (rb_channel0, rb_channel1, rb_channel2, rb_channel3)
>>> kernel_sizes_tuple = (((5, 5), ) * 2 + ((1, 1), ), ) * 4
>>> strides_tuple = ((1, 1, 1), ) * 4
>>> use_bns_tuple = ((True, True, True), ) * 3 + ((False, False, False), )
>>> acts_tuple = (("relu", None, None), ) * 4
>>> model = ppsci.arch.Generator(("in",), ("out",), in_channel, out_channels_tuple, kernel_sizes_tuple, strides_tuple, use_bns_tuple, acts_tuple)
>>> batch_size = 4
>>> height = 64
>>> width = 64
>>> input_data = paddle.randn([batch_size, in_channel, height, width])
>>> input_dict = {'in': input_data}
>>> output_data = model(input_dict)
>>> print(output_data['out'].shape)
[4, 1, 64, 64]

Source code in ppsci/arch/gan.py

class Generator(base.Arch):
    """Generator Net of GAN. Attention, the net using a kind of variant of ResBlock which is
        unique to "tempoGAN" example but not an open source network.

    Args:
        input_keys (Tuple[str, ...]): Name of input keys, such as ("input1", "input2").
        output_keys (Tuple[str, ...]): Name of output keys, such as ("output1", "output2").
        in_channel (int): Number of input channels of the first conv layer.
        out_channels_tuple (Tuple[Tuple[int, ...], ...]): Number of output channels of all conv layers,
            such as [[out_res0_conv0, out_res0_conv1], [out_res1_conv0, out_res1_conv1]]
        kernel_sizes_tuple (Tuple[Tuple[int, ...], ...]): Number of kernel_size of all conv layers,
            such as [[kernel_size_res0_conv0, kernel_size_res0_conv1], [kernel_size_res1_conv0, kernel_size_res1_conv1]]
        strides_tuple (Tuple[Tuple[int, ...], ...]): Number of stride of all conv layers,
            such as [[stride_res0_conv0, stride_res0_conv1], [stride_res1_conv0, stride_res1_conv1]]
        use_bns_tuple (Tuple[Tuple[bool, ...], ...]): Whether to use the batch_norm layer after each conv layer.
        acts_tuple (Tuple[Tuple[str, ...], ...]): Whether to use the activation layer after each conv layer. If so, witch activation to use,
            such as [[act_res0_conv0, act_res0_conv1], [act_res1_conv0, act_res1_conv1]]

    Examples:
        >>> import ppsci
        >>> in_channel = 1
        >>> rb_channel0 = (2, 8, 8)
        >>> rb_channel1 = (128, 128, 128)
        >>> rb_channel2 = (32, 8, 8)
        >>> rb_channel3 = (2, 1, 1)
        >>> out_channels_tuple = (rb_channel0, rb_channel1, rb_channel2, rb_channel3)
        >>> kernel_sizes_tuple = (((5, 5), ) * 2 + ((1, 1), ), ) * 4
        >>> strides_tuple = ((1, 1, 1), ) * 4
        >>> use_bns_tuple = ((True, True, True), ) * 3 + ((False, False, False), )
        >>> acts_tuple = (("relu", None, None), ) * 4
        >>> model = ppsci.arch.Generator(("in",), ("out",), in_channel, out_channels_tuple, kernel_sizes_tuple, strides_tuple, use_bns_tuple, acts_tuple)
        >>> batch_size = 4
        >>> height = 64
        >>> width = 64
        >>> input_data = paddle.randn([batch_size, in_channel, height, width])
        >>> input_dict = {'in': input_data}
        >>> output_data = model(input_dict)
        >>> print(output_data['out'].shape)
        [4, 1, 64, 64]
    """

    def __init__(
        self,
        input_keys: Tuple[str, ...],
        output_keys: Tuple[str, ...],
        in_channel: int,
        out_channels_tuple: Tuple[Tuple[int, ...], ...],
        kernel_sizes_tuple: Tuple[Tuple[int, ...], ...],
        strides_tuple: Tuple[Tuple[int, ...], ...],
        use_bns_tuple: Tuple[Tuple[bool, ...], ...],
        acts_tuple: Tuple[Tuple[str, ...], ...],
    ):
        super().__init__()
        self.input_keys = input_keys
        self.output_keys = output_keys
        self.in_channel = in_channel
        self.out_channels_tuple = out_channels_tuple
        self.kernel_sizes_tuple = kernel_sizes_tuple
        self.strides_tuple = strides_tuple
        self.use_bns_tuple = use_bns_tuple
        self.acts_tuple = acts_tuple

        self.init_blocks()

    def init_blocks(self):
        blocks_list = []
        for i in range(len(self.out_channels_tuple)):
            in_channel = (
                self.in_channel if i == 0 else self.out_channels_tuple[i - 1][-1]
            )
            blocks_list.append(
                VariantResBlock(
                    in_channel=in_channel,
                    out_channels=self.out_channels_tuple[i],
                    kernel_sizes=self.kernel_sizes_tuple[i],
                    strides=self.strides_tuple[i],
                    use_bns=self.use_bns_tuple[i],
                    acts=self.acts_tuple[i],
                    mean=0.0,
                    std=0.04,
                    value=0.1,
                )
            )
        self.blocks = nn.LayerList(blocks_list)

    def forward_tensor(self, x):
        y = x
        for block in self.blocks:
            y = block(y)
        return y

    def forward(self, x):
        if self._input_transform is not None:
            x = self._input_transform(x)

        y = self.concat_to_tensor(x, self.input_keys, axis=-1)
        y = self.forward_tensor(y)
        y = self.split_to_dict(y, self.output_keys, axis=-1)

        if self._output_transform is not None:
            y = self._output_transform(x, y)
        return y

HEDeepONets

Bases: Arch

Physical information deep operator networks.

Parameters:

Name	Type	Description	Default
heat_input_keys	Tuple[str, ...]	Name of input data for heat boundary.	required
cold_input_keys	Tuple[str, ...]	Name of input data for cold boundary.	required
trunk_input_keys	Tuple[str, ...]	Name of input data for trunk net.	required
output_keys	Tuple[str, ...]	Output name of predicted temperature.	required
heat_num_loc	int	Number of sampled input data for heat boundary.	required
cold_num_loc	int	Number of sampled input data for cold boundary.	required
num_features	int	Number of features extracted from heat boundary, same for cold boundary and trunk net.	required
branch_num_layers	int	Number of hidden layers of branch net.	required
trunk_num_layers	int	Number of hidden layers of trunk net.	required
branch_hidden_size	Union[int, Tuple[int, ...]]	Number of hidden size of branch net. An integer for all layers, or list of integer specify each layer's size.	required
trunk_hidden_size	Union[int, Tuple[int, ...]]	Number of hidden size of trunk net. An integer for all layers, or list of integer specify each layer's size.	required
branch_skip_connection	bool	Whether to use skip connection for branch net. Defaults to False.	False
trunk_skip_connection	bool	Whether to use skip connection for trunk net. Defaults to False.	False
branch_activation	str	Name of activation function for branch net. Defaults to "tanh".	'tanh'
trunk_activation	str	Name of activation function for trunk net. Defaults to "tanh".	'tanh'
branch_weight_norm	bool	Whether to apply weight norm on parameter(s) for branch net. Defaults to False.	False
trunk_weight_norm	bool	Whether to apply weight norm on parameter(s) for trunk net. Defaults to False.	False
use_bias	bool	Whether to add bias on predicted G(u)(y). Defaults to True.	True

Examples:

>>> import ppsci
>>> model = ppsci.arch.HEDeepONets(
...     ('qm_h',),
...     ('qm_c',),
...     ("x",'t'),
...     ("T_h",'T_c','T_w'),
...     1,
...     1,
...     100,
...     9,
...     6,
...     256,
...     128,
...     branch_activation="swish",
...     trunk_activation="swish",
...     use_bias=True,
... )

Source code in ppsci/arch/he_deeponets.py

class HEDeepONets(base.Arch):
    """Physical information deep operator networks.

    Args:
        heat_input_keys (Tuple[str, ...]): Name of input data for heat boundary.
        cold_input_keys (Tuple[str, ...]): Name of input data for cold boundary.
        trunk_input_keys (Tuple[str, ...]): Name of input data for trunk net.
        output_keys (Tuple[str, ...]): Output name of predicted temperature.
        heat_num_loc (int): Number of sampled input data for heat boundary.
        cold_num_loc (int): Number of sampled input data for cold boundary.
        num_features (int): Number of features extracted from heat boundary, same for cold boundary and trunk net.
        branch_num_layers (int): Number of hidden layers of branch net.
        trunk_num_layers (int): Number of hidden layers of trunk net.
        branch_hidden_size (Union[int, Tuple[int, ...]]): Number of hidden size of branch net.
            An integer for all layers, or list of integer specify each layer's size.
        trunk_hidden_size (Union[int, Tuple[int, ...]]): Number of hidden size of trunk net.
            An integer for all layers, or list of integer specify each layer's size.
        branch_skip_connection (bool, optional): Whether to use skip connection for branch net. Defaults to False.
        trunk_skip_connection (bool, optional): Whether to use skip connection for trunk net. Defaults to False.
        branch_activation (str, optional): Name of activation function for branch net. Defaults to "tanh".
        trunk_activation (str, optional): Name of activation function for trunk net. Defaults to "tanh".
        branch_weight_norm (bool, optional): Whether to apply weight norm on parameter(s) for branch net. Defaults to False.
        trunk_weight_norm (bool, optional): Whether to apply weight norm on parameter(s) for trunk net. Defaults to False.
        use_bias (bool, optional): Whether to add bias on predicted G(u)(y). Defaults to True.

    Examples:
        >>> import ppsci
        >>> model = ppsci.arch.HEDeepONets(
        ...     ('qm_h',),
        ...     ('qm_c',),
        ...     ("x",'t'),
        ...     ("T_h",'T_c','T_w'),
        ...     1,
        ...     1,
        ...     100,
        ...     9,
        ...     6,
        ...     256,
        ...     128,
        ...     branch_activation="swish",
        ...     trunk_activation="swish",
        ...     use_bias=True,
        ... )
    """

    def __init__(
        self,
        heat_input_keys: Tuple[str, ...],
        cold_input_keys: Tuple[str, ...],
        trunk_input_keys: Tuple[str, ...],
        output_keys: Tuple[str, ...],
        heat_num_loc: int,
        cold_num_loc: int,
        num_features: int,
        branch_num_layers: int,
        trunk_num_layers: int,
        branch_hidden_size: Union[int, Tuple[int, ...]],
        trunk_hidden_size: Union[int, Tuple[int, ...]],
        branch_skip_connection: bool = False,
        trunk_skip_connection: bool = False,
        branch_activation: str = "tanh",
        trunk_activation: str = "tanh",
        branch_weight_norm: bool = False,
        trunk_weight_norm: bool = False,
        use_bias: bool = True,
    ):
        super().__init__()
        self.trunk_input_keys = trunk_input_keys
        self.heat_input_keys = heat_input_keys
        self.cold_input_keys = cold_input_keys
        self.input_keys = (
            self.trunk_input_keys + self.heat_input_keys + self.cold_input_keys
        )
        self.output_keys = output_keys
        self.num_features = num_features

        self.heat_net = mlp.MLP(
            self.heat_input_keys,
            ("h",),
            branch_num_layers,
            branch_hidden_size,
            branch_activation,
            branch_skip_connection,
            branch_weight_norm,
            input_dim=heat_num_loc,
            output_dim=num_features * len(self.output_keys),
        )

        self.cold_net = mlp.MLP(
            self.cold_input_keys,
            ("c",),
            branch_num_layers,
            branch_hidden_size,
            branch_activation,
            branch_skip_connection,
            branch_weight_norm,
            input_dim=cold_num_loc,
            output_dim=num_features * len(self.output_keys),
        )

        self.trunk_net = mlp.MLP(
            self.trunk_input_keys,
            ("t",),
            trunk_num_layers,
            trunk_hidden_size,
            trunk_activation,
            trunk_skip_connection,
            trunk_weight_norm,
            input_dim=len(self.trunk_input_keys),
            output_dim=num_features * len(self.output_keys),
        )
        self.trunk_act = act_mod.get_activation(trunk_activation)
        self.heat_act = act_mod.get_activation(branch_activation)
        self.cold_act = act_mod.get_activation(branch_activation)

        self.use_bias = use_bias
        if use_bias:
            # register bias to parameter for updating in optimizer and storage
            self.b = self.create_parameter(
                shape=(len(self.output_keys),),
                attr=nn.initializer.Constant(0.0),
            )

    def forward(self, x):
        if self._input_transform is not None:
            x = self._input_transform(x)

        # Branch net to encode the input function
        heat_features = self.heat_net(x)[self.heat_net.output_keys[0]]
        cold_features = self.cold_net(x)[self.cold_net.output_keys[0]]
        # Trunk net to encode the domain of the output function
        y_features = self.trunk_net(x)[self.trunk_net.output_keys[0]]
        y_features = self.trunk_act(y_features)
        # Dot product
        G_u_h = paddle.sum(
            heat_features[:, : self.num_features]
            * y_features[:, : self.num_features]
            * cold_features[:, : self.num_features],
            axis=1,
            keepdim=True,
        )
        G_u_c = paddle.sum(
            heat_features[:, self.num_features : 2 * self.num_features]
            * y_features[:, self.num_features : 2 * self.num_features]
            * cold_features[:, self.num_features : 2 * self.num_features],
            axis=1,
            keepdim=True,
        )
        G_u_w = paddle.sum(
            heat_features[:, 2 * self.num_features :]
            * y_features[:, 2 * self.num_features :]
            * cold_features[:, 2 * self.num_features :],
            axis=1,
            keepdim=True,
        )
        # Add bias
        if self.use_bias:
            G_u_h += self.b[0]
            G_u_c += self.b[1]
            G_u_w += self.b[2]

        result_dict = {
            self.output_keys[0]: G_u_h,
            self.output_keys[1]: G_u_c,
            self.output_keys[2]: G_u_w,
        }
        if self._output_transform is not None:
            result_dict = self._output_transform(x, result_dict)

        return result_dict

IFMMLP

Bases: Arch

Understanding the limitations of deep models for molecular property prediction: Insights and solutions. [Xia, Jun, et al. Advances in Neural Information Processing Systems 36 (2023): 64774-64792.]https://openreview.net/forum?id=NLFqlDeuzt)

Code reference: https://github.com/junxia97/IFM

Parameters:

Name	Type	Description	Default
input_keys	Tuple[str, ...]	Name of input keys, such as ("input", ).	required
output_keys	Tuple[str, ...]	Name of output keys, such as ("pred", ).	required
hidden_units	List[int]	Units num in hidden layers.	required
embed_name	str	Embed name used in arch, such as "IMF", "None".	required
inputs	int	Input dim.	required
outputs	int	Output dim.	required
d_out	int	Embedding output dim for some architecture.	required
sigma	float	Hyper parameter for some architecture.	required
dp_ratio	float	Dropout ratio.	required
reg	bool	Regularization flag.	required
first_omega_0	float	Frequency factor used in first layer.	required
hidden_omega_0	float	Frequency factor used in hidden layer.	required

Source code in ppsci/arch/ifm_mlp.py

class IFMMLP(base.Arch):
    """Understanding the limitations of deep models for molecular property prediction: Insights and solutions.
    [Xia, Jun, et al. Advances in Neural Information Processing Systems 36 (2023): 64774-64792.]https://openreview.net/forum?id=NLFqlDeuzt)

    Code reference: https://github.com/junxia97/IFM

    Args:
        input_keys (Tuple[str, ...]): Name of input keys, such as ("input", ).
        output_keys (Tuple[str, ...]): Name of output keys, such as ("pred", ).
        hidden_units (List[int]): Units num in hidden layers.
        embed_name (str): Embed name used in arch, such as "IMF", "None".
        inputs (int): Input dim.
        outputs (int): Output dim.
        d_out (int): Embedding output dim for some architecture.
        sigma (float): Hyper parameter for some architecture.
        dp_ratio (float): Dropout ratio.
        reg (bool): Regularization flag.
        first_omega_0 (float): Frequency factor used in first layer.
        hidden_omega_0 (float): Frequency factor used in hidden layer.
    """

    def __init__(
        self,
        input_keys: Tuple[str, ...],
        output_keys: Tuple[str, ...],
        hidden_units: List[int],
        embed_name: str,
        inputs: int,
        outputs: int,
        d_out: int,
        sigma: float,
        dp_ratio: float,
        reg: bool,
        first_omega_0: float,
        hidden_omega_0: float,
    ):
        super().__init__()
        self.input_keys = input_keys
        self.output_keys = output_keys

        # initialization
        if embed_name == "None":
            my_model = MyDNN(
                inputs=inputs,
                hidden_units=hidden_units,
                dp_ratio=dp_ratio,
                outputs=outputs,
                reg=reg,
            )
        elif embed_name == "LE":
            my_model = LE_DNN(
                inputs=inputs,
                hidden_units=hidden_units,
                d_out=d_out + 1,
                dp_ratio=dp_ratio,
                outputs=outputs,
                reg=reg,
            )
        elif embed_name == "LSIM":
            my_model = LSIM_DNN(
                inputs=inputs,
                hidden_units=hidden_units,
                d_out=d_out + 1,
                sigma=sigma,
                dp_ratio=dp_ratio,
                outputs=outputs,
                reg=reg,
            )
        elif embed_name == "IFM":
            my_model = IFM_DNN(
                inputs=inputs,
                hidden_units=hidden_units,
                outputs=outputs,
                dp_ratio=dp_ratio,
                first_omega_0=first_omega_0,
                hidden_omega_0=hidden_omega_0,
                reg=reg,
            )
        elif embed_name == "GM":
            my_model = GM_DNN(
                inputs=inputs,
                hidden_units=hidden_units,
                d_out=d_out + 1,
                sigma=sigma + 1,
                dp_ratio=dp_ratio,
                outputs=outputs,
                reg=reg,
            )
        elif embed_name == "SIM":
            my_model = SIM_DNN(
                inputs=inputs,
                hidden_units=hidden_units,
                d_out=d_out + 1,
                sigma=sigma + 1,
                dp_ratio=dp_ratio,
                outputs=outputs,
                reg=reg,
            )
        else:
            raise ValueError("Invalid Embedding Name")

        self.model = my_model

    def forward(self, x):
        Xs = x[self.input_keys[0]]
        ret = self.model(Xs)
        return {self.output_keys[0]: ret}

LatentNO

Bases: Arch

Source code in ppsci/arch/latent_no.py

class LatentNO(base.Arch):
    def __init__(
        self,
        n_block: int,
        n_mode: int,
        n_dim: int,
        n_head: int,
        n_layer: int,
        trunk_dim: int,
        branch_dim: int,
        out_dim: int,
        input_keys: Tuple[str, ...],
        output_keys: Tuple[str, ...],
    ):
        """
        Latent Neural Operator (LatentNO).

        Args:
            n_block (int): Number of attention blocks.
            n_mode (int): Number of latent modes.
            n_dim (int): Hidden dimension size.
            n_head (int): Number of attention heads.
            n_layer (int): Number of layers in MLP.
            trunk_dim (int): Dimension of trunk input.
            branch_dim (int): Dimension of branch input.
            out_dim (int): Dimension of output.
            input_keys (Tuple[str, ...]): Name of input keys.
            output_keys (Tuple[str, ...]): Name of output keys.
        """
        super().__init__()
        self.input_keys = input_keys
        self.output_keys = output_keys
        self.trunk_dim = trunk_dim
        self.trunk_mlp = LatentMLP(trunk_dim, n_dim, n_dim, n_layer)
        self.branch_mlp = LatentMLP(branch_dim, n_dim, n_dim, n_layer)
        self.mode_mlp = LatentMLP(n_dim, n_dim, n_mode, n_layer)
        self.out_mlp = LatentMLP(n_dim, n_dim, out_dim, n_layer)

        self.attn_blocks = paddle.nn.Sequential(
            *[AttentionBlock(n_mode, n_dim, n_head) for _ in range(n_block)]
        )
        self.apply(self._init_weights)

    def _init_weights(self, module):
        if isinstance(module, paddle.nn.Linear):
            initializer.linear_init_(module)
        elif isinstance(module, paddle.nn.Conv2D):
            initializer.conv_init_(module)
        elif isinstance(module, paddle.nn.LayerNorm):
            initializer.ones_(module.weight)
            initializer.zeros_(module.bias)

    def forward(self, inputs: dict[str, paddle.Tensor]) -> dict[str, paddle.Tensor]:
        """
        Forward pass of LatentNO.

        Args:
            inputs (dict[str, paddle.Tensor]):
                Dictionary with keys:
                    - "x": Trunk input tensor of shape (B, N, trunk_dim).
                    - "y1": Branch input tensor of shape (B, N, branch_dim).

        Returns:
            dict[str, paddle.Tensor]: Dictionary containing:
                - "y2": Output tensor of shape (B, N, out_dim).
        """
        x = inputs[self.input_keys[0]]  # trunk input
        y = inputs[self.input_keys[1]]  # branch input

        x = self.trunk_mlp(x)
        y = self.branch_mlp(y)

        score = self.mode_mlp(x)
        score_encode = paddle.nn.functional.softmax(score, axis=1)
        score_decode = paddle.nn.functional.softmax(score, axis=-1)

        z = paddle.matmul(paddle.transpose(score_encode, perm=[0, 2, 1]), y)
        for block in self.attn_blocks:
            z = block(z)

        r = paddle.matmul(score_decode, z)
        r = self.out_mlp(r)

        return {self.output_keys[0]: r}

__init__(n_block, n_mode, n_dim, n_head, n_layer, trunk_dim, branch_dim, out_dim, input_keys, output_keys)

Latent Neural Operator (LatentNO).

Parameters:

Name	Type	Description	Default
n_block	int	Number of attention blocks.	required
n_mode	int	Number of latent modes.	required
n_dim	int	Hidden dimension size.	required
n_head	int	Number of attention heads.	required
n_layer	int	Number of layers in MLP.	required
trunk_dim	int	Dimension of trunk input.	required
branch_dim	int	Dimension of branch input.	required
out_dim	int	Dimension of output.	required
input_keys	Tuple[str, ...]	Name of input keys.	required
output_keys	Tuple[str, ...]	Name of output keys.	required

Source code in ppsci/arch/latent_no.py

def __init__(
    self,
    n_block: int,
    n_mode: int,
    n_dim: int,
    n_head: int,
    n_layer: int,
    trunk_dim: int,
    branch_dim: int,
    out_dim: int,
    input_keys: Tuple[str, ...],
    output_keys: Tuple[str, ...],
):
    """
    Latent Neural Operator (LatentNO).

    Args:
        n_block (int): Number of attention blocks.
        n_mode (int): Number of latent modes.
        n_dim (int): Hidden dimension size.
        n_head (int): Number of attention heads.
        n_layer (int): Number of layers in MLP.
        trunk_dim (int): Dimension of trunk input.
        branch_dim (int): Dimension of branch input.
        out_dim (int): Dimension of output.
        input_keys (Tuple[str, ...]): Name of input keys.
        output_keys (Tuple[str, ...]): Name of output keys.
    """
    super().__init__()
    self.input_keys = input_keys
    self.output_keys = output_keys
    self.trunk_dim = trunk_dim
    self.trunk_mlp = LatentMLP(trunk_dim, n_dim, n_dim, n_layer)
    self.branch_mlp = LatentMLP(branch_dim, n_dim, n_dim, n_layer)
    self.mode_mlp = LatentMLP(n_dim, n_dim, n_mode, n_layer)
    self.out_mlp = LatentMLP(n_dim, n_dim, out_dim, n_layer)

    self.attn_blocks = paddle.nn.Sequential(
        *[AttentionBlock(n_mode, n_dim, n_head) for _ in range(n_block)]
    )
    self.apply(self._init_weights)

forward(inputs)

Forward pass of LatentNO.

Parameters:

Name	Type	Description	Default
inputs	dict[str, Tensor]	Dictionary with keys: - "x": Trunk input tensor of shape (B, N, trunk_dim). - "y1": Branch input tensor of shape (B, N, branch_dim).	required

Returns:

Type	Description
dict[str, Tensor]	dict[str, paddle.Tensor]: Dictionary containing: - "y2": Output tensor of shape (B, N, out_dim).

Source code in ppsci/arch/latent_no.py

def forward(self, inputs: dict[str, paddle.Tensor]) -> dict[str, paddle.Tensor]:
    """
    Forward pass of LatentNO.

    Args:
        inputs (dict[str, paddle.Tensor]):
            Dictionary with keys:
                - "x": Trunk input tensor of shape (B, N, trunk_dim).
                - "y1": Branch input tensor of shape (B, N, branch_dim).

    Returns:
        dict[str, paddle.Tensor]: Dictionary containing:
            - "y2": Output tensor of shape (B, N, out_dim).
    """
    x = inputs[self.input_keys[0]]  # trunk input
    y = inputs[self.input_keys[1]]  # branch input

    x = self.trunk_mlp(x)
    y = self.branch_mlp(y)

    score = self.mode_mlp(x)
    score_encode = paddle.nn.functional.softmax(score, axis=1)
    score_decode = paddle.nn.functional.softmax(score, axis=-1)

    z = paddle.matmul(paddle.transpose(score_encode, perm=[0, 2, 1]), y)
    for block in self.attn_blocks:
        z = block(z)

    r = paddle.matmul(score_decode, z)
    r = self.out_mlp(r)

    return {self.output_keys[0]: r}

LNO

Bases: Arch

Laplace Neural Operator net.

Parameters:

Name	Type	Description	Default
input_keys	Tuple[str, ...]	Name of input keys, such as ("input1", "input2").	required
output_keys	Tuple[str, ...]	Name of output keys, such as ("output1", "output2").	required
width	int	Tensor width of Laplace Layer.	required
modes	Tuple[int, ...]	Number of modes to use for contraction in Laplace domain during training.	required
T	Tensor	Linspace of time dimension.	required
data	Tuple[Tensor, ...]	Linspaces of other dimensions.	None
in_features	int	Number of input channels of the first layer.. Defaults to 1.	1
hidden_features	int	Number of channels of the fully-connected layer. Defaults to 64.	64
activation	str	The activation function. Defaults to "sin".	'sin'
use_norm	bool	Whether to use normalization layers. Defaults to True.	True
use_grid	bool	Whether to create grid. Defaults to False.	False

Source code in ppsci/arch/lno.py

class LNO(base.Arch):
    """Laplace Neural Operator net.

    Args:
        input_keys (Tuple[str, ...]): Name of input keys, such as ("input1", "input2").
        output_keys (Tuple[str, ...]): Name of output keys, such as ("output1", "output2").
        width (int): Tensor width of Laplace Layer.
        modes (Tuple[int, ...]): Number of modes to use for contraction in Laplace domain during training.
        T (paddle.Tensor): Linspace of time dimension.
        data (Tuple[paddle.Tensor, ...]): Linspaces of other dimensions.
        in_features (int, optional): Number of input channels of the first layer.. Defaults to 1.
        hidden_features (int, optional): Number of channels of the fully-connected layer. Defaults to 64.
        activation (str, optional): The activation function. Defaults to "sin".
        use_norm (bool, optional): Whether to use normalization layers. Defaults to True.
        use_grid (bool, optional): Whether to create grid. Defaults to False.
    """

    def __init__(
        self,
        input_keys: Tuple[str, ...],
        output_keys: Tuple[str, ...],
        width: int,
        modes: Tuple[int, ...],
        T: paddle.Tensor,
        data: Optional[Tuple[paddle.Tensor, ...]] = None,
        in_features: int = 1,
        hidden_features: int = 64,
        activation: str = "sin",
        use_norm: bool = True,
        use_grid: bool = False,
    ):
        super().__init__()
        self.input_keys = input_keys
        self.output_keys = output_keys
        self.width = width
        self.modes = modes
        self.dims = len(modes)
        assert self.dims <= 3, "Only 3 dims and lower of modes are supported now."

        if data is None:
            data = ()
        assert (
            self.dims == len(data) + 1
        ), f"Dims of modes is {self.dims} but only {len(data)} dims(except T) of data received."

        self.fc0 = nn.Linear(in_features=in_features, out_features=self.width)
        self.laplace = Laplace(self.width, self.width, self.modes, T, data)
        self.conv = getattr(nn, f"Conv{self.dims}D")(
            in_channels=self.width,
            out_channels=self.width,
            kernel_size=1,
            data_format="NCDHW",
        )
        if use_norm:
            self.norm = getattr(nn, f"InstanceNorm{self.dims}D")(
                num_features=self.width,
                weight_attr=False,
                bias_attr=False,
            )
        self.fc1 = nn.Linear(in_features=self.width, out_features=hidden_features)
        self.fc2 = nn.Linear(in_features=hidden_features, out_features=1)
        self.act = act_mod.get_activation(activation)

        self.use_norm = use_norm
        self.use_grid = use_grid

    def get_grid(self, shape):
        batchsize, size_t, size_x, size_y = shape[0], shape[1], shape[2], shape[3]
        gridt = paddle.linspace(0, 1, size_t)
        gridt = gridt.reshape([1, size_t, 1, 1, 1]).tile(
            [batchsize, 1, size_x, size_y, 1]
        )
        gridx = paddle.linspace(0, 1, size_x)
        gridx = gridx.reshape([1, 1, size_x, 1, 1]).tile(
            [batchsize, size_t, 1, size_y, 1]
        )
        gridy = paddle.linspace(0, 1, size_y)
        gridy = gridy.reshape([1, 1, 1, size_y, 1]).tile(
            [batchsize, size_t, size_x, 1, 1]
        )
        return paddle.concat([gridt, gridx, gridy], axis=-1)

    def transpoe_to_NCDHW(self, x):
        perm = [0, self.dims + 1] + list(range(1, self.dims + 1))
        return paddle.transpose(x, perm=perm)

    def transpoe_to_NDHWC(self, x):
        perm = [0] + list(range(2, self.dims + 2)) + [1]
        return paddle.transpose(x, perm=perm)

    def forward_tensor(self, x):
        if self.use_grid:
            grid = self.get_grid(x.shape)
            x = paddle.concat([x, grid], axis=-1)
        x = self.fc0(x)
        x = self.transpoe_to_NCDHW(x)

        if self.use_norm:
            x1 = self.norm(self.laplace(self.norm(x)))
        else:
            x1 = self.laplace(x)

        x2 = self.conv(x)
        x = x1 + x2

        x = self.transpoe_to_NDHWC(x)

        x = self.fc1(x)
        x = self.act(x)
        x = self.fc2(x)
        return x

    def forward(self, x):
        if self._input_transform is not None:
            x = self._input_transform(x)

        y = self.concat_to_tensor(x, self.input_keys, axis=-1)
        y = self.forward_tensor(y)
        y = self.split_to_dict(y, self.output_keys, axis=-1)

        if self._output_transform is not None:
            y = self._output_transform(x, y)
        return y

LorenzEmbedding

Bases: Arch

Embedding Koopman model for the Lorenz ODE system.

Parameters:

Name	Type	Description	Default
input_keys	Tuple[str, ...]	Input keys, such as ("states",).	required
output_keys	Tuple[str, ...]	Output keys, such as ("pred_states", "recover_states").	required
mean	Optional[Tuple[float, ...]]	Mean of training dataset. Defaults to None.	None
std	Optional[Tuple[float, ...]]	Standard Deviation of training dataset. Defaults to None.	None
input_size	int	Size of input data. Defaults to 3.	3
hidden_size	int	Number of hidden size. Defaults to 500.	500
embed_size	int	Number of embedding size. Defaults to 32.	32
drop	float	Probability of dropout the units. Defaults to 0.0.	0.0

Examples:

>>> import ppsci
>>> model = ppsci.arch.LorenzEmbedding(
...     input_keys=("x", "y"),
...     output_keys=("u", "v"),
...     input_size=3,
...     hidden_size=500,
...     embed_size=32,
...     drop=0.0,
...     mean=None,
...     std=None,
... )
>>> x_shape = [8, 3, 2]
>>> y_shape = [8, 3, 1]
>>> input_dict = {"x": paddle.rand(x_shape),
...               "y": paddle.rand(y_shape)}
>>> output_dict = model(input_dict)
>>> print(output_dict["u"].shape)
[8, 2, 3]
>>> print(output_dict["v"].shape)
[8, 3, 3]

Source code in ppsci/arch/embedding_koopman.py

class LorenzEmbedding(base.Arch):
    """Embedding Koopman model for the Lorenz ODE system.

    Args:
        input_keys (Tuple[str, ...]): Input keys, such as ("states",).
        output_keys (Tuple[str, ...]): Output keys, such as ("pred_states", "recover_states").
        mean (Optional[Tuple[float, ...]]): Mean of training dataset. Defaults to None.
        std (Optional[Tuple[float, ...]]): Standard Deviation of training dataset. Defaults to None.
        input_size (int, optional): Size of input data. Defaults to 3.
        hidden_size (int, optional): Number of hidden size. Defaults to 500.
        embed_size (int, optional): Number of embedding size. Defaults to 32.
        drop (float, optional):  Probability of dropout the units. Defaults to 0.0.

    Examples:
        >>> import ppsci
        >>> model = ppsci.arch.LorenzEmbedding(
        ...     input_keys=("x", "y"),
        ...     output_keys=("u", "v"),
        ...     input_size=3,
        ...     hidden_size=500,
        ...     embed_size=32,
        ...     drop=0.0,
        ...     mean=None,
        ...     std=None,
        ... )
        >>> x_shape = [8, 3, 2]
        >>> y_shape = [8, 3, 1]
        >>> input_dict = {"x": paddle.rand(x_shape),
        ...               "y": paddle.rand(y_shape)}
        >>> output_dict = model(input_dict)
        >>> print(output_dict["u"].shape)
        [8, 2, 3]
        >>> print(output_dict["v"].shape)
        [8, 3, 3]
    """

    def __init__(
        self,
        input_keys: Tuple[str, ...],
        output_keys: Tuple[str, ...],
        mean: Optional[Tuple[float, ...]] = None,
        std: Optional[Tuple[float, ...]] = None,
        input_size: int = 3,
        hidden_size: int = 500,
        embed_size: int = 32,
        drop: float = 0.0,
    ):
        super().__init__()
        self.input_keys = input_keys
        self.output_keys = output_keys
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.embed_size = embed_size

        # build observable network
        self.encoder_net = self.build_encoder(input_size, hidden_size, embed_size, drop)
        # build koopman operator
        self.k_diag, self.k_ut = self.build_koopman_operator(embed_size)
        # build recovery network
        self.decoder_net = self.build_decoder(input_size, hidden_size, embed_size)

        mean = [0.0, 0.0, 0.0] if mean is None else mean
        std = [1.0, 1.0, 1.0] if std is None else std
        self.register_buffer("mean", paddle.to_tensor(mean).reshape([1, 3]))
        self.register_buffer("std", paddle.to_tensor(std).reshape([1, 3]))

        self.apply(self._init_weights)

    def _init_weights(self, m: nn.Layer):
        if isinstance(m, nn.Linear):
            k = 1 / m.weight.shape[0]
            uniform = Uniform(-(k**0.5), k**0.5)
            uniform(m.weight)
            if m.bias is not None:
                uniform(m.bias)
        elif isinstance(m, nn.LayerNorm):
            zeros_(m.bias)
            ones_(m.weight)

    def build_encoder(
        self, input_size: int, hidden_size: int, embed_size: int, drop: float = 0.0
    ):
        net = nn.Sequential(
            nn.Linear(input_size, hidden_size),
            nn.ReLU(),
            nn.Linear(hidden_size, embed_size),
            nn.LayerNorm(embed_size),
            nn.Dropout(drop),
        )
        return net

    def build_decoder(self, input_size: int, hidden_size: int, embed_size: int):
        net = nn.Sequential(
            nn.Linear(embed_size, hidden_size),
            nn.ReLU(),
            nn.Linear(hidden_size, input_size),
        )
        return net

    def build_koopman_operator(self, embed_size: int):
        # Learned Koopman operator
        data = paddle.linspace(1, 0, embed_size)
        k_diag = paddle.create_parameter(
            shape=data.shape,
            dtype=paddle.get_default_dtype(),
            default_initializer=nn.initializer.Assign(data),
        )

        data = 0.1 * paddle.rand([2 * embed_size - 3])
        k_ut = paddle.create_parameter(
            shape=data.shape,
            dtype=paddle.get_default_dtype(),
            default_initializer=nn.initializer.Assign(data),
        )
        return k_diag, k_ut

    def encoder(self, x: paddle.Tensor):
        x = self._normalize(x)
        g = self.encoder_net(x)
        return g

    def decoder(self, g: paddle.Tensor):
        out = self.decoder_net(g)
        x = self._unnormalize(out)
        return x

    def koopman_operation(self, embed_data: paddle.Tensor, k_matrix: paddle.Tensor):
        # Apply Koopman operation
        embed_pred_data = paddle.bmm(
            k_matrix.expand(
                [embed_data.shape[0], k_matrix.shape[0], k_matrix.shape[1]]
            ),
            embed_data.transpose([0, 2, 1]),
        ).transpose([0, 2, 1])
        return embed_pred_data

    def _normalize(self, x: paddle.Tensor):
        return (x - self.mean) / self.std

    def _unnormalize(self, x: paddle.Tensor):
        return self.std * x + self.mean

    def get_koopman_matrix(self):
        # # Koopman operator
        k_ut_tensor = self.k_ut * 1
        k_ut_tensor = paddle.diag(
            k_ut_tensor[0 : self.embed_size - 1], offset=1
        ) + paddle.diag(k_ut_tensor[self.embed_size - 1 :], offset=2)
        k_matrix = k_ut_tensor + (-1) * k_ut_tensor.t()
        k_matrix = k_matrix + paddle.diag(self.k_diag)
        return k_matrix

    def forward_tensor(self, x):
        k_matrix = self.get_koopman_matrix()
        embed_data = self.encoder(x)
        recover_data = self.decoder(embed_data)

        embed_pred_data = self.koopman_operation(embed_data, k_matrix)
        pred_data = self.decoder(embed_pred_data)

        return (pred_data[:, :-1, :], recover_data, k_matrix)

    @staticmethod
    def split_to_dict(data_tensors: Tuple[paddle.Tensor, ...], keys: Tuple[str, ...]):
        return {key: data_tensors[i] for i, key in enumerate(keys)}

    def forward(self, x):
        if self._input_transform is not None:
            x = self._input_transform(x)

        x_tensor = self.concat_to_tensor(x, self.input_keys, axis=-1)
        y = self.forward_tensor(x_tensor)
        y = self.split_to_dict(y, self.output_keys)

        if self._output_transform is not None:
            y = self._output_transform(x, y)
        return y

Meteoformer

Bases: Arch

Meteoformer is a class that represents a Spatial-Temporal Transformer model designed for short-to-medium-term weather prediction with multiple meteorological variables.

Parameters:

Name	Type	Description	Default
input_keys	Tuple[str, ...]	A tuple of input keys.	required
output_keys	Tuple[str, ...]	A tuple of output keys.	required
shape_in	Tuple[int, ...]	The shape of the input data (T, C, H, W), where T is the number of time steps, C is the number of channels, H and W are the spatial dimensions.	required
hid_S	int	The number of hidden channels in the spatial encoder.	64
hid_T	int	The number of hidden units in the temporal encoder.	256
N_S	int	The number of spatial transformer layers.	4
N_T	int	The number of temporal transformer layers.	4
incep_ker	Tuple[int, ...]	The kernel sizes used in the inception block.	(3, 5, 7, 11)
groups	int	The number of groups for grouped convolutions.	8
num_classes	int	The number of predicted meteorological variables.	12

Examples:

>>> import paddle
>>> import ppsci
>>> model = ppsci.arch.Meteoformer(
...     input_keys=("input",),
...     output_keys=("output",),
...     shape_in=(6, 12, 192, 256),
...     hid_S=64,
...     hid_T=256,
...     N_S=4,
...     N_T=4,
...     incep_ker=(3, 5, 7, 11),
...     groups=8,
...     num_classes=4,
... )
>>> input_dict = {"input": paddle.rand([8, 6, 12, 192, 256])}
>>> output_dict = model(input_dict)
>>> print(output_dict["output"].shape)
[8, 6, 12, 192, 256]

Source code in ppsci/arch/meteoformer.py

class Meteoformer(base.Arch):
    """
    Meteoformer is a class that represents a Spatial-Temporal Transformer model designed for short-to-medium-term weather prediction with multiple meteorological variables.

    Args:
        input_keys (Tuple[str, ...]): A tuple of input keys.
        output_keys (Tuple[str, ...]): A tuple of output keys.
        shape_in (Tuple[int, ...]): The shape of the input data (T, C, H, W), where
            T is the number of time steps, C is the number of channels,
            H and W are the spatial dimensions.
        hid_S (int): The number of hidden channels in the spatial encoder.
        hid_T (int): The number of hidden units in the temporal encoder.
        N_S (int): The number of spatial transformer layers.
        N_T (int): The number of temporal transformer layers.
        incep_ker (Tuple[int, ...]): The kernel sizes used in the inception block.
        groups (int): The number of groups for grouped convolutions.
        num_classes (int): The number of predicted meteorological variables.

    Examples:
        >>> import paddle
        >>> import ppsci
        >>> model = ppsci.arch.Meteoformer(
        ...     input_keys=("input",),
        ...     output_keys=("output",),
        ...     shape_in=(6, 12, 192, 256),
        ...     hid_S=64,
        ...     hid_T=256,
        ...     N_S=4,
        ...     N_T=4,
        ...     incep_ker=(3, 5, 7, 11),
        ...     groups=8,
        ...     num_classes=4,
        ... )
        >>> input_dict = {"input": paddle.rand([8, 6, 12, 192, 256])}
        >>> output_dict = model(input_dict)
        >>> print(output_dict["output"].shape)
        [8, 6, 12, 192, 256]
    """

    def __init__(
        self,
        input_keys: Tuple[str, ...],
        output_keys: Tuple[str, ...],
        shape_in: Tuple[int, ...],
        hid_S: int = 64,
        hid_T: int = 256,
        N_S: int = 4,
        N_T: int = 4,
        incep_ker: Tuple[int, ...] = (3, 5, 7, 11),
        groups: int = 8,
        num_classes: int = 12,
    ):
        super().__init__()
        self.input_keys = input_keys
        self.output_keys = output_keys
        self.num_classes = num_classes

        T, C, H, W = shape_in
        self.enc = Encoder(C, hid_S, N_S)
        self.hid1 = MidXnet(T * hid_S, hid_T // 2, N_T, incep_ker, groups)
        self.dec = Decoder(T * hid_S, T * self.num_classes, N_S)

    def forward(self, x):
        if self._input_transform is not None:
            x = self._input_transform(x)

        x = self.concat_to_tensor(x, self.input_keys)

        B, T, C, H, W = x.shape
        x = x.reshape([B * T, C, H, W])

        # encoded
        embed = self.enc(x)
        _, C_4, H_4, W_4 = embed[-1].shape

        # translator
        z = embed[-1].reshape([B, T, C_4, H_4, W_4])
        hid = self.hid1(z)
        hid = hid.transpose(perm=[0, 2, 1]).reshape([B, -1, H_4, W_4])

        # decoded
        y = self.dec(hid, embed[0])
        y = y.reshape([B, T, self.num_classes, H, W])

        y = self.split_to_dict(y, self.output_keys)
        if self._output_transform is not None:
            y = self._output_transform(x, y)

        return y  # {self.output_keys[0]: Y}

MLP

Bases: Arch

Multi layer perceptron network.

Parameters:

Name	Type	Description	Default
input_keys	Tuple[str, ...]	Name of input keys, such as ("x", "y", "z").	required
output_keys	Tuple[str, ...]	Name of output keys, such as ("u", "v", "w").	required
num_layers	int	Number of hidden layers.	required
hidden_size	Union[int, Tuple[int, ...]]	Number of hidden size. An integer for all layers, or list of integer specify each layer's size.	required
activation	str	Name of activation function. Defaults to "tanh".	'tanh'
skip_connection	bool	Whether to use skip connection. Defaults to False.	False
weight_norm	bool	Whether to apply weight norm on parameter(s). Defaults to False.	False
input_dim	Optional[int]	Number of input's dimension. Defaults to None.	None
output_dim	Optional[int]	Number of output's dimension. Defaults to None.	None
periods	Optional[Dict[int, Tuple[float, bool]]]	Period of each input key, input in given channel will be period embedded if specified, each tuple of periods list is [period, trainable]. Defaults to None.	None
fourier	Optional[Dict[str, Union[float, int]]]	Random fourier feature embedding, e.g. {'dim': 256, 'scale': 1.0}. Defaults to None.	None
random_weight	Optional[Dict[str, float]]	Mean and std of random weight factorization layer, e.g. {"mean": 0.5, "std: 0.1"}. Defaults to None.	None

Examples:

>>> import paddle
>>> import ppsci
>>> model = ppsci.arch.MLP(
...     input_keys=("x", "y"),
...     output_keys=("u", "v"),
...     num_layers=5,
...     hidden_size=128
... )
>>> input_dict = {"x": paddle.rand([64, 1]),
...               "y": paddle.rand([64, 1])}
>>> output_dict = model(input_dict)
>>> print(output_dict["u"].shape)
[64, 1]
>>> print(output_dict["v"].shape)
[64, 1]

Source code in ppsci/arch/mlp.py

class MLP(base.Arch):
    """Multi layer perceptron network.

    Args:
        input_keys (Tuple[str, ...]): Name of input keys, such as ("x", "y", "z").
        output_keys (Tuple[str, ...]): Name of output keys, such as ("u", "v", "w").
        num_layers (int): Number of hidden layers.
        hidden_size (Union[int, Tuple[int, ...]]): Number of hidden size.
            An integer for all layers, or list of integer specify each layer's size.
        activation (str, optional): Name of activation function. Defaults to "tanh".
        skip_connection (bool, optional): Whether to use skip connection. Defaults to False.
        weight_norm (bool, optional): Whether to apply weight norm on parameter(s). Defaults to False.
        input_dim (Optional[int]): Number of input's dimension. Defaults to None.
        output_dim (Optional[int]): Number of output's dimension. Defaults to None.
        periods (Optional[Dict[int, Tuple[float, bool]]]): Period of each input key,
            input in given channel will be period embedded if specified, each tuple of
            periods list is [period, trainable]. Defaults to None.
        fourier (Optional[Dict[str, Union[float, int]]]): Random fourier feature embedding,
            e.g. {'dim': 256, 'scale': 1.0}. Defaults to None.
        random_weight (Optional[Dict[str, float]]): Mean and std of random weight
            factorization layer, e.g. {"mean": 0.5, "std: 0.1"}. Defaults to None.

    Examples:
        >>> import paddle
        >>> import ppsci
        >>> model = ppsci.arch.MLP(
        ...     input_keys=("x", "y"),
        ...     output_keys=("u", "v"),
        ...     num_layers=5,
        ...     hidden_size=128
        ... )
        >>> input_dict = {"x": paddle.rand([64, 1]),
        ...               "y": paddle.rand([64, 1])}
        >>> output_dict = model(input_dict)
        >>> print(output_dict["u"].shape)
        [64, 1]
        >>> print(output_dict["v"].shape)
        [64, 1]
    """

    def __init__(
        self,
        input_keys: Tuple[str, ...],
        output_keys: Tuple[str, ...],
        num_layers: int,
        hidden_size: Union[int, Tuple[int, ...]],
        activation: str = "tanh",
        skip_connection: bool = False,
        weight_norm: bool = False,
        input_dim: Optional[int] = None,
        output_dim: Optional[int] = None,
        periods: Optional[Dict[int, Tuple[float, bool]]] = None,
        fourier: Optional[Dict[str, Union[float, int]]] = None,
        random_weight: Optional[Dict[str, float]] = None,
    ):
        super().__init__()
        self.input_keys = input_keys
        self.output_keys = output_keys
        self.linears = []
        self.acts = []
        self.periods = periods
        self.fourier = fourier
        if periods:
            self.period_emb = PeriodEmbedding(periods)

        if isinstance(hidden_size, (tuple, list)):
            if num_layers is not None:
                raise ValueError(
                    "num_layers should be None when hidden_size is specified"
                )
        elif isinstance(hidden_size, int):
            if not isinstance(num_layers, int):
                raise ValueError(
                    "num_layers should be an int when hidden_size is an int"
                )
            hidden_size = [hidden_size] * num_layers
        else:
            raise ValueError(
                f"hidden_size should be list of int or int, but got {type(hidden_size)}"
            )

        # initialize FC layer(s)
        cur_size = len(self.input_keys) if input_dim is None else input_dim
        if input_dim is None and periods:
            # period embedded channel(s) will be doubled automatically
            # if input_dim is not specified
            cur_size += len(periods)

        if fourier:
            self.fourier_emb = FourierEmbedding(
                cur_size, fourier["dim"], fourier["scale"]
            )
            cur_size = fourier["dim"]

        for i, _size in enumerate(hidden_size):
            if weight_norm:
                self.linears.append(WeightNormLinear(cur_size, _size))
            elif random_weight:
                self.linears.append(
                    RandomWeightFactorization(
                        cur_size,
                        _size,
                        mean=random_weight["mean"],
                        std=random_weight["std"],
                    )
                )
            else:
                self.linears.append(nn.Linear(cur_size, _size))

            # initialize activation function
            self.acts.append(
                act_mod.get_activation(activation)
                if activation != "stan"
                else act_mod.get_activation(activation)(_size)
            )
            # special initialization for certain activation
            # TODO: Adapt code below to a more elegant style
            if activation == "siren":
                if i == 0:
                    act_mod.Siren.init_for_first_layer(self.linears[-1])
                else:
                    act_mod.Siren.init_for_hidden_layer(self.linears[-1])

            cur_size = _size

        self.linears = nn.LayerList(self.linears)
        self.acts = nn.LayerList(self.acts)
        if random_weight:
            self.last_fc = RandomWeightFactorization(
                cur_size,
                len(self.output_keys) if output_dim is None else output_dim,
                mean=random_weight["mean"],
                std=random_weight["std"],
            )
        else:
            self.last_fc = nn.Linear(
                cur_size,
                len(self.output_keys) if output_dim is None else output_dim,
            )

        self.skip_connection = skip_connection

    def forward_tensor(self, x):
        y = x
        skip = None
        for i, linear in enumerate(self.linears):
            y = linear(y)
            if self.skip_connection and i % 2 == 0:
                if skip is not None:
                    skip = y
                    y = y + skip
                else:
                    skip = y
            y = self.acts[i](y)

        y = self.last_fc(y)

        return y

    def forward(self, x):
        if self._input_transform is not None:
            x = self._input_transform(x)

        if self.periods:
            x = self.period_emb(x)

        y = self.concat_to_tensor(x, self.input_keys, axis=-1)

        if self.fourier:
            y = self.fourier_emb(y)

        y = self.forward_tensor(y)
        y = self.split_to_dict(y, self.output_keys, axis=-1)

        if self._output_transform is not None:
            y = self._output_transform(x, y)
        return y

ModelList

Bases: Arch

ModelList layer which wrap more than one model that shares inputs.

Parameters:

Name	Type	Description	Default
model_list	Tuple[Arch, ...]	Model(s) nested in tuple.	required

Examples:

>>> import paddle
>>> import ppsci
>>> model1 = ppsci.arch.MLP(("x", "y"), ("u", "v"), 10, 128)
>>> model2 = ppsci.arch.MLP(("x", "y"), ("w", "p"), 5, 128)
>>> model = ppsci.arch.ModelList((model1, model2))
>>> input_dict = {"x": paddle.rand([64, 64, 1]),"y": paddle.rand([64, 64, 1])}
>>> output_dict = model(input_dict)
>>> for k, v in output_dict.items():
...     print(k, v.shape)
u [64, 64, 1]
v [64, 64, 1]
w [64, 64, 1]
p [64, 64, 1]

Source code in ppsci/arch/model_list.py

class ModelList(base.Arch):
    """ModelList layer which wrap more than one model that shares inputs.

    Args:
        model_list (Tuple[base.Arch, ...]): Model(s) nested in tuple.

    Examples:
        >>> import paddle
        >>> import ppsci
        >>> model1 = ppsci.arch.MLP(("x", "y"), ("u", "v"), 10, 128)
        >>> model2 = ppsci.arch.MLP(("x", "y"), ("w", "p"), 5, 128)
        >>> model = ppsci.arch.ModelList((model1, model2))
        >>> input_dict = {"x": paddle.rand([64, 64, 1]),"y": paddle.rand([64, 64, 1])}
        >>> output_dict = model(input_dict)
        >>> for k, v in output_dict.items():
        ...     print(k, v.shape)
        u [64, 64, 1]
        v [64, 64, 1]
        w [64, 64, 1]
        p [64, 64, 1]
    """

    def __init__(
        self,
        model_list: Tuple[base.Arch, ...],
    ):
        super().__init__()
        self.input_keys = sum([model.input_keys for model in model_list], ())
        self.input_keys = set(self.input_keys)

        output_keys_set = set()
        for model in model_list:
            if len(output_keys_set & set(model.output_keys)):
                raise ValueError(
                    "output_keys of model from model_list should be unique,"
                    f"but got duplicate keys: {output_keys_set & set(model.output_keys)}"
                )
            output_keys_set = output_keys_set | set(model.output_keys)
        self.output_keys = tuple(output_keys_set)

        self.model_list = nn.LayerList(model_list)

    def forward(self, x):
        y_all = {}
        for model in self.model_list:
            y = model(x)
            y_all.update(y)

        return y_all

ModifiedMLP

Bases: Arch

Modified Multi layer perceptron network.

Understanding and mitigating gradient pathologies in physics-informed neural networks. https://arxiv.org/pdf/2001.04536.pdf.

Parameters:

Name	Type	Description	Default
input_keys	Tuple[str, ...]	Name of input keys, such as ("x", "y", "z").	required
output_keys	Tuple[str, ...]	Name of output keys, such as ("u", "v", "w").	required
num_layers	int	Number of hidden layers.	required
hidden_size	int	Number of hidden size, an integer for all layers.	required
activation	str	Name of activation function. Defaults to "tanh".	'tanh'
skip_connection	bool	Whether to use skip connection. Defaults to False.	False
weight_norm	bool	Whether to apply weight norm on parameter(s). Defaults to False.	False
input_dim	Optional[int]	Number of input's dimension. Defaults to None.	None
output_dim	Optional[int]	Number of output's dimension. Defaults to None.	None

Examples:

>>> import paddle
>>> import ppsci
>>> model = ppsci.arch.ModifiedMLP(
...     input_keys=("x", "y"),
...     output_keys=("u", "v"),
...     num_layers=5,
...     hidden_size=128
... )
>>> input_dict = {"x": paddle.rand([64, 1]),
...               "y": paddle.rand([64, 1])}
>>> output_dict = model(input_dict)
>>> print(output_dict["u"].shape)
[64, 1]
>>> print(output_dict["v"].shape)
[64, 1]

Source code in ppsci/arch/mlp.py

class ModifiedMLP(base.Arch):
    """Modified Multi layer perceptron network.

    Understanding and mitigating gradient pathologies in physics-informed
    neural networks. https://arxiv.org/pdf/2001.04536.pdf.

    Args:
        input_keys (Tuple[str, ...]): Name of input keys, such as ("x", "y", "z").
        output_keys (Tuple[str, ...]): Name of output keys, such as ("u", "v", "w").
        num_layers (int): Number of hidden layers.
        hidden_size (int): Number of hidden size, an integer for all layers.
        activation (str, optional): Name of activation function. Defaults to "tanh".
        skip_connection (bool, optional): Whether to use skip connection. Defaults to False.
        weight_norm (bool, optional): Whether to apply weight norm on parameter(s). Defaults to False.
        input_dim (Optional[int]): Number of input's dimension. Defaults to None.
        output_dim (Optional[int]): Number of output's dimension. Defaults to None.

    Examples:
        >>> import paddle
        >>> import ppsci
        >>> model = ppsci.arch.ModifiedMLP(
        ...     input_keys=("x", "y"),
        ...     output_keys=("u", "v"),
        ...     num_layers=5,
        ...     hidden_size=128
        ... )
        >>> input_dict = {"x": paddle.rand([64, 1]),
        ...               "y": paddle.rand([64, 1])}
        >>> output_dict = model(input_dict)
        >>> print(output_dict["u"].shape)
        [64, 1]
        >>> print(output_dict["v"].shape)
        [64, 1]
    """

    def __init__(
        self,
        input_keys: Tuple[str, ...],
        output_keys: Tuple[str, ...],
        num_layers: int,
        hidden_size: int,
        activation: str = "tanh",
        skip_connection: bool = False,
        weight_norm: bool = False,
        input_dim: Optional[int] = None,
        output_dim: Optional[int] = None,
        periods: Optional[Dict[int, Tuple[float, bool]]] = None,
        fourier: Optional[Dict[str, Union[float, int]]] = None,
        random_weight: Optional[Dict[str, float]] = None,
    ):
        super().__init__()
        self.input_keys = input_keys
        self.output_keys = output_keys
        self.linears = []
        self.acts = []
        self.periods = periods
        self.fourier = fourier
        if periods:
            self.period_emb = PeriodEmbedding(periods)
        if isinstance(hidden_size, int):
            if not isinstance(num_layers, int):
                raise ValueError("num_layers should be an int")
            hidden_size = [hidden_size] * num_layers
        else:
            raise ValueError(f"hidden_size should be int, but got {type(hidden_size)}")

        # initialize FC layer(s)
        cur_size = len(self.input_keys) if input_dim is None else input_dim
        if input_dim is None and periods:
            # period embedded channel(s) will be doubled automatically
            # if input_dim is not specified
            cur_size += len(periods)

        if fourier:
            self.fourier_emb = FourierEmbedding(
                cur_size, fourier["dim"], fourier["scale"]
            )
            cur_size = fourier["dim"]

        self.embed_u = nn.Sequential(
            (
                WeightNormLinear(cur_size, hidden_size[0])
                if weight_norm
                else (
                    nn.Linear(cur_size, hidden_size[0])
                    if random_weight is None
                    else RandomWeightFactorization(
                        cur_size,
                        hidden_size[0],
                        mean=random_weight["mean"],
                        std=random_weight["std"],
                    )
                )
            ),
            (
                act_mod.get_activation(activation)
                if activation != "stan"
                else act_mod.get_activation(activation)(hidden_size[0])
            ),
        )
        self.embed_v = nn.Sequential(
            (
                WeightNormLinear(cur_size, hidden_size[0])
                if weight_norm
                else (
                    nn.Linear(cur_size, hidden_size[0])
                    if random_weight is None
                    else RandomWeightFactorization(
                        cur_size,
                        hidden_size[0],
                        mean=random_weight["mean"],
                        std=random_weight["std"],
                    )
                )
            ),
            (
                act_mod.get_activation(activation)
                if activation != "stan"
                else act_mod.get_activation(activation)(hidden_size[0])
            ),
        )

        for i, _size in enumerate(hidden_size):
            if weight_norm:
                self.linears.append(WeightNormLinear(cur_size, _size))
            elif random_weight:
                self.linears.append(
                    RandomWeightFactorization(
                        cur_size,
                        _size,
                        mean=random_weight["mean"],
                        std=random_weight["std"],
                    )
                )
            else:
                self.linears.append(nn.Linear(cur_size, _size))

            # initialize activation function
            self.acts.append(
                act_mod.get_activation(activation)
                if activation != "stan"
                else act_mod.get_activation(activation)(_size)
            )
            # special initialization for certain activation
            # TODO: Adapt code below to a more elegant style
            if activation == "siren":
                if i == 0:
                    act_mod.Siren.init_for_first_layer(self.linears[-1])
                else:
                    act_mod.Siren.init_for_hidden_layer(self.linears[-1])

            cur_size = _size

        self.linears = nn.LayerList(self.linears)
        self.acts = nn.LayerList(self.acts)
        if random_weight:
            self.last_fc = RandomWeightFactorization(
                cur_size,
                len(self.output_keys) if output_dim is None else output_dim,
                mean=random_weight["mean"],
                std=random_weight["std"],
            )
        else:
            self.last_fc = nn.Linear(
                cur_size,
                len(self.output_keys) if output_dim is None else output_dim,
            )

        self.skip_connection = skip_connection

    def forward_tensor(self, x):
        u = self.embed_u(x)
        v = self.embed_v(x)

        y = x
        skip = None
        for i, linear in enumerate(self.linears):
            y = linear(y)
            y = self.acts[i](y)
            y = y * u + (1 - y) * v
            if self.skip_connection and i % 2 == 0:
                if skip is not None:
                    skip = y
                    y = y + skip
                else:
                    skip = y

        y = self.last_fc(y)

        return y

    def forward(self, x):
        x_identity = x
        if self._input_transform is not None:
            x = self._input_transform(x)

        if self.periods:
            x = self.period_emb(x)

        y = self.concat_to_tensor(x, self.input_keys, axis=-1)

        if self.fourier:
            y = self.fourier_emb(y)

        y = self.forward_tensor(y)
        y = self.split_to_dict(y, self.output_keys, axis=-1)

        if self._output_transform is not None:
            y = self._output_transform(x_identity, y)
        return y

MoleculeModel

Bases: Layer

Model with an MPN encoder followed by feed-forward layers.

Parameters:

Name	Type	Description	Default
cfg	DictConfig	Hydra DictConfig used to build internal TrainArgs.	required

Source code in ppsci/arch/chemprop_molecule.py

class MoleculeModel(paddle.nn.Layer):
    """Model with an MPN encoder followed by feed-forward layers.

    Args:
        cfg: Hydra `DictConfig` used to build internal `TrainArgs`.
    """

    def __init__(self, cfg: DictConfig):
        super(MoleculeModel, self).__init__()
        args = self.build_from_cfg(cfg)
        self.classification = args.dataset_type == "classification"
        self.multiclass = args.dataset_type == "multiclass"
        self.loss_function = args.loss_function
        if hasattr(args, "train_class_sizes"):
            self.train_class_sizes = args.train_class_sizes
        else:
            self.train_class_sizes = None
        if self.classification or self.multiclass:
            self.no_training_normalization = args.loss_function in [
                "cross_entropy",
                "binary_cross_entropy",
            ]
        self.output_size = args.num_tasks
        if self.multiclass:
            self.output_size *= args.multiclass_num_classes
        if self.loss_function == "mve":
            self.output_size *= 2
        if self.loss_function == "dirichlet" and self.classification:
            self.output_size *= 2
        if self.loss_function == "evidential":
            self.output_size *= 4
        if self.classification:
            self.sigmoid = paddle.nn.Sigmoid()
        if self.multiclass:
            self.multiclass_softmax = paddle.nn.Softmax(axis=2)
        if self.loss_function in ["mve", "evidential", "dirichlet"]:
            self.softplus = paddle.nn.Softplus()
        self.create_encoder(args)
        self.create_ffn(args)
        initialize_weights(self)

    def _make_args(
        self,
        dataset_type,
        epochs,
        use_gpu,
        fingerprint_type,
        property_name,
        train_smiles=None,
        train_fingerprints=None,
    ):

        # Create args
        arg_list = [
            "--data_path",
            "foo.csv",
            "--dataset_type",
            dataset_type,
            "--save_dir",
            "foo",
            "--epochs",
            str(epochs),
            "--quiet",
        ] + ([] if use_gpu else ["--no_cuda"])

        if fingerprint_type == "morgan":
            arg_list += ["--features_generator", "morgan"]
        elif fingerprint_type == "rdkit":
            arg_list += [
                "--features_generator",
                "rdkit_2d_normalized",
                "--no_features_scaling",
            ]
        elif fingerprint_type is None:
            pass
        else:
            raise ValueError(f'Fingerprint type "{fingerprint_type}" is not supported.')

        args = TrainArgs().parse_args(arg_list)
        args.task_names = [property_name]
        if train_smiles is not None:
            args.train_data_size = len(train_smiles)

        if fingerprint_type is not None:
            args.features_size = train_fingerprints.shape[1]
        return args

    def build_from_cfg(self, cfg: DictConfig):
        args = self._make_args(
            dataset_type=cfg.DATA.dataset_type,  # "classification",
            epochs=cfg.TRAIN.epochs,  # 1,
            use_gpu=cfg.TRAIN.use_gpu,
            fingerprint_type=cfg.DATA.fingerprint_type,  # None,
            property_name=cfg.DATA.property_column,  # "antibiotic_activity"
            train_smiles=None,
            train_fingerprints=None,
        )
        return args

    def create_encoder(self, args: TrainArgs) -> None:
        """Create the message passing encoder for the model.

        Args:
            args: Model arguments in a `TrainArgs` instance.
        """
        self.encoder = MPN(args)
        if args.checkpoint_frzn is not None:
            if args.freeze_first_only:
                for param in list(self.encoder.encoder.children())[0].parameters():
                    param.stop_gradient = not False
            else:
                for param in self.encoder.parameters():
                    param.stop_gradient = not False

    def create_ffn(self, args: TrainArgs) -> None:
        """Create the feed-forward layers for the model.

        Args:
            args: Model arguments in a `TrainArgs` instance.
        """
        self.multiclass = args.dataset_type == "multiclass"
        if self.multiclass:
            self.num_classes = args.multiclass_num_classes
        if args.features_only:
            first_linear_dim = args.features_size
        else:
            if args.reaction_solvent:
                first_linear_dim = args.hidden_size + args.hidden_size_solvent
            else:
                first_linear_dim = args.hidden_size * args.number_of_molecules
            if args.use_input_features:
                first_linear_dim += args.features_size
        if args.atom_descriptors == "descriptor":
            first_linear_dim += args.atom_descriptors_size
        dropout = paddle.nn.Dropout(p=args.dropout)
        activation = get_activation_function(args.activation)
        if args.ffn_num_layers == 1:
            ffn = [
                dropout,
                paddle.nn.Linear(
                    in_features=first_linear_dim, out_features=self.output_size
                ),
            ]
        else:
            ffn = [
                dropout,
                paddle.nn.Linear(
                    in_features=first_linear_dim, out_features=args.ffn_hidden_size
                ),
            ]
            for _ in range(args.ffn_num_layers - 2):
                ffn.extend(
                    [
                        activation,
                        dropout,
                        paddle.nn.Linear(
                            in_features=args.ffn_hidden_size,
                            out_features=args.ffn_hidden_size,
                        ),
                    ]
                )
            ffn.extend(
                [
                    activation,
                    dropout,
                    paddle.nn.Linear(
                        in_features=args.ffn_hidden_size, out_features=self.output_size
                    ),
                ]
            )
        if args.dataset_type == "spectra":
            if args.spectra_activation == "softplus":
                spectra_activation = paddle.nn.Softplus()
            else:

                class nn_exp(paddle.nn.Layer):
                    def __init__(self):
                        super(nn_exp, self).__init__()

                    def forward(self, x):
                        return paddle.exp(x=x)

                spectra_activation = nn_exp()
            ffn.append(spectra_activation)
        self.ffn = paddle.nn.Sequential(*ffn)
        if args.checkpoint_frzn is not None:
            if args.frzn_ffn_layers > 0:
                for param in list(self.ffn.parameters())[0 : 2 * args.frzn_ffn_layers]:
                    param.stop_gradient = not False

    def fingerprint(
        self,
        batch: Union[
            List[List[str]],
            List[List[Chem.Mol]],
            List[List[Tuple[Chem.Mol, Chem.Mol]]],
            List[BatchMolGraph],
        ],
        features_batch: List[np.ndarray] = None,
        atom_descriptors_batch: List[np.ndarray] = None,
        atom_features_batch: List[np.ndarray] = None,
        bond_features_batch: List[np.ndarray] = None,
        fingerprint_type: str = "MPN",
    ) -> paddle.Tensor:
        """Return latent representations (fingerprints) from intermediate stages.

        Args:
            batch: List of lists of SMILES or RDKit molecules, or a list of `BatchMolGraph`.
            features_batch: Optional list of numpy arrays with additional features.
            atom_descriptors_batch: Optional list of numpy arrays with atom descriptors.
            atom_features_batch: Optional list of numpy arrays with atom features.
            bond_features_batch: Optional list of numpy arrays with bond features.
            fingerprint_type: Which latent representation to return. Supported: "MPN" (output of MPN) or
                "last_FFN" (input to the final readout layer).

        Returns:
            `paddle.Tensor`: Fingerprint vectors.

        Raises:
            ValueError: If `fingerprint_type` is unsupported.
        """
        if fingerprint_type == "MPN":
            return self.encoder(
                batch,
                features_batch,
                atom_descriptors_batch,
                atom_features_batch,
                bond_features_batch,
            )
        elif fingerprint_type == "last_FFN":
            return self.ffn[:-1](
                self.encoder(
                    batch,
                    features_batch,
                    atom_descriptors_batch,
                    atom_features_batch,
                    bond_features_batch,
                )
            )
        else:
            raise ValueError(f"Unsupported fingerprint type {fingerprint_type}.")

    def forward(
        self,
        batch: Union[
            List[List[str]],
            List[List[Chem.Mol]],
            List[List[Tuple[Chem.Mol, Chem.Mol]]],
            List[BatchMolGraph],
        ],
        features_batch: List[np.ndarray] = None,
        atom_descriptors_batch: List[np.ndarray] = None,
        atom_features_batch: List[np.ndarray] = None,
        bond_features_batch: List[np.ndarray] = None,
    ) -> paddle.float32:
        """Run the model forward pass on input.

        Args:
            batch: Dictionary with keys `mol_batch`, `features_batch`,
                `atom_descriptors_batch`, `atom_features_batch`, `bond_features_batch`.
            features_batch: Unused (kept for API symmetry), see `batch`.
            atom_descriptors_batch: Unused, see `batch`.
            atom_features_batch: Unused, see `batch`.
            bond_features_batch: Unused, see `batch`.

        Returns:
            dict: `{"pred": tensor}` predictions (and transformed for special losses).
        """

        mol_batch = batch["mol_batch"]
        features_batch = batch["features_batch"]
        atom_descriptors_batch = batch["atom_descriptors_batch"]
        atom_features_batch = batch["atom_features_batch"]
        bond_features_batch = batch["bond_features_batch"]
        batch = mol_batch

        output = self.ffn(
            self.encoder(
                batch,
                features_batch,
                atom_descriptors_batch,
                atom_features_batch,
                bond_features_batch,
            )
        )
        if (
            self.classification
            and not (self.training and self.no_training_normalization)
            and self.loss_function != "dirichlet"
        ):
            output = self.sigmoid(output)
        if self.multiclass:
            output = output.reshape((tuple(output.shape)[0], -1, self.num_classes))
            if (
                not (self.training and self.no_training_normalization)
                and self.loss_function != "dirichlet"
            ):
                output = self.multiclass_softmax(output)
        if self.loss_function == "mve":
            means, variances = split(
                x=output, num_or_sections=tuple(output.shape)[1] // 2, axis=1
            )
            variances = self.softplus(variances)
            output = paddle.concat([means, variances], axis=1)

        if self.loss_function == "evidential":
            means, lambdas, alphas, betas = split(
                x=output, num_or_sections=tuple(output.shape)[1] // 4, axis=1
            )
            lambdas = self.softplus(lambdas)
            alphas = self.softplus(alphas) + 1
            betas = self.softplus(betas)
            output = paddle.concat(x=[means, lambdas, alphas, betas], axis=1)
        if self.loss_function == "dirichlet":
            output = paddle.nn.functional.softplus(x=output) + 1
        return {"pred": output}

create_encoder(args)

Create the message passing encoder for the model.

Parameters:

Name	Type	Description	Default
args	TrainArgs	Model arguments in a TrainArgs instance.	required

Source code in ppsci/arch/chemprop_molecule.py

def create_encoder(self, args: TrainArgs) -> None:
    """Create the message passing encoder for the model.

    Args:
        args: Model arguments in a `TrainArgs` instance.
    """
    self.encoder = MPN(args)
    if args.checkpoint_frzn is not None:
        if args.freeze_first_only:
            for param in list(self.encoder.encoder.children())[0].parameters():
                param.stop_gradient = not False
        else:
            for param in self.encoder.parameters():
                param.stop_gradient = not False

create_ffn(args)

Create the feed-forward layers for the model.

Parameters:

Name	Type	Description	Default
args	TrainArgs	Model arguments in a TrainArgs instance.	required

Source code in ppsci/arch/chemprop_molecule.py

def create_ffn(self, args: TrainArgs) -> None:
    """Create the feed-forward layers for the model.

    Args:
        args: Model arguments in a `TrainArgs` instance.
    """
    self.multiclass = args.dataset_type == "multiclass"
    if self.multiclass:
        self.num_classes = args.multiclass_num_classes
    if args.features_only:
        first_linear_dim = args.features_size
    else:
        if args.reaction_solvent:
            first_linear_dim = args.hidden_size + args.hidden_size_solvent
        else:
            first_linear_dim = args.hidden_size * args.number_of_molecules
        if args.use_input_features:
            first_linear_dim += args.features_size
    if args.atom_descriptors == "descriptor":
        first_linear_dim += args.atom_descriptors_size
    dropout = paddle.nn.Dropout(p=args.dropout)
    activation = get_activation_function(args.activation)
    if args.ffn_num_layers == 1:
        ffn = [
            dropout,
            paddle.nn.Linear(
                in_features=first_linear_dim, out_features=self.output_size
            ),
        ]
    else:
        ffn = [
            dropout,
            paddle.nn.Linear(
                in_features=first_linear_dim, out_features=args.ffn_hidden_size
            ),
        ]
        for _ in range(args.ffn_num_layers - 2):
            ffn.extend(
                [
                    activation,
                    dropout,
                    paddle.nn.Linear(
                        in_features=args.ffn_hidden_size,
                        out_features=args.ffn_hidden_size,
                    ),
                ]
            )
        ffn.extend(
            [
                activation,
                dropout,
                paddle.nn.Linear(
                    in_features=args.ffn_hidden_size, out_features=self.output_size
                ),
            ]
        )
    if args.dataset_type == "spectra":
        if args.spectra_activation == "softplus":
            spectra_activation = paddle.nn.Softplus()
        else:

            class nn_exp(paddle.nn.Layer):
                def __init__(self):
                    super(nn_exp, self).__init__()

                def forward(self, x):
                    return paddle.exp(x=x)

            spectra_activation = nn_exp()
        ffn.append(spectra_activation)
    self.ffn = paddle.nn.Sequential(*ffn)
    if args.checkpoint_frzn is not None:
        if args.frzn_ffn_layers > 0:
            for param in list(self.ffn.parameters())[0 : 2 * args.frzn_ffn_layers]:
                param.stop_gradient = not False

fingerprint(batch, features_batch=None, atom_descriptors_batch=None, atom_features_batch=None, bond_features_batch=None, fingerprint_type='MPN')

Return latent representations (fingerprints) from intermediate stages.

Parameters:

Name	Type	Description	Default
batch	Union[List[List[str]], List[List[Mol]], List[List[Tuple[Mol, Mol]]], List[BatchMolGraph]]	List of lists of SMILES or RDKit molecules, or a list of BatchMolGraph.	required
features_batch	List[ndarray]	Optional list of numpy arrays with additional features.	None
atom_descriptors_batch	List[ndarray]	Optional list of numpy arrays with atom descriptors.	None
atom_features_batch	List[ndarray]	Optional list of numpy arrays with atom features.	None
bond_features_batch	List[ndarray]	Optional list of numpy arrays with bond features.	None
fingerprint_type	str	Which latent representation to return. Supported: "MPN" (output of MPN) or "last_FFN" (input to the final readout layer).	'MPN'

Returns:

Type	Description
Tensor	paddle.Tensor: Fingerprint vectors.

Raises:

Type	Description
ValueError	If fingerprint_type is unsupported.

Source code in ppsci/arch/chemprop_molecule.py

def fingerprint(
    self,
    batch: Union[
        List[List[str]],
        List[List[Chem.Mol]],
        List[List[Tuple[Chem.Mol, Chem.Mol]]],
        List[BatchMolGraph],
    ],
    features_batch: List[np.ndarray] = None,
    atom_descriptors_batch: List[np.ndarray] = None,
    atom_features_batch: List[np.ndarray] = None,
    bond_features_batch: List[np.ndarray] = None,
    fingerprint_type: str = "MPN",
) -> paddle.Tensor:
    """Return latent representations (fingerprints) from intermediate stages.

    Args:
        batch: List of lists of SMILES or RDKit molecules, or a list of `BatchMolGraph`.
        features_batch: Optional list of numpy arrays with additional features.
        atom_descriptors_batch: Optional list of numpy arrays with atom descriptors.
        atom_features_batch: Optional list of numpy arrays with atom features.
        bond_features_batch: Optional list of numpy arrays with bond features.
        fingerprint_type: Which latent representation to return. Supported: "MPN" (output of MPN) or
            "last_FFN" (input to the final readout layer).

    Returns:
        `paddle.Tensor`: Fingerprint vectors.

    Raises:
        ValueError: If `fingerprint_type` is unsupported.
    """
    if fingerprint_type == "MPN":
        return self.encoder(
            batch,
            features_batch,
            atom_descriptors_batch,
            atom_features_batch,
            bond_features_batch,
        )
    elif fingerprint_type == "last_FFN":
        return self.ffn[:-1](
            self.encoder(
                batch,
                features_batch,
                atom_descriptors_batch,
                atom_features_batch,
                bond_features_batch,
            )
        )
    else:
        raise ValueError(f"Unsupported fingerprint type {fingerprint_type}.")

forward(batch, features_batch=None, atom_descriptors_batch=None, atom_features_batch=None, bond_features_batch=None)

Run the model forward pass on input.

Parameters:

Name	Type	Description	Default
batch	Union[List[List[str]], List[List[Mol]], List[List[Tuple[Mol, Mol]]], List[BatchMolGraph]]	Dictionary with keys mol_batch, features_batch, atom_descriptors_batch, atom_features_batch, bond_features_batch.	required
features_batch	List[ndarray]	Unused (kept for API symmetry), see batch.	None
atom_descriptors_batch	List[ndarray]	Unused, see batch.	None
atom_features_batch	List[ndarray]	Unused, see batch.	None
bond_features_batch	List[ndarray]	Unused, see batch.	None

Returns:

Name	Type	Description
dict	float32	{"pred": tensor} predictions (and transformed for special losses).

Source code in ppsci/arch/chemprop_molecule.py

def forward(
    self,
    batch: Union[
        List[List[str]],
        List[List[Chem.Mol]],
        List[List[Tuple[Chem.Mol, Chem.Mol]]],
        List[BatchMolGraph],
    ],
    features_batch: List[np.ndarray] = None,
    atom_descriptors_batch: List[np.ndarray] = None,
    atom_features_batch: List[np.ndarray] = None,
    bond_features_batch: List[np.ndarray] = None,
) -> paddle.float32:
    """Run the model forward pass on input.

    Args:
        batch: Dictionary with keys `mol_batch`, `features_batch`,
            `atom_descriptors_batch`, `atom_features_batch`, `bond_features_batch`.
        features_batch: Unused (kept for API symmetry), see `batch`.
        atom_descriptors_batch: Unused, see `batch`.
        atom_features_batch: Unused, see `batch`.
        bond_features_batch: Unused, see `batch`.

    Returns:
        dict: `{"pred": tensor}` predictions (and transformed for special losses).
    """

    mol_batch = batch["mol_batch"]
    features_batch = batch["features_batch"]
    atom_descriptors_batch = batch["atom_descriptors_batch"]
    atom_features_batch = batch["atom_features_batch"]
    bond_features_batch = batch["bond_features_batch"]
    batch = mol_batch

    output = self.ffn(
        self.encoder(
            batch,
            features_batch,
            atom_descriptors_batch,
            atom_features_batch,
            bond_features_batch,
        )
    )
    if (
        self.classification
        and not (self.training and self.no_training_normalization)
        and self.loss_function != "dirichlet"
    ):
        output = self.sigmoid(output)
    if self.multiclass:
        output = output.reshape((tuple(output.shape)[0], -1, self.num_classes))
        if (
            not (self.training and self.no_training_normalization)
            and self.loss_function != "dirichlet"
        ):
            output = self.multiclass_softmax(output)
    if self.loss_function == "mve":
        means, variances = split(
            x=output, num_or_sections=tuple(output.shape)[1] // 2, axis=1
        )
        variances = self.softplus(variances)
        output = paddle.concat([means, variances], axis=1)

    if self.loss_function == "evidential":
        means, lambdas, alphas, betas = split(
            x=output, num_or_sections=tuple(output.shape)[1] // 4, axis=1
        )
        lambdas = self.softplus(lambdas)
        alphas = self.softplus(alphas) + 1
        betas = self.softplus(betas)
        output = paddle.concat(x=[means, lambdas, alphas, betas], axis=1)
    if self.loss_function == "dirichlet":
        output = paddle.nn.functional.softplus(x=output) + 1
    return {"pred": output}

NowcastNet

Bases: Arch

The NowcastNet model.

Parameters:

Name	Type	Description	Default
input_keys	Tuple[str, ...]	Name of input keys, such as ("input",).	required
output_keys	Tuple[str, ...]	Name of output keys, such as ("output",).	required
input_length	int	Input length. Defaults to 9.	9
total_length	int	Total length. Defaults to 29.	29
image_height	int	Image height. Defaults to 512.	512
image_width	int	Image width. Defaults to 512.	512
image_ch	int	Image channel. Defaults to 2.	2
ngf	int	Noise Projector input length. Defaults to 32.	32

Examples:

>>> import ppsci
>>> model = ppsci.arch.NowcastNet(("input", ), ("output", ))
>>> input_data = paddle.rand([1, 9, 512, 512, 2])
>>> input_dict = {"input": input_data}
>>> output_dict = model(input_dict)
>>> print(output_dict["output"].shape)
[1, 20, 512, 512, 1]

Source code in ppsci/arch/nowcastnet.py

class NowcastNet(base.Arch):
    """The NowcastNet model.

    Args:
        input_keys (Tuple[str, ...]): Name of input keys, such as ("input",).
        output_keys (Tuple[str, ...]): Name of output keys, such as ("output",).
        input_length (int, optional): Input length. Defaults to 9.
        total_length (int, optional): Total length. Defaults to 29.
        image_height (int, optional): Image height. Defaults to 512.
        image_width (int, optional): Image width. Defaults to 512.
        image_ch (int, optional): Image channel. Defaults to 2.
        ngf (int, optional): Noise Projector input length. Defaults to 32.

    Examples:
        >>> import ppsci
        >>> model = ppsci.arch.NowcastNet(("input", ), ("output", ))
        >>> input_data = paddle.rand([1, 9, 512, 512, 2])
        >>> input_dict = {"input": input_data}
        >>> output_dict = model(input_dict)
        >>> print(output_dict["output"].shape)
        [1, 20, 512, 512, 1]
    """

    def __init__(
        self,
        input_keys: Tuple[str, ...],
        output_keys: Tuple[str, ...],
        input_length: int = 9,
        total_length: int = 29,
        image_height: int = 512,
        image_width: int = 512,
        image_ch: int = 2,
        ngf: int = 32,
    ):
        super().__init__()
        self.input_keys = input_keys
        self.output_keys = output_keys

        self.input_length = input_length
        self.total_length = total_length
        self.image_height = image_height
        self.image_width = image_width
        self.image_ch = image_ch
        self.ngf = ngf

        configs = collections.namedtuple(
            "Object", ["ngf", "evo_ic", "gen_oc", "ic_feature"]
        )
        configs.ngf = self.ngf
        configs.evo_ic = self.total_length - self.input_length
        configs.gen_oc = self.total_length - self.input_length
        configs.ic_feature = self.ngf * 10

        self.pred_length = self.total_length - self.input_length
        self.evo_net = Evolution_Network(self.input_length, self.pred_length, base_c=32)
        self.gen_enc = Generative_Encoder(self.total_length, base_c=self.ngf)
        self.gen_dec = Generative_Decoder(configs)
        self.proj = Noise_Projector(self.ngf)
        sample_tensor = paddle.zeros(shape=[1, 1, self.image_height, self.image_width])
        self.grid = make_grid(sample_tensor)

    @staticmethod
    def split_to_dict(data_tensors: Tuple[paddle.Tensor, ...], keys: Tuple[str, ...]):
        return {key: data_tensors[i] for i, key in enumerate(keys)}

    def forward(self, x):
        if self._input_transform is not None:
            x = self._input_transform(x)

        x_tensor = self.concat_to_tensor(x, self.input_keys)

        y = []
        out = self.forward_tensor(x_tensor)
        y.append(out)
        y = self.split_to_dict(y, self.output_keys)

        if self._output_transform is not None:
            y = self._output_transform(x, y)
        return y

    def forward_tensor(self, x):
        all_frames = x[:, :, :, :, :1]
        frames = all_frames.transpose(perm=[0, 1, 4, 2, 3])
        batch = frames.shape[0]
        height = frames.shape[3]
        width = frames.shape[4]
        # Input Frames
        input_frames = frames[:, : self.input_length]
        input_frames = input_frames.reshape((batch, self.input_length, height, width))
        # Evolution Network
        intensity, motion = self.evo_net(input_frames)
        motion_ = motion.reshape((batch, self.pred_length, 2, height, width))
        intensity_ = intensity.reshape((batch, self.pred_length, 1, height, width))
        series = []
        last_frames = all_frames[:, self.input_length - 1 : self.input_length, :, :, 0]
        grid = self.grid.tile((batch, 1, 1, 1))
        for i in range(self.pred_length):
            last_frames = warp(
                last_frames, motion_[:, i], grid, mode="nearest", padding_mode="border"
            )
            last_frames = last_frames + intensity_[:, i]
            series.append(last_frames)
        evo_result = paddle.concat(x=series, axis=1)
        evo_result = evo_result / 128
        # Generative Network
        evo_feature = self.gen_enc(paddle.concat(x=[input_frames, evo_result], axis=1))
        noise = paddle.randn(shape=[batch, self.ngf, height // 32, width // 32])
        noise = self.proj(noise)
        ngf = noise.shape[1]
        noise_feature = (
            noise.reshape((batch, -1, 4, 4, 8, 8))
            .transpose(perm=[0, 1, 4, 5, 2, 3])
            .reshape((batch, ngf // 16, height // 8, width // 8))
        )
        feature = paddle.concat(x=[evo_feature, noise_feature], axis=1)
        gen_result = self.gen_dec(feature, evo_result)
        return gen_result.unsqueeze(axis=-1)

Preformer

Bases: Arch

Preformer is a class that represents a Spatial-Temporal Transformer model designed for short-term precipitation forecasting with multiple meteorological variables.

Parameters:

Name	Type	Description	Default
input_keys	Tuple[str, ...]	A tuple of input keys.	required
output_keys	Tuple[str, ...]	A tuple of output keys.	required
shape_in	Tuple[int, ...]	The shape of the input data (T, C, H, W), where T is the number of time steps, C is the number of channels, H and W are the spatial dimensions.	required
hid_S	int	The number of hidden channels in the spatial encoder.	64
hid_T	int	The number of hidden units in the temporal encoder.	256
N_S	int	The number of spatial transformer layers.	4
N_T	int	The number of temporal transformer layers.	8
incep_ker	Tuple[int, ...]	The kernel sizes used in the inception block.	(3, 5, 7, 11)
groups	int	The number of groups for grouped convolutions.	8
num_classes	int	The number of predicted meteorological variables.	1

Examples:

>>> import paddle
>>> import ppsci
>>> model = ppsci.arch.Preformer(
...     input_keys=("input",),
...     output_keys=("output",),
...     shape_in=(6, 12, 192, 256),
...     hid_S=64,
...     hid_T=256,
...     N_S=4,
...     N_T=4,
...     incep_ker=(3, 5, 7, 11),
...     groups=8,
...     num_classes=4,
... )
>>> input_dict = {"input": paddle.rand([8, 6, 12, 192, 256])}
>>> output_dict = model(input_dict)
>>> print(output_dict["output"].shape)
[8, 6, 1, 192, 256]

Source code in ppsci/arch/preformer.py

class Preformer(base.Arch):
    """
    Preformer is a class that represents a Spatial-Temporal Transformer model designed for short-term precipitation forecasting with multiple meteorological variables.

    Args:
        input_keys (Tuple[str, ...]): A tuple of input keys.
        output_keys (Tuple[str, ...]): A tuple of output keys.
        shape_in (Tuple[int, ...]): The shape of the input data (T, C, H, W), where
            T is the number of time steps, C is the number of channels,
            H and W are the spatial dimensions.
        hid_S (int): The number of hidden channels in the spatial encoder.
        hid_T (int): The number of hidden units in the temporal encoder.
        N_S (int): The number of spatial transformer layers.
        N_T (int): The number of temporal transformer layers.
        incep_ker (Tuple[int, ...]): The kernel sizes used in the inception block.
        groups (int): The number of groups for grouped convolutions.
        num_classes (int): The number of predicted meteorological variables.

    Examples:
        >>> import paddle
        >>> import ppsci
        >>> model = ppsci.arch.Preformer(
        ...     input_keys=("input",),
        ...     output_keys=("output",),
        ...     shape_in=(6, 12, 192, 256),
        ...     hid_S=64,
        ...     hid_T=256,
        ...     N_S=4,
        ...     N_T=4,
        ...     incep_ker=(3, 5, 7, 11),
        ...     groups=8,
        ...     num_classes=4,
        ... )
        >>> input_dict = {"input": paddle.rand([8, 6, 12, 192, 256])}
        >>> output_dict = model(input_dict)
        >>> print(output_dict["output"].shape)
        [8, 6, 1, 192, 256]
    """

    def __init__(
        self,
        input_keys: Tuple[str, ...],
        output_keys: Tuple[str, ...],
        shape_in: Tuple[int, ...],
        hid_S: int = 64,
        hid_T: int = 256,
        N_S: int = 4,
        N_T: int = 8,
        incep_ker: Tuple[int, ...] = (3, 5, 7, 11),
        groups: int = 8,
        num_classes: int = 1,
    ):
        super().__init__()

        self.input_keys = input_keys
        self.output_keys = output_keys

        T, C, H, W = shape_in
        self.enc = Encoder(C, hid_S, N_S)
        self.hid1 = MidXnet(T * hid_S, hid_T // 2, N_T, incep_ker, groups)
        self.dec = Decoder(T * hid_S, T * num_classes, N_S)

    def forward(self, x_raw):
        x_raw = x_raw[self.input_keys[0]]

        B, T, C, H, W = x_raw.shape
        x = x_raw.reshape([B * T, C, H, W])

        # encoded
        embed = self.enc(x)
        _, C_4, H_4, W_4 = embed[-1].shape

        # translator
        z = embed[-1].reshape([B, T, C_4, H_4, W_4])
        hid = self.hid1(z)
        hid = hid.transpose(perm=[0, 2, 1]).reshape([B, -1, H_4, W_4])

        # decoded
        Y = self.dec(hid, embed[0])
        Y = Y.reshape([B, T, 1, H, W])

        Y = nn.functional.softplus(Y)

        return {self.output_keys[0]: Y}

RegDGCNN

Bases: Layer

Deep Graph Convolutional Neural Network for Regression Tasks (RegDGCNN) designed to process 3D point cloud data.

This network architecture extracts hierarchical features from point clouds using graph-based convolutions, enabling effective learning of spatial structures.

Parameters:

Name	Type	Description	Default
input_keys	Tuple[str, ...]	Keys for input data fields.	required
label_keys	Tuple[str, ...]	Keys for label data fields.	required
weight_keys	Tuple[str, ...]	Keys for weight data fields.	required
args	dict	Configuration parameters including: - 'k' (int): Number of neighbors for graph convolution. - 'emb_dims' (int): Embedding dimensions for feature aggregation. - 'dropout' (float): Dropout rate for regularization.	required
output_channels	int	Number of output channels. Defaults to 1.	1

Source code in ppsci/arch/regdgcnn.py

class RegDGCNN(paddle.nn.Layer):
    """Deep Graph Convolutional Neural Network for Regression Tasks (RegDGCNN) designed to process 3D point cloud data.

    This network architecture extracts hierarchical features from point clouds using graph-based convolutions,
    enabling effective learning of spatial structures.

    Args:
        input_keys (Tuple[str, ...]): Keys for input data fields.
        label_keys (Tuple[str, ...]): Keys for label data fields.
        weight_keys (Tuple[str, ...]): Keys for weight data fields.
        args (dict): Configuration parameters including:
            - 'k' (int): Number of neighbors for graph convolution.
            - 'emb_dims' (int): Embedding dimensions for feature aggregation.
            - 'dropout' (float): Dropout rate for regularization.
        output_channels (int, optional): Number of output channels. Defaults to 1.
    """

    def __init__(
        self,
        input_keys: Tuple[str, ...],
        label_keys: Tuple[str, ...],
        weight_keys: Tuple[str, ...],
        args: dict,
        output_channels=1,
    ):

        super(RegDGCNN, self).__init__()
        self.input_keys = input_keys
        self.label_keys = label_keys
        self.weight_keys = weight_keys
        self.args = args
        self.k = args["k"]
        self.bn1 = paddle.nn.BatchNorm2D(num_features=256)
        self.bn2 = paddle.nn.BatchNorm2D(num_features=512)
        self.bn3 = paddle.nn.BatchNorm2D(num_features=512)
        self.bn4 = paddle.nn.BatchNorm2D(num_features=1024)
        self.bn5 = paddle.nn.BatchNorm1D(num_features=args["emb_dims"])
        self.conv1 = paddle.nn.Sequential(
            paddle.nn.Conv2D(
                in_channels=6, out_channels=256, kernel_size=1, bias_attr=False
            ),
            self.bn1,
            paddle.nn.LeakyReLU(negative_slope=0.2),
        )
        self.conv2 = paddle.nn.Sequential(
            paddle.nn.Conv2D(
                in_channels=256 * 2, out_channels=512, kernel_size=1, bias_attr=False
            ),
            self.bn2,
            paddle.nn.LeakyReLU(negative_slope=0.2),
        )
        self.conv3 = paddle.nn.Sequential(
            paddle.nn.Conv2D(
                in_channels=512 * 2, out_channels=512, kernel_size=1, bias_attr=False
            ),
            self.bn3,
            paddle.nn.LeakyReLU(negative_slope=0.2),
        )
        self.conv4 = paddle.nn.Sequential(
            paddle.nn.Conv2D(
                in_channels=512 * 2, out_channels=1024, kernel_size=1, bias_attr=False
            ),
            self.bn4,
            paddle.nn.LeakyReLU(negative_slope=0.2),
        )
        self.conv5 = paddle.nn.Sequential(
            paddle.nn.Conv1D(
                in_channels=2304,
                out_channels=args["emb_dims"],
                kernel_size=1,
                bias_attr=False,
            ),
            self.bn5,
            paddle.nn.LeakyReLU(negative_slope=0.2),
        )
        self.linear1 = paddle.nn.Linear(
            in_features=args["emb_dims"] * 2, out_features=128, bias_attr=False
        )
        self.bn6 = paddle.nn.BatchNorm1D(num_features=128)
        self.dp1 = paddle.nn.Dropout(p=args["dropout"])
        self.linear2 = paddle.nn.Linear(in_features=128, out_features=64)
        self.bn7 = paddle.nn.BatchNorm1D(num_features=64)
        self.dp2 = paddle.nn.Dropout(p=args["dropout"])
        self.linear3 = paddle.nn.Linear(in_features=64, out_features=32)
        self.bn8 = paddle.nn.BatchNorm1D(num_features=32)
        self.dp3 = paddle.nn.Dropout(p=args["dropout"])
        self.linear4 = paddle.nn.Linear(in_features=32, out_features=16)
        self.bn9 = paddle.nn.BatchNorm1D(num_features=16)
        self.dp4 = paddle.nn.Dropout(p=args["dropout"])
        self.linear5 = paddle.nn.Linear(in_features=16, out_features=output_channels)

    def forward(self, x: paddle.Tensor) -> Dict[str, paddle.Tensor]:
        """
        Forward pass of the model to process input data and predict outputs.

        Args:
            x (paddle.Tensor): Input tensor representing a batch of point clouds.

        Returns:
            Dict[str, paddle.Tensor]: Model predictions for the input batch.

        """

        x = x[self.input_keys[0]]
        batch_size = x.shape[0]
        x = x.transpose(perm=[0, 2, 1])

        x = get_graph_feature(x, k=self.k)
        x = self.conv1(x)
        x1 = x.max(axis=-1, keepdim=False)
        x = get_graph_feature(x1, k=self.k)
        x = self.conv2(x)
        x2 = x.max(axis=-1, keepdim=False)
        x = get_graph_feature(x2, k=self.k)
        x = self.conv3(x)
        x3 = x.max(axis=-1, keepdim=False)
        x = get_graph_feature(x3, k=self.k)
        x = self.conv4(x)
        x4 = x.max(axis=-1, keepdim=False)
        x = paddle.concat(x=(x1, x2, x3, x4), axis=1)
        x = self.conv5(x)
        x1 = paddle.nn.functional.adaptive_max_pool1d(x=x, output_size=1).reshape(
            [batch_size, -1]
        )
        x2 = paddle.nn.functional.adaptive_avg_pool1d(x=x, output_size=1).reshape(
            [batch_size, -1]
        )
        x = paddle.concat(x=(x1, x2), axis=1)
        x = paddle.nn.functional.leaky_relu(
            x=self.bn6(self.linear1(x)), negative_slope=0.2
        )
        x = self.dp1(x)
        x = paddle.nn.functional.leaky_relu(
            x=self.bn7(self.linear2(x)), negative_slope=0.2
        )
        x = self.dp2(x)
        x = paddle.nn.functional.leaky_relu(
            x=self.bn8(self.linear3(x)), negative_slope=0.2
        )
        x = self.dp3(x)
        x = paddle.nn.functional.leaky_relu(
            x=self.bn9(self.linear4(x)), negative_slope=0.2
        )
        x = self.dp4(x)
        x = self.linear5(x)
        return {self.label_keys[0]: x}

forward(x)

Forward pass of the model to process input data and predict outputs.

Parameters:

Name	Type	Description	Default
x	Tensor	Input tensor representing a batch of point clouds.	required

Returns:

Type	Description
Dict[str, Tensor]	Dict[str, paddle.Tensor]: Model predictions for the input batch.

Source code in ppsci/arch/regdgcnn.py

def forward(self, x: paddle.Tensor) -> Dict[str, paddle.Tensor]:
    """
    Forward pass of the model to process input data and predict outputs.

    Args:
        x (paddle.Tensor): Input tensor representing a batch of point clouds.

    Returns:
        Dict[str, paddle.Tensor]: Model predictions for the input batch.

    """

    x = x[self.input_keys[0]]
    batch_size = x.shape[0]
    x = x.transpose(perm=[0, 2, 1])

    x = get_graph_feature(x, k=self.k)
    x = self.conv1(x)
    x1 = x.max(axis=-1, keepdim=False)
    x = get_graph_feature(x1, k=self.k)
    x = self.conv2(x)
    x2 = x.max(axis=-1, keepdim=False)
    x = get_graph_feature(x2, k=self.k)
    x = self.conv3(x)
    x3 = x.max(axis=-1, keepdim=False)
    x = get_graph_feature(x3, k=self.k)
    x = self.conv4(x)
    x4 = x.max(axis=-1, keepdim=False)
    x = paddle.concat(x=(x1, x2, x3, x4), axis=1)
    x = self.conv5(x)
    x1 = paddle.nn.functional.adaptive_max_pool1d(x=x, output_size=1).reshape(
        [batch_size, -1]
    )
    x2 = paddle.nn.functional.adaptive_avg_pool1d(x=x, output_size=1).reshape(
        [batch_size, -1]
    )
    x = paddle.concat(x=(x1, x2), axis=1)
    x = paddle.nn.functional.leaky_relu(
        x=self.bn6(self.linear1(x)), negative_slope=0.2
    )
    x = self.dp1(x)
    x = paddle.nn.functional.leaky_relu(
        x=self.bn7(self.linear2(x)), negative_slope=0.2
    )
    x = self.dp2(x)
    x = paddle.nn.functional.leaky_relu(
        x=self.bn8(self.linear3(x)), negative_slope=0.2
    )
    x = self.dp3(x)
    x = paddle.nn.functional.leaky_relu(
        x=self.bn9(self.linear4(x)), negative_slope=0.2
    )
    x = self.dp4(x)
    x = self.linear5(x)
    return {self.label_keys[0]: x}

RegPointNet

Bases: Layer

PointNet-based regression model for 3D point cloud data.

This network architecture is designed to process 3D point cloud data using a series of convolutional layers, followed by fully connected layers, enabling effective learning of spatial structures and features.

Parameters:

Name	Type	Description	Default
input_keys	Tuple[str, ...]	Keys for input data fields.	required
output_keys	Tuple[str, ...]	Keys for output data fields.	required
weight_keys	Tuple[str, ...]	Keys for weight data fields.	required
args	dict	Configuration parameters including: - 'emb_dims' (int): Dimensionality of the embedding space. - 'dropout' (float): Dropout probability.	required

Source code in ppsci/arch/regpointnet.py

class RegPointNet(paddle.nn.Layer):
    """
    PointNet-based regression model for 3D point cloud data.

    This network architecture is designed to process 3D point cloud data using a series of convolutional layers,
    followed by fully connected layers, enabling effective learning of spatial structures and features.

    Args:
        input_keys (Tuple[str, ...]): Keys for input data fields.
        output_keys (Tuple[str, ...]): Keys for output data fields.
        weight_keys (Tuple[str, ...]): Keys for weight data fields.
        args (dict): Configuration parameters including:
            - 'emb_dims' (int): Dimensionality of the embedding space.
            - 'dropout' (float): Dropout probability.
    """

    def __init__(
        self,
        input_keys: Tuple[str, ...],
        output_keys: Tuple[str, ...],
        weight_keys: Tuple[str, ...],
        args,
    ):
        super().__init__()
        self.input_keys = input_keys
        self.output_keys = output_keys
        self.weight_keys = weight_keys
        self.args = args
        self.conv1 = paddle.nn.Conv1D(
            in_channels=3, out_channels=512, kernel_size=1, bias_attr=False
        )
        self.conv2 = paddle.nn.Conv1D(
            in_channels=512, out_channels=1024, kernel_size=1, bias_attr=False
        )
        self.conv3 = paddle.nn.Conv1D(
            in_channels=1024, out_channels=1024, kernel_size=1, bias_attr=False
        )
        self.conv4 = paddle.nn.Conv1D(
            in_channels=1024, out_channels=1024, kernel_size=1, bias_attr=False
        )
        self.conv5 = paddle.nn.Conv1D(
            in_channels=1024, out_channels=1024, kernel_size=1, bias_attr=False
        )
        self.conv6 = paddle.nn.Conv1D(
            in_channels=1024,
            out_channels=args["emb_dims"],
            kernel_size=1,
            bias_attr=False,
        )
        self.bn1 = paddle.nn.BatchNorm1D(num_features=512)
        self.bn2 = paddle.nn.BatchNorm1D(num_features=1024)
        self.bn3 = paddle.nn.BatchNorm1D(num_features=1024)
        self.bn4 = paddle.nn.BatchNorm1D(num_features=1024)
        self.bn5 = paddle.nn.BatchNorm1D(num_features=1024)
        self.bn6 = paddle.nn.BatchNorm1D(num_features=args["emb_dims"])
        self.dropout_conv = paddle.nn.Dropout(p=args["dropout"])
        self.dropout_linear = paddle.nn.Dropout(p=args["dropout"])
        self.conv_shortcut = paddle.nn.Conv1D(
            in_channels=3, out_channels=args["emb_dims"], kernel_size=1, bias_attr=False
        )
        self.bn_shortcut = paddle.nn.BatchNorm1D(num_features=args["emb_dims"])
        self.linear1 = paddle.nn.Linear(
            in_features=args["emb_dims"], out_features=512, bias_attr=False
        )
        self.bn7 = paddle.nn.BatchNorm1D(num_features=512)
        self.linear2 = paddle.nn.Linear(
            in_features=512, out_features=256, bias_attr=False
        )
        self.bn8 = paddle.nn.BatchNorm1D(num_features=256)
        self.linear3 = paddle.nn.Linear(in_features=256, out_features=128)
        self.bn9 = paddle.nn.BatchNorm1D(num_features=128)
        self.linear4 = paddle.nn.Linear(in_features=128, out_features=64)
        self.bn10 = paddle.nn.BatchNorm1D(num_features=64)
        self.final_linear = paddle.nn.Linear(in_features=64, out_features=1)

    def forward(self, x: Dict[str, paddle.Tensor]) -> Dict[str, paddle.Tensor]:
        """
        Forward pass of the network.

        Args:
            x (Dict[str, paddle.Tensor]): Input tensor of shape (batch_size, 3, num_points).

        Returns:
            Dict[str, paddle.Tensor]: A dictionary where the key is the first element of `self.output_keys`
                                       and the value is the output tensor of the predicted scalar value.
        """

        x: paddle.Tensor = x[self.input_keys[0]]

        x_processed = x.transpose(perm=[0, 2, 1])

        shortcut = self.bn_shortcut(self.conv_shortcut(x_processed))
        x = paddle.nn.functional.relu(x=self.bn1(self.conv1(x_processed)))
        x = self.dropout_conv(x)
        x = paddle.nn.functional.relu(x=self.bn2(self.conv2(x)))
        x = self.dropout_conv(x)
        x = paddle.nn.functional.relu(x=self.bn3(self.conv3(x)))
        x = self.dropout_conv(x)
        x = paddle.nn.functional.relu(x=self.bn4(self.conv4(x)))
        x = self.dropout_conv(x)
        x = paddle.nn.functional.relu(x=self.bn5(self.conv5(x)))
        x = self.dropout_conv(x)
        x = paddle.nn.functional.relu(x=self.bn6(self.conv6(x)))
        x = x + shortcut
        x = paddle.nn.functional.adaptive_max_pool1d(x=x, output_size=1).squeeze(
            axis=-1
        )
        x = paddle.nn.functional.relu(x=self.bn7(self.linear1(x)))
        x = paddle.nn.functional.relu(x=self.bn8(self.linear2(x)))
        x = paddle.nn.functional.relu(x=self.bn9(self.linear3(x)))
        x = paddle.nn.functional.relu(x=self.bn10(self.linear4(x)))
        x = self.final_linear(x)
        return {self.output_keys[0]: x}

forward(x)

Forward pass of the network.

Parameters:

Name	Type	Description	Default
x	Dict[str, Tensor]	Input tensor of shape (batch_size, 3, num_points).	required

Returns:

Type	Description
Dict[str, Tensor]	Dict[str, paddle.Tensor]: A dictionary where the key is the first element of self.output_keys and the value is the output tensor of the predicted scalar value.

Source code in ppsci/arch/regpointnet.py

def forward(self, x: Dict[str, paddle.Tensor]) -> Dict[str, paddle.Tensor]:
    """
    Forward pass of the network.

    Args:
        x (Dict[str, paddle.Tensor]): Input tensor of shape (batch_size, 3, num_points).

    Returns:
        Dict[str, paddle.Tensor]: A dictionary where the key is the first element of `self.output_keys`
                                   and the value is the output tensor of the predicted scalar value.
    """

    x: paddle.Tensor = x[self.input_keys[0]]

    x_processed = x.transpose(perm=[0, 2, 1])

    shortcut = self.bn_shortcut(self.conv_shortcut(x_processed))
    x = paddle.nn.functional.relu(x=self.bn1(self.conv1(x_processed)))
    x = self.dropout_conv(x)
    x = paddle.nn.functional.relu(x=self.bn2(self.conv2(x)))
    x = self.dropout_conv(x)
    x = paddle.nn.functional.relu(x=self.bn3(self.conv3(x)))
    x = self.dropout_conv(x)
    x = paddle.nn.functional.relu(x=self.bn4(self.conv4(x)))
    x = self.dropout_conv(x)
    x = paddle.nn.functional.relu(x=self.bn5(self.conv5(x)))
    x = self.dropout_conv(x)
    x = paddle.nn.functional.relu(x=self.bn6(self.conv6(x)))
    x = x + shortcut
    x = paddle.nn.functional.adaptive_max_pool1d(x=x, output_size=1).squeeze(
        axis=-1
    )
    x = paddle.nn.functional.relu(x=self.bn7(self.linear1(x)))
    x = paddle.nn.functional.relu(x=self.bn8(self.linear2(x)))
    x = paddle.nn.functional.relu(x=self.bn9(self.linear3(x)))
    x = paddle.nn.functional.relu(x=self.bn10(self.linear4(x)))
    x = self.final_linear(x)
    return {self.output_keys[0]: x}

SFNONet

Bases: Arch

N-Dimensional Tensorized Fourier Neural Operator.

Parameters:

Name	Type	Description	Default
input_keys	Tuple[str, ...]	Name of input keys, such as ("input",).	required
output_keys	Tuple[str, ...]	Name of output keys, such as ("output",).	required
n_modes	Tuple[int, ...]	Number of modes to keep in Fourier Layer, along each dimension The dimensionality of the SFNO is inferred from `len(n_modes)	required
hidden_channels	int	Width of the FNO (i.e. number of channels)	required
in_channels	int	Number of input channels. Defaults to 3.	3
out_channels	int	Number of output channels. Defaults to 1.	1
lifting_channels	int	Number of hidden channels of the lifting block of the FNO. Defaults to 256.	256
projection_channels	int	Number of hidden channels of the projection block of the FNO. Defaults to 256.	256
n_layers	int	Number of Fourier Layers. Defaults to 4.	4
use_mlp	bool	Whether to use an MLP layer after each FNO block. Defaults to False.	False
mlp	Dict[str, float]	Parameters of the MLP. {'expansion': float, 'dropout': float}. Defaults to None.	None
non_linearity	functional	Non-Linearity module to use. Defaults to F.gelu.	gelu
norm	str	Normalization layer to use. Defaults to None.	None
ada_in_features	(int, optional)	The input channels of the adaptive normalization.Defaults to None.	None
preactivation	bool	Whether to use resnet-style preactivation. Defaults to False.	False
fno_skip	str	Type of skip connection to use,{'linear', 'identity', 'soft-gating'}. Defaults to "soft-gating".	'linear'
separable	bool	Whether to use a depthwise separable spectral convolution. Defaults to False.	False
factorization	str	Tensor factorization of the parameters weight to use. * If None, a dense tensor parametrizes the Spectral convolutions. * Otherwise, the specified tensor factorization is used. Defaults to "Tucker".	None
rank	float	Rank of the tensor factorization of the Fourier weights. Defaults to 1.0.	1.0
joint_factorization	bool	Whether all the Fourier Layers should be parametrized by a single tensor (vs one per layer). Defaults to False.	False
implementation	str	{'factorized', 'reconstructed'}, optional. Defaults to "factorized". If factorization is not None, forward mode to use:: * reconstructed : the full weight tensor is reconstructed from the factorization and used for the forward pass. * factorized : the input is directly contracted with the factors of the decomposition.	'factorized'
domain_padding	Optional[list]	Whether to use percentage of padding. Defaults to None.	None
domain_padding_mode	str	{'symmetric', 'one-sided'}, optional How to perform domain padding, by default 'one-sided'. Defaults to "one-sided".	'one-sided'
fft_norm	str	The normalization mode for the FFT. Defaults to "forward".	'forward'
patching_levels	int	Number of patching levels to use. Defaults to 0.	0

Source code in ppsci/arch/sfnonet.py

class SFNONet(base.Arch):
    """N-Dimensional Tensorized Fourier Neural Operator.

    Args:
        input_keys (Tuple[str, ...]): Name of input keys, such as ("input",).
        output_keys (Tuple[str, ...]): Name of output keys, such as ("output",).
        n_modes (Tuple[int, ...]): Number of modes to keep in Fourier Layer, along each dimension
            The dimensionality of the SFNO is inferred from ``len(n_modes)`
        hidden_channels (int): Width of the FNO (i.e. number of channels)
        in_channels (int, optional): Number of input channels. Defaults to 3.
        out_channels (int, optional): Number of output channels. Defaults to 1.
        lifting_channels (int, optional): Number of hidden channels of the lifting block of the FNO.
            Defaults to 256.
        projection_channels (int, optional): Number of hidden channels of the projection block of the FNO.
            Defaults to 256.
        n_layers (int, optional): Number of Fourier Layers. Defaults to 4.
        use_mlp (bool, optional): Whether to use an MLP layer after each FNO block. Defaults to False.
        mlp (Dict[str, float], optional): Parameters of the MLP. {'expansion': float, 'dropout': float}.
            Defaults to None.
        non_linearity (nn.functional, optional): Non-Linearity module to use. Defaults to F.gelu.
        norm (str, optional): Normalization layer to use. Defaults to None.
        ada_in_features (int,optional): The input channels of the adaptive normalization.Defaults to None.
        preactivation (bool, optional): Whether to use resnet-style preactivation. Defaults to False.
        fno_skip (str, optional): Type of skip connection to use,{'linear', 'identity', 'soft-gating'}.
            Defaults to "soft-gating".
        separable (bool, optional): Whether to use a depthwise separable spectral convolution.
            Defaults to  False.
        factorization (str, optional): Tensor factorization of the parameters weight to use.
            * If None, a dense tensor parametrizes the Spectral convolutions.
            * Otherwise, the specified tensor factorization is used. Defaults to "Tucker".
        rank (float, optional): Rank of the tensor factorization of the Fourier weights. Defaults to 1.0.
        joint_factorization (bool, optional): Whether all the Fourier Layers should be parametrized by a
            single tensor (vs one per layer). Defaults to False.
        implementation (str, optional): {'factorized', 'reconstructed'}, optional. Defaults to "factorized".
            If factorization is not None, forward mode to use::
            * `reconstructed` : the full weight tensor is reconstructed from the factorization and used for the forward pass.
            * `factorized` : the input is directly contracted with the factors of the decomposition.
        domain_padding (Optional[list], optional): Whether to use percentage of padding. Defaults to None.
        domain_padding_mode (str, optional): {'symmetric', 'one-sided'}, optional
            How to perform domain padding, by default 'one-sided'. Defaults to "one-sided".
        fft_norm (str, optional): The normalization mode for the FFT. Defaults to "forward".
        patching_levels (int, optional): Number of patching levels to use. Defaults to 0.
    """

    def __init__(
        self,
        input_keys: Tuple[str, ...],
        output_keys: Tuple[str, ...],
        n_modes: Tuple[int, ...],
        hidden_channels: int,
        in_channels: int = 3,
        out_channels: int = 1,
        lifting_channels: int = 256,
        projection_channels: int = 256,
        n_layers: int = 4,
        use_mlp: bool = False,
        mlp: Optional[Dict[str, float]] = None,
        max_n_modes: int = None,
        non_linearity: nn.functional = F.gelu,
        stabilizer: str = None,
        norm: str = None,
        ada_in_features: Optional[int] = None,
        preactivation: bool = False,
        fno_skip: str = "linear",
        mlp_skip: str = "soft-gating",
        separable: bool = False,
        factorization: str = None,
        rank: float = 1.0,
        joint_factorization: bool = False,
        implementation: str = "factorized",
        domain_padding: Optional[list] = None,
        domain_padding_mode: str = "one-sided",
        fft_norm: str = "forward",
        patching_levels: int = 0,
        **kwargs,
    ):
        super().__init__()
        self.input_keys = input_keys
        self.output_keys = output_keys

        self.n_dim = len(n_modes)
        self.n_modes = n_modes
        self.hidden_channels = hidden_channels
        self.lifting_channels = lifting_channels
        self.projection_channels = projection_channels
        self.in_channels = in_channels
        if patching_levels:
            self.in_channels = self.in_channels * patching_levels + 1
        self.out_channels = out_channels
        self.n_layers = n_layers
        self.joint_factorization = joint_factorization
        self.non_linearity = non_linearity
        self.rank = rank
        self.factorization = factorization
        self.fno_skip = (fno_skip,)
        self.mlp_skip = (mlp_skip,)
        self.fft_norm = fft_norm
        self.implementation = implementation
        self.separable = separable
        self.preactivation = preactivation
        self.stabilizer = stabilizer
        if domain_padding is not None and (
            (isinstance(domain_padding, list) and sum(domain_padding) > 0)
            or (isinstance(domain_padding, (float, int)) and domain_padding > 0)
        ):
            self.domain_padding = fno_block.DomainPadding(
                domain_padding=domain_padding, padding_mode=domain_padding_mode
            )
        else:
            self.domain_padding = None
        self.domain_padding_mode = domain_padding_mode

        self.fno_blocks = fno_block.FNOBlocks(
            in_channels=hidden_channels,
            out_channels=hidden_channels,
            n_modes=self.n_modes,
            n_layers=n_layers,
            max_n_modes=max_n_modes,
            use_mlp=use_mlp,
            mlp=mlp,
            non_linearity=non_linearity,
            stabilizer=stabilizer,
            norm=norm,
            ada_in_features=ada_in_features,
            preactivation=preactivation,
            fno_skip=fno_skip,
            mlp_skip=mlp_skip,
            separable=separable,
            factorization=factorization,
            rank=rank,
            SpectralConv=SphericalConv,
            joint_factorization=joint_factorization,
            implementation=implementation,
            fft_norm=fft_norm,
        )
        # if lifting_channels is passed, make lifting an MLP
        # with a hidden layer of size lifting_channels
        if self.lifting_channels:
            self.lifting = fno_block.MLP(
                in_channels=in_channels,
                out_channels=self.hidden_channels,
                hidden_channels=self.lifting_channels,
                n_layers=2,
                n_dim=self.n_dim,
            )
        # otherwise, make it a linear layer
        else:
            self.lifting = fno_block.MLP(
                in_channels=in_channels,
                out_channels=self.hidden_channels,
                hidden_channels=self.hidden_channels,
                n_layers=1,
                n_dim=self.n_dim,
            )
        self.projection = fno_block.MLP(
            in_channels=self.hidden_channels,
            out_channels=out_channels,
            hidden_channels=self.projection_channels,
            n_layers=2,
            n_dim=self.n_dim,
            non_linearity=non_linearity,
        )

    def forward(self, x):
        """SFNO's forward pass"""
        x = self.concat_to_tensor(x, self.input_keys)

        x = self.lifting(x)
        if self.domain_padding is not None:
            x = self.domain_padding.pad(x)
        # x is 0.4 * [1, 32, 16, 16], passed
        for index in range(self.n_layers):
            x = self.fno_blocks(x, index)

        if self.domain_padding is not None:
            x = self.domain_padding.unpad(x)
        out = self.projection(x)

        return {self.output_keys[0]: out}

forward(x)

SFNO's forward pass

Source code in ppsci/arch/sfnonet.py

def forward(self, x):
    """SFNO's forward pass"""
    x = self.concat_to_tensor(x, self.input_keys)

    x = self.lifting(x)
    if self.domain_padding is not None:
        x = self.domain_padding.pad(x)
    # x is 0.4 * [1, 32, 16, 16], passed
    for index in range(self.n_layers):
        x = self.fno_blocks(x, index)

    if self.domain_padding is not None:
        x = self.domain_padding.unpad(x)
    out = self.projection(x)

    return {self.output_keys[0]: out}

SPINN

Bases: Arch

Separable Physics-Informed Neural Networks.

Parameters:

Name	Type	Description	Default
input_keys	Tuple[str, ...]	Keys of input variables.	required
output_keys	Tuple[str, ...]	Keys of output variables.	required
r	int	Number of features for each output dimension.	required
num_layers	int	Number of layers.	required
hidden_size	Union[int, Tuple[int, ...]]	Size of hidden layer.	required
activation	str	Name of activation function.	'tanh'
skip_connection	bool	Whether to use skip connection.	False
weight_norm	bool	Whether to use weight normalization.	False
periods	Optional[Dict[int, Tuple[float, bool]]]	Periodicity of input variables.	None
fourier	Optional[Dict[str, Union[float, int]]]	Frequency of input variables.	None
random_weight	Optional[Dict[str, float]]	Random weight of linear layer.	None

Examples:

>>> from ppsci.arch import SPINN
>>> model = SPINN(
...     input_keys=('x', 'y', 'z'),
...     output_keys=('u', 'v'),
...     r=32,
...     num_layers=4,
...     hidden_size=32,
... )
>>> input_dict = {"x": paddle.rand([3, 1]),
...               "y": paddle.rand([4, 1]),
...               "z": paddle.rand([5, 1])}
>>> output_dict = model(input_dict)
>>> print(output_dict["u"].shape)
[3, 4, 5, 1]
>>> print(output_dict["v"].shape)
[3, 4, 5, 1]

Source code in ppsci/arch/spinn.py

class SPINN(base.Arch):
    """Separable Physics-Informed Neural Networks.

    Args:
        input_keys (Tuple[str, ...]): Keys of input variables.
        output_keys (Tuple[str, ...]): Keys of output variables.
        r (int): Number of features for each output dimension.
        num_layers (int): Number of layers.
        hidden_size (Union[int, Tuple[int, ...]]): Size of hidden layer.
        activation (str, optional): Name of activation function.
        skip_connection (bool, optional): Whether to use skip connection.
        weight_norm (bool, optional): Whether to use weight normalization.
        periods (Optional[Dict[int, Tuple[float, bool]]], optional): Periodicity of input variables.
        fourier (Optional[Dict[str, Union[float, int]]], optional): Frequency of input variables.
        random_weight (Optional[Dict[str, float]], optional): Random weight of linear layer.

    Examples:
        >>> from ppsci.arch import SPINN
        >>> model = SPINN(
        ...     input_keys=('x', 'y', 'z'),
        ...     output_keys=('u', 'v'),
        ...     r=32,
        ...     num_layers=4,
        ...     hidden_size=32,
        ... )
        >>> input_dict = {"x": paddle.rand([3, 1]),
        ...               "y": paddle.rand([4, 1]),
        ...               "z": paddle.rand([5, 1])}
        >>> output_dict = model(input_dict)
        >>> print(output_dict["u"].shape)
        [3, 4, 5, 1]
        >>> print(output_dict["v"].shape)
        [3, 4, 5, 1]
    """

    def __init__(
        self,
        input_keys: Tuple[str, ...],
        output_keys: Tuple[str, ...],
        r: int,
        num_layers: int,
        hidden_size: Union[int, Tuple[int, ...]],
        activation: str = "tanh",
        skip_connection: bool = False,
        weight_norm: bool = False,
        periods: Optional[Dict[int, Tuple[float, bool]]] = None,
        fourier: Optional[Dict[str, Union[float, int]]] = None,
        random_weight: Optional[Dict[str, float]] = None,
    ):

        super().__init__()
        self.input_keys = input_keys
        self.output_keys = output_keys
        self.r = r
        input_dim = len(self.input_keys)

        self.branch_nets = nn.LayerList()
        for i in range(input_dim):
            self.branch_nets.append(
                ModifiedMLP(
                    input_keys=(input_keys[i],),
                    output_keys=("f",),
                    num_layers=num_layers,
                    hidden_size=hidden_size,
                    activation=activation,
                    skip_connection=skip_connection,
                    weight_norm=weight_norm,
                    output_dim=r * len(output_keys),
                    periods=periods,
                    fourier=fourier,
                    random_weight=random_weight,
                )
            )

        self._init_weights()

    def _init_weights(self):
        for m in self.sublayers(True):
            if isinstance(m, nn.Linear):
                initializer.glorot_normal_(m.weight)
                initializer.zeros_(m.bias)

    def _tensor_contraction(self, x: paddle.Tensor, y: paddle.Tensor) -> paddle.Tensor:
        """Tensor contraction between two tensors along the last channel.

        Args:
            x (Tensor): Input tensor with shape [*N, C].
            y (Tensor): Input tensor with shape [*M, C]

        Returns:
            Tensor: Output tensor with shape [*N, *M, C].
        """
        x_ndim = x.ndim
        y_ndim = y.ndim
        out_dim = x_ndim + y_ndim - 1

        # Align the dimensions of x and y to out_dim
        if x_ndim < out_dim:
            # Add singleton dimensions to x at the end of dimensions
            x = x.unsqueeze([-2] * (out_dim - x_ndim))
        if y_ndim < out_dim:
            # Add singleton dimensions to y at the begin of dimensions
            y = y.unsqueeze([0] * (out_dim - y_ndim))

        # Multiply x and y with implicit broadcasting
        out = x * y

        return out

    def forward_tensor(self, *xs) -> List[paddle.Tensor]:
        # forward each dim branch
        feature_f = []
        for i, input_var in enumerate(xs):
            input_i = {self.input_keys[i]: input_var}
            output_f_i = self.branch_nets[i](input_i)
            feature_f.append(output_f_i["f"])  # [B, r*output_dim]

        output = []
        for i, key in enumerate(self.output_keys):
            st, ed = i * self.r, (i + 1) * self.r
            # do tensor contraction and sum over all branch outputs
            if ed - st == self.r:
                output_i = feature_f[0]
            else:
                output_i = feature_f[0][:, st:ed]

            for j in range(1, len(self.input_keys)):
                if ed - st == self.r:
                    output_ii = feature_f[j]
                else:
                    output_ii = feature_f[j][:, st:ed]
                output_i = self._tensor_contraction(output_i, output_ii)

            output_i = output_i.sum(-1, keepdim=True)
            output.append(output_i)

        return output

    def forward(self, x):
        if self._input_transform is not None:
            x = self._input_transform(x)

        output = self.forward_tensor(*[x[key] for key in self.input_keys])

        output = {key: output[i] for i, key in enumerate(self.output_keys)}

        if self._output_transform is not None:
            output = self._output_transform(x, output)

        return output

STAFNet

Bases: Arch

Spatio-Temporal Attention Fusion Network (STAFNet).

STAFNet is a neural network architecture for spatio-temporal data prediction, combining attention mechanisms and convolutional neural networks to capture spatio-temporal dependencies.

Parameters:

Name	Type	Description	Default
input_keys	Tuple[str, ...]	Keys of input variables.	required
output_keys	Tuple[str, ...]	Keys of output variables.	required
seq_len	int	Sequence length.	required
pred_len	int	Prediction length.	required
aq_gat_node_num	int	Number of nodes in the air quality GAT.	required
aq_gat_node_features	int	Number of features for each node in the air quality GAT.	required
mete_gat_node_num	int	Number of nodes in the meteorological GAT.	required
mete_gat_node_features	int	Number of features for each node in the meteorological GAT.	required
gat_hidden_dim	int	Hidden dimension of the GAT.	required
gat_edge_dim	int	Edge dimension of the GAT.	required
d_model	int	Dimension of the model.	required
n_heads	int	Number of attention heads.	required
e_layers	int	Number of encoder layers.	required
enc_in	int	Encoder input dimension.	required
dec_in	int	Decoder input dimension.	required
freq	str	Frequency for positional encoding.	required
embed	str	Embedding type.	required
d_ff	int	Dimension of feedforward network.	required
top_k	int	Number of top frequencies to consider.	required
num_kernels	int	Number of kernels in Inception block.	required
dropout	float	Dropout rate.	required
output_attention	bool	Whether to output attention.	required
factor	int	Factor for attention mechanism.	required
c_out	int	Output channels.	required

Examples:

>>> from ppsci.arch import STAFNet
>>> model = STAFNet(
...     input_keys=('x', 'y', 'z'),
...     output_keys=('u', 'v'),
...     seq_len=72,
...     pred_len=48,
...     aq_gat_node_num=10,
...     aq_gat_node_features=16,
...     mete_gat_node_num=10,
...     mete_gat_node_features=16,
...     gat_hidden_dim=32,
...     gat_edge_dim=8,
...     d_model=512,
...     n_heads=8,
...     e_layers=3,
...     enc_in=7,
...     dec_in=7,
...     freq='h',
...     embed='fixed',
...     d_ff=2048,
...     top_k=5,
...     num_kernels=32,
...     dropout=0.1,
...     output_attention=False,
...     factor=5,
...     c_out=7,
... )
>>> input_dict = {"x": paddle.rand([32, 72, 7]),
...               "y": paddle.rand([32, 72, 7]),
...               "z": paddle.rand([32, 72, 7])}
>>> output_dict = model(input_dict)
>>> print(output_dict["u"].shape)
[32, 48, 7]
>>> print(output_dict["v"].shape)
[32, 48, 7]

Source code in ppsci/arch/stafnet.py

class STAFNet(base.Arch):
    """Spatio-Temporal Attention Fusion Network (STAFNet).

    STAFNet is a neural network architecture for spatio-temporal data prediction, combining attention mechanisms and convolutional neural networks to capture spatio-temporal dependencies.

    Args:
        input_keys (Tuple[str, ...]): Keys of input variables.
        output_keys (Tuple[str, ...]): Keys of output variables.
        seq_len (int): Sequence length.
        pred_len (int): Prediction length.
        aq_gat_node_num (int): Number of nodes in the air quality GAT.
        aq_gat_node_features (int): Number of features for each node in the air quality GAT.
        mete_gat_node_num (int): Number of nodes in the meteorological GAT.
        mete_gat_node_features (int): Number of features for each node in the meteorological GAT.
        gat_hidden_dim (int): Hidden dimension of the GAT.
        gat_edge_dim (int): Edge dimension of the GAT.
        d_model (int): Dimension of the model.
        n_heads (int): Number of attention heads.
        e_layers (int): Number of encoder layers.
        enc_in (int): Encoder input dimension.
        dec_in (int): Decoder input dimension.
        freq (str): Frequency for positional encoding.
        embed (str): Embedding type.
        d_ff (int): Dimension of feedforward network.
        top_k (int): Number of top frequencies to consider.
        num_kernels (int): Number of kernels in Inception block.
        dropout (float): Dropout rate.
        output_attention (bool): Whether to output attention.
        factor (int): Factor for attention mechanism.
        c_out (int): Output channels.

    Examples:
        >>> from ppsci.arch import STAFNet
        >>> model = STAFNet(
        ...     input_keys=('x', 'y', 'z'),
        ...     output_keys=('u', 'v'),
        ...     seq_len=72,
        ...     pred_len=48,
        ...     aq_gat_node_num=10,
        ...     aq_gat_node_features=16,
        ...     mete_gat_node_num=10,
        ...     mete_gat_node_features=16,
        ...     gat_hidden_dim=32,
        ...     gat_edge_dim=8,
        ...     d_model=512,
        ...     n_heads=8,
        ...     e_layers=3,
        ...     enc_in=7,
        ...     dec_in=7,
        ...     freq='h',
        ...     embed='fixed',
        ...     d_ff=2048,
        ...     top_k=5,
        ...     num_kernels=32,
        ...     dropout=0.1,
        ...     output_attention=False,
        ...     factor=5,
        ...     c_out=7,
        ... )
        >>> input_dict = {"x": paddle.rand([32, 72, 7]),
        ...               "y": paddle.rand([32, 72, 7]),
        ...               "z": paddle.rand([32, 72, 7])}
        >>> output_dict = model(input_dict)
        >>> print(output_dict["u"].shape)
        [32, 48, 7]
        >>> print(output_dict["v"].shape)
        [32, 48, 7]
    """

    def __init__(
        self,
        input_keys: Tuple[str, ...],
        output_keys: Tuple[str, ...],
        seq_len: int,
        pred_len: int,
        aq_gat_node_num: int,
        aq_gat_node_features: int,
        mete_gat_node_num: int,
        mete_gat_node_features: int,
        gat_hidden_dim: int,
        gat_edge_dim: int,
        d_model: int,
        n_heads: int,
        e_layers: int,
        enc_in: int,
        dec_in: int,
        freq: str,
        embed: str,
        d_ff: int,
        top_k: int,
        num_kernels: int,
        dropout: float,
        output_attention: bool,
        factor: int,
        c_out: int,
    ):
        super(STAFNet, self).__init__()
        self.input_keys = input_keys
        self.output_keys = output_keys
        self.output_attention = output_attention
        self.seq_len = seq_len
        self.pred_len = pred_len
        self.dec_in = dec_in
        self.gat_embed_dim = enc_in
        self.enc_embedding = DataEmbedding(enc_in, d_model, embed, freq, dropout)
        self.aq_gat_node_num = aq_gat_node_num
        self.aq_gat_node_features = aq_gat_node_features
        self.aq_GAT = GAT_Encoder(
            aq_gat_node_features,
            gat_hidden_dim,
            gat_edge_dim,
            self.gat_embed_dim,
            dropout,
        )
        self.mete_gat_node_num = mete_gat_node_num
        self.mete_gat_node_features = mete_gat_node_features
        self.mete_GAT = GAT_Encoder(
            mete_gat_node_features,
            gat_hidden_dim,
            gat_edge_dim,
            self.gat_embed_dim,
            dropout,
        )
        self.pos_fc = paddle.nn.Linear(
            in_features=2, out_features=self.gat_embed_dim, bias_attr=True
        )
        self.fusion_Attention = AttentionLayer(
            ProbAttention(
                False,
                factor,
                attention_dropout=dropout,
                output_attention=self.output_attention,
            ),
            self.gat_embed_dim,
            n_heads,
        )
        self.model = paddle.nn.LayerList(
            sublayers=[
                TimesBlock(seq_len, pred_len, top_k, num_kernels, d_model, d_ff)
                for _ in range(e_layers)
            ]
        )
        self.layer = e_layers
        self.layer_norm = paddle.nn.LayerNorm(normalized_shape=d_model)
        self.predict_linear = paddle.nn.Linear(
            in_features=self.seq_len, out_features=self.pred_len + self.seq_len
        )
        self.projection = paddle.nn.Linear(
            in_features=d_model, out_features=c_out, bias_attr=True
        )
        self.output_attention = output_attention

    def aq_gat(self, G):
        g_batch = G.num_graph
        batch_size = int(g_batch / self.seq_len)
        gat_output = self.aq_GAT(
            G, G.node_feat["feature"][:, -self.aq_gat_node_features :]
        )
        gat_output = gat_output.reshape(
            (batch_size, self.seq_len, self.aq_gat_node_num, self.gat_embed_dim)
        )
        gat_output = paddle.flatten(x=gat_output, start_axis=0, stop_axis=1)
        return gat_output

    def mete_gat(self, G):
        g_batch = G.num_graph
        batch_size = int(g_batch / self.seq_len)
        gat_output = self.mete_GAT(
            G, G.node_feat["feature"][:, -self.mete_gat_node_features :]
        )
        gat_output = gat_output.reshape(
            (batch_size, self.seq_len, self.mete_gat_node_num, self.gat_embed_dim)
        )
        gat_output = paddle.flatten(x=gat_output, start_axis=0, stop_axis=1)
        return gat_output

    def norm_pos(self, A, B):
        A_mean = paddle.mean(A, axis=0)
        A_std = paddle.std(A, axis=0)

        A_norm = (A - A_mean) / A_std
        B_norm = (B - A_mean) / A_std
        return A_norm, B_norm

    def forward(self, Data, mask=None):
        train_data = Data["aq_train_data"]
        batch_size = train_data.shape[0]
        x = train_data
        perm_1 = list(range(x.ndim))
        perm_1[1] = 2
        perm_1[2] = 1
        train_data = paddle.transpose(x=x, perm=perm_1)
        train_data = paddle.flatten(x=train_data, start_axis=0, stop_axis=1)
        x_enc = train_data[:, : self.seq_len, -self.dec_in :]
        x_mark_enc = train_data[:, : self.seq_len, 1:6]

        means = x_enc.mean(axis=1, keepdim=True).detach()
        x_enc = x_enc - means
        stdev = paddle.sqrt(
            x=paddle.var(x=x_enc, axis=1, keepdim=True, unbiased=False) + 1e-05
        )
        x_enc /= stdev
        # enc_out = self.enc_embedding(aq_gat_output, x_mark_enc)
        enc_out = self.enc_embedding(x_enc, x_mark_enc)
        enc_out = self.predict_linear(enc_out.transpose(perm=[0, 2, 1])).transpose(
            perm=[0, 2, 1]
        )
        for i in range(self.layer):
            enc_out, period_list, period_weight = self.model[i](enc_out)
            enc_out = self.layer_norm(enc_out)
        dec_out = self.projection(enc_out)
        dec_out = dec_out * paddle.tile(
            stdev[:, 0, :].unsqueeze(axis=1), (1, self.pred_len + self.seq_len, 1)
        )
        dec_out = dec_out + paddle.tile(
            means[:, 0, :].unsqueeze(axis=1), (1, self.pred_len + self.seq_len, 1)
        )

        return {
            self.output_keys[0]: dec_out[:, -self.pred_len :, -7:].reshape(
                (batch_size, self.pred_len, -1, 7)
            )
        }

TFNO1dNet

Bases: FNONet

1D Fourier Neural Operator.

Parameters:

Name	Type	Description	Default
input_keys	Tuple[str, ...]	Name of input keys, such as ("input",).	required
output_keys	Tuple[str, ...]	Name of output keys, such as ("output",).	required
n_modes_height	Tuple[int, ...]	Number of Fourier modes to keep along the height, along each dimension.	required
hidden_channels	int	Width of the FNO (i.e. number of channels).	required
in_channels	int	Number of input channels. Defaults to 3.	3
out_channels	int	Number of output channels. Defaults to 1.	1
lifting_channels	int	Number of hidden channels of the lifting block of the FNO. Defaults to 256.	256
projection_channels	int	Number of hidden channels of the projection block of the FNO. Defaults to 256.	256
n_layers	int	Number of Fourier Layers. Defaults to 4.	4
use_mlp	bool	Whether to use an MLP layer after each FNO block. Defaults to False.	False
mlp	dict[str, float]	Parameters of the MLP. {'expansion': float, 'dropout': float}. Defaults to None.	None
non_linearity	functional	Non-Linearity module to use. Defaults to F.gelu.	gelu
norm	module	Normalization layer to use. Defaults to None.	None
preactivation	bool	Whether to use resnet-style preactivation. Defaults to False.	False
skip	str	Type of skip connection to use,{'linear', 'identity', 'soft-gating'}. Defaults to "soft-gating".	'soft-gating'
separable	bool	Whether to use a depthwise separable spectral convolution. Defaults to False.	False
factorization	str	Tensor factorization of the parameters weight to use. * If None, a dense tensor parametrizes the Spectral convolutions. * Otherwise, the specified tensor factorization is used. Defaults to "Tucker".	'Tucker'
rank	float	Rank of the tensor factorization of the Fourier weights. Defaults to 1.0.	1.0
joint_factorization	bool	Whether all the Fourier Layers should be parametrized by a single tensor (vs one per layer). Defaults to False.	False
implementation	str	{'factorized', 'reconstructed'}, optional. Defaults to "factorized". If factorization is not None, forward mode to use:: * reconstructed : the full weight tensor is reconstructed from the factorization and used for the forward pass. * factorized : the input is directly contracted with the factors of the decomposition.	'factorized'
domain_padding	Optional[Union[list, float, int]]	Whether to use percentage of padding. Defaults to None.	None
domain_padding_mode	str	{'symmetric', 'one-sided'}, optional How to perform domain padding, by default 'one-sided'. Defaults to "one-sided".	'one-sided'
fft_norm	str	The normalization mode for the FFT. Defaults to "forward".	'forward'
patching_levels	int	Number of patching levels to use. Defaults to 0.	0
SpectralConv	layer	Spectral convolution layer to use. Defaults to fno_block.FactorizedSpectralConv.	FactorizedSpectralConv

Source code in ppsci/arch/tfnonet.py

class TFNO1dNet(FNONet):
    """1D Fourier Neural Operator.

    Args:
        input_keys (Tuple[str, ...]): Name of input keys, such as ("input",).
        output_keys (Tuple[str, ...]): Name of output keys, such as ("output",).
        n_modes_height (Tuple[int, ...]): Number of Fourier modes to keep along the height, along each
            dimension.
        hidden_channels (int): Width of the FNO (i.e. number of channels).
        in_channels (int, optional): Number of input channels. Defaults to 3.
        out_channels (int, optional): Number of output channels. Defaults to 1.
        lifting_channels (int, optional): Number of hidden channels of the lifting block of the FNO.
            Defaults to 256.
        projection_channels (int, optional): Number of hidden channels of the projection block of the FNO.
            Defaults to 256.
        n_layers (int, optional): Number of Fourier Layers. Defaults to 4.
        use_mlp (bool, optional): Whether to use an MLP layer after each FNO block. Defaults to False.
        mlp (dict[str, float], optional): Parameters of the MLP. {'expansion': float, 'dropout': float}.
            Defaults to None.
        non_linearity (nn.functional, optional): Non-Linearity module to use. Defaults to F.gelu.
        norm (F.module, optional): Normalization layer to use. Defaults to None.
        preactivation (bool, optional): Whether to use resnet-style preactivation. Defaults to False.
        skip (str, optional): Type of skip connection to use,{'linear', 'identity', 'soft-gating'}.
            Defaults to "soft-gating".
        separable (bool, optional): Whether to use a depthwise separable spectral convolution.
            Defaults to  False.
        factorization (str, optional): Tensor factorization of the parameters weight to use.
            * If None, a dense tensor parametrizes the Spectral convolutions.
            * Otherwise, the specified tensor factorization is used. Defaults to "Tucker".
        rank (float, optional): Rank of the tensor factorization of the Fourier weights. Defaults to 1.0.
        joint_factorization (bool, optional): Whether all the Fourier Layers should be parametrized by a
            single tensor (vs one per layer). Defaults to False.
        implementation (str, optional): {'factorized', 'reconstructed'}, optional. Defaults to "factorized".
            If factorization is not None, forward mode to use::
            * `reconstructed` : the full weight tensor is reconstructed from the factorization and used for the forward pass.
            * `factorized` : the input is directly contracted with the factors of the decomposition.
        domain_padding (Optional[Union[list, float, int]], optional): Whether to use percentage of padding.
            Defaults to None.
        domain_padding_mode (str, optional): {'symmetric', 'one-sided'}, optional
            How to perform domain padding, by default 'one-sided'. Defaults to "one-sided".
        fft_norm (str, optional): The normalization mode for the FFT. Defaults to "forward".
        patching_levels (int, optional): Number of patching levels to use. Defaults to 0.
        SpectralConv (nn.layer, optional): Spectral convolution layer to use.
            Defaults to fno_block.FactorizedSpectralConv.
    """

    def __init__(
        self,
        input_keys: Tuple[str, ...],
        output_keys: Tuple[str, ...],
        n_modes_height: Tuple[int, ...],
        hidden_channels: int,
        in_channels: int = 3,
        out_channels: int = 1,
        lifting_channels: int = 256,
        projection_channels: int = 256,
        n_layers: int = 4,
        non_linearity: nn.functional = F.gelu,
        use_mlp: bool = False,
        mlp: Optional[Dict[str, float]] = None,
        norm: str = None,
        skip: str = "soft-gating",
        separable: bool = False,
        preactivation: bool = False,
        factorization: str = "Tucker",
        rank: float = 1.0,
        joint_factorization: bool = False,
        implementation: str = "factorized",
        domain_padding: Optional[Union[list, float, int]] = None,
        domain_padding_mode: str = "one-sided",
        fft_norm: str = "forward",
        patching_levels: int = 0,
        SpectralConv: nn.Layer = fno_block.FactorizedSpectralConv,
        **kwargs,
    ):
        super().__init__(
            input_keys=input_keys,
            output_keys=output_keys,
            n_modes=(n_modes_height,),
            hidden_channels=hidden_channels,
            in_channels=in_channels,
            out_channels=out_channels,
            lifting_channels=lifting_channels,
            projection_channels=projection_channels,
            n_layers=n_layers,
            non_linearity=non_linearity,
            use_mlp=use_mlp,
            mlp=mlp,
            norm=norm,
            skip=skip,
            separable=separable,
            preactivation=preactivation,
            factorization=factorization,
            rank=rank,
            joint_factorization=joint_factorization,
            implementation=implementation,
            domain_padding=domain_padding,
            domain_padding_mode=domain_padding_mode,
            fft_norm=fft_norm,
            patching_levels=patching_levels,
            SpectralConv=SpectralConv,
        )
        self.n_modes_height = n_modes_height

TFNO2dNet

Bases: FNONet

2D Fourier Neural Operator.

Parameters:

Name	Type	Description	Default
input_keys	Tuple[str, ...]	Name of input keys, such as ("input",).	required
output_keys	Tuple[str, ...]	Name of output keys, such as ("output",).	required
n_modes_height	int	Number of Fourier modes to keep along the height.	required
n_modes_width	int	Number of modes to keep in Fourier Layer, along the width.	required
hidden_channels	int	Width of the FNO (i.e. number of channels).	required
in_channels	int	Number of input channels. Defaults to 3.	3
out_channels	int	Number of output channels. Defaults to 1.	1
lifting_channels	int	Number of hidden channels of the lifting block of the FNO. Defaults to 256.	256
projection_channels	int	Number of hidden channels of the projection block of the FNO. Defaults to 256.	256
n_layers	int	Number of Fourier Layers. Defaults to 4.	4
use_mlp	bool	Whether to use an MLP layer after each FNO block. Defaults to False.	False
mlp	Dict[str, float]	Parameters of the MLP. {'expansion': float, 'dropout': float}. Defaults to None.	None
non_linearity	Layer	Non-Linearity module to use. Defaults to F.gelu.	gelu
norm	module	Normalization layer to use. Defaults to None.	None
preactivation	bool	Whether to use resnet-style preactivation. Defaults to False.	False
skip	str	Type of skip connection to use,{'linear', 'identity', 'soft-gating'}. Defaults to "soft-gating".	'soft-gating'
separable	bool	Whether to use a depthwise separable spectral convolution. Defaults to False.	False
factorization	str	Tensor factorization of the parameters weight to use. * If None, a dense tensor parametrizes the Spectral convolutions. * Otherwise, the specified tensor factorization is used. Defaults to "Tucker".	'Tucker'
rank	float	Rank of the tensor factorization of the Fourier weights. Defaults to 1.0.	1.0
joint_factorization	bool	Whether all the Fourier Layers should be parametrized by a single tensor (vs one per layer). Defaults to False.	False
implementation	str	{'factorized', 'reconstructed'}, optional. Defaults to "factorized". If factorization is not None, forward mode to use:: * reconstructed : the full weight tensor is reconstructed from the factorization and used for the forward pass. * factorized : the input is directly contracted with the factors of the decomposition.	'factorized'
domain_padding	Union[list, float, int]	Whether to use percentage of padding. Defaults to None.	None
domain_padding_mode	str	{'symmetric', 'one-sided'}, optional How to perform domain padding, by default 'one-sided'. Defaults to "one-sided".	'one-sided'
fft_norm	str	The normalization mode for the FFT. Defaults to "forward".	'forward'
patching_levels	int	Number of patching levels to use. Defaults to 0.	0
SpectralConv	layer	Spectral convolution layer to use. Defaults to fno_block.FactorizedSpectralConv.	FactorizedSpectralConv

Source code in ppsci/arch/tfnonet.py

class TFNO2dNet(FNONet):
    """2D Fourier Neural Operator.

    Args:
       input_keys (Tuple[str, ...]): Name of input keys, such as ("input",).
       output_keys (Tuple[str, ...]): Name of output keys, such as ("output",).
       n_modes_height (int): Number of Fourier modes to keep along the height.
       n_modes_width (int): Number of modes to keep in Fourier Layer, along the width.
       hidden_channels (int): Width of the FNO (i.e. number of channels).
       in_channels (int, optional): Number of input channels. Defaults to 3.
       out_channels (int, optional): Number of output channels. Defaults to 1.
       lifting_channels (int, optional): Number of hidden channels of the lifting block of the FNO.
           Defaults to 256.
       projection_channels (int, optional): Number of hidden channels of the projection block of the FNO.
           Defaults to 256.
       n_layers (int, optional): Number of Fourier Layers. Defaults to 4.
       use_mlp (bool, optional): Whether to use an MLP layer after each FNO block. Defaults to False.
       mlp (Dict[str, float], optional): Parameters of the MLP. {'expansion': float, 'dropout': float}.
           Defaults to None.
       non_linearity (nn.Layer, optional): Non-Linearity module to use. Defaults to F.gelu.
       norm (F.module, optional): Normalization layer to use. Defaults to None.
       preactivation (bool, optional): Whether to use resnet-style preactivation. Defaults to False.
       skip (str, optional): Type of skip connection to use,{'linear', 'identity', 'soft-gating'}.
           Defaults to "soft-gating".
       separable (bool, optional): Whether to use a depthwise separable spectral convolution.
           Defaults to  False.
       factorization (str, optional): Tensor factorization of the parameters weight to use.
           * If None, a dense tensor parametrizes the Spectral convolutions.
           * Otherwise, the specified tensor factorization is used. Defaults to "Tucker".
       rank (float, optional): Rank of the tensor factorization of the Fourier weights. Defaults to 1.0.
       joint_factorization (bool, optional): Whether all the Fourier Layers should be parametrized by a
           single tensor (vs one per layer). Defaults to False.
       implementation (str, optional): {'factorized', 'reconstructed'}, optional. Defaults to "factorized".
           If factorization is not None, forward mode to use::
           * `reconstructed` : the full weight tensor is reconstructed from the factorization and used for the forward pass.
           * `factorized` : the input is directly contracted with the factors of the decomposition.
       domain_padding (Union[list,float,int], optional): Whether to use percentage of padding. Defaults to
            None.
       domain_padding_mode (str, optional): {'symmetric', 'one-sided'}, optional
           How to perform domain padding, by default 'one-sided'. Defaults to "one-sided".
       fft_norm (str, optional): The normalization mode for the FFT. Defaults to "forward".
       patching_levels (int, optional): Number of patching levels to use. Defaults to 0.
       SpectralConv (nn.layer, optional): Spectral convolution layer to use.
            Defaults to fno_block.FactorizedSpectralConv.
    """

    def __init__(
        self,
        input_keys: Tuple[str, ...],
        output_keys: Tuple[str, ...],
        n_modes_height: int,
        n_modes_width: int,
        hidden_channels: int,
        in_channels: int = 3,
        out_channels: int = 1,
        lifting_channels: int = 256,
        projection_channels: int = 256,
        n_layers: int = 4,
        non_linearity: nn.functional = F.gelu,
        use_mlp: bool = False,
        mlp: Optional[Dict[str, float]] = None,
        norm: str = None,
        skip: str = "soft-gating",
        separable: bool = False,
        preactivation: bool = False,
        factorization: str = "Tucker",
        rank: float = 1.0,
        joint_factorization: bool = False,
        implementation: str = "factorized",
        domain_padding: Optional[Union[list, float, int]] = None,
        domain_padding_mode: str = "one-sided",
        fft_norm: str = "forward",
        patching_levels: int = 0,
        SpectralConv: nn.layer = fno_block.FactorizedSpectralConv,
        **kwargs,
    ):
        super().__init__(
            input_keys=input_keys,
            output_keys=output_keys,
            n_modes=(n_modes_height, n_modes_width),
            hidden_channels=hidden_channels,
            in_channels=in_channels,
            out_channels=out_channels,
            lifting_channels=lifting_channels,
            projection_channels=projection_channels,
            n_layers=n_layers,
            non_linearity=non_linearity,
            use_mlp=use_mlp,
            mlp=mlp,
            norm=norm,
            skip=skip,
            separable=separable,
            preactivation=preactivation,
            factorization=factorization,
            rank=rank,
            joint_factorization=joint_factorization,
            implementation=implementation,
            domain_padding=domain_padding,
            domain_padding_mode=domain_padding_mode,
            fft_norm=fft_norm,
            patching_levels=patching_levels,
            SpectralConv=SpectralConv,
        )
        self.n_modes_height = n_modes_height
        self.n_modes_width = n_modes_width

TFNO3dNet

Bases: FNONet

3D Fourier Neural Operator.

Parameters:

Name	Type	Description	Default
input_keys	Tuple[str, ...]	Name of input keys, such as ("input",).	required
output_keys	Tuple[str, ...]	Name of output keys, such as ("output",).	required
n_modes_height	int	Number of Fourier modes to keep along the height.	required
n_modes_width	int	Number of modes to keep in Fourier Layer, along the width.	required
n_modes_depth	int	Number of Fourier modes to keep along the depth.	required
hidden_channels	int	Width of the FNO (i.e. number of channels).	required
in_channels	int	Number of input channels. Defaults to 3.	3
out_channels	int	Number of output channels. Defaults to 1.	1
lifting_channels	int	Number of hidden channels of the lifting block of the FNO. Defaults to 256.	256
projection_channels	int	Number of hidden channels of the projection block of the FNO. Defaults to 256.	256
n_layers	int	Number of Fourier Layers. Defaults to 4.	4
use_mlp	bool	Whether to use an MLP layer after each FNO block. Defaults to False.	False
mlp	Dict[str, float]	Parameters of the MLP. {'expansion': float, 'dropout': float}. Defaults to None.	None
non_linearity	Layer	Non-Linearity module to use. Defaults to F.gelu.	gelu
norm	module	Normalization layer to use. Defaults to None.	None
preactivation	bool	Whether to use resnet-style preactivation. Defaults to False.	False
skip	str	Type of skip connection to use,{'linear', 'identity', 'soft-gating'}. Defaults to "soft-gating".	'soft-gating'
separable	bool	Whether to use a depthwise separable spectral convolution. Defaults to False.	False
factorization	str	Tensor factorization of the parameters weight to use. * If None, a dense tensor parametrizes the Spectral convolutions. * Otherwise, the specified tensor factorization is used. Defaults to "Tucker".	'Tucker'
rank	float	Rank of the tensor factorization of the Fourier weights. Defaults to 1.0.	1.0
joint_factorization	bool	Whether all the Fourier Layers should be parametrized by a single tensor (vs one per layer). Defaults to False.	False
implementation	str	{'factorized', 'reconstructed'}, optional. Defaults to "factorized". If factorization is not None, forward mode to use:: * reconstructed : the full weight tensor is reconstructed from the factorization and used for the forward pass. * factorized : the input is directly contracted with the factors of the decomposition.	'factorized'
domain_padding	str	Whether to use percentage of padding. Defaults to None.	None
domain_padding_mode	str	{'symmetric', 'one-sided'}, optional How to perform domain padding, by default 'one-sided'. Defaults to "one-sided".	'one-sided'
fft_norm	str	The normalization mode for the FFT. Defaults to "forward".	'forward'
patching_levels	int	Number of patching levels to use. Defaults to 0.	0
SpectralConv	layer	Spectral convolution layer to use. Defaults to fno_block. FactorizedSpectralConv.	FactorizedSpectralConv

Source code in ppsci/arch/tfnonet.py

class TFNO3dNet(FNONet):
    """3D Fourier Neural Operator.

    Args:
        input_keys (Tuple[str, ...]): Name of input keys, such as ("input",).
        output_keys (Tuple[str, ...]): Name of output keys, such as ("output",).
        n_modes_height (int): Number of Fourier modes to keep along the height.
        n_modes_width (int): Number of modes to keep in Fourier Layer, along the width.
        n_modes_depth (int): Number of Fourier modes to keep along the depth.
        hidden_channels (int): Width of the FNO (i.e. number of channels).
        in_channels (int, optional): Number of input channels. Defaults to 3.
        out_channels (int, optional): Number of output channels. Defaults to 1.
        lifting_channels (int, optional): Number of hidden channels of the lifting block of the FNO.
            Defaults to 256.
        projection_channels (int, optional): Number of hidden channels of the projection block of the FNO.
            Defaults to 256.
        n_layers (int, optional): Number of Fourier Layers. Defaults to 4.
        use_mlp (bool, optional): Whether to use an MLP layer after each FNO block. Defaults to False.
        mlp (Dict[str, float], optional): Parameters of the MLP. {'expansion': float, 'dropout': float}.
            Defaults to None.
        non_linearity (nn.Layer, optional): Non-Linearity module to use. Defaults to F.gelu.
        norm (F.module, optional): Normalization layer to use. Defaults to None.
        preactivation (bool, optional): Whether to use resnet-style preactivation. Defaults to False.
        skip (str, optional): Type of skip connection to use,{'linear', 'identity', 'soft-gating'}.
            Defaults to "soft-gating".
        separable (bool, optional): Whether to use a depthwise separable spectral convolution.
            Defaults to  False.
        factorization (str, optional): Tensor factorization of the parameters weight to use.
            * If None, a dense tensor parametrizes the Spectral convolutions.
            * Otherwise, the specified tensor factorization is used. Defaults to "Tucker".
        rank (float, optional): Rank of the tensor factorization of the Fourier weights. Defaults to 1.0.
        joint_factorization (bool, optional): Whether all the Fourier Layers should be parametrized by a
            single tensor (vs one per layer). Defaults to False.
        implementation (str, optional): {'factorized', 'reconstructed'}, optional. Defaults to "factorized".
            If factorization is not None, forward mode to use::
            * `reconstructed` : the full weight tensor is reconstructed from the factorization and used for the forward pass.
            * `factorized` : the input is directly contracted with the factors of the decomposition.
        domain_padding (str, optional): Whether to use percentage of padding. Defaults to None.
        domain_padding_mode (str, optional): {'symmetric', 'one-sided'}, optional
            How to perform domain padding, by default 'one-sided'. Defaults to "one-sided".
        fft_norm (str, optional): The normalization mode for the FFT. Defaults to "forward".
        patching_levels (int, optional): Number of patching levels to use. Defaults to 0.
        SpectralConv (nn.layer, optional): Spectral convolution layer to use. Defaults to fno_block.
            FactorizedSpectralConv.
    """

    def __init__(
        self,
        input_keys: Tuple[str, ...],
        output_keys: Tuple[str, ...],
        n_modes_height: int,
        n_modes_width: int,
        n_modes_depth: int,
        hidden_channels: int,
        in_channels: int = 3,
        out_channels: int = 1,
        lifting_channels: int = 256,
        projection_channels: int = 256,
        n_layers: int = 4,
        non_linearity: nn.functional = F.gelu,
        use_mlp: bool = False,
        mlp: Optional[Dict[str, float]] = None,
        norm: str = None,
        skip: str = "soft-gating",
        separable: bool = False,
        preactivation: bool = False,
        factorization: str = "Tucker",
        rank: float = 1.0,
        joint_factorization: bool = False,
        implementation: str = "factorized",
        domain_padding: Optional[Union[list, float, int]] = None,
        domain_padding_mode: str = "one-sided",
        fft_norm: str = "forward",
        patching_levels: int = 0,
        SpectralConv: nn.layer = fno_block.FactorizedSpectralConv,
        **kwargs,
    ):
        super().__init__(
            input_keys=input_keys,
            output_keys=output_keys,
            n_modes=(n_modes_height, n_modes_width, n_modes_depth),
            hidden_channels=hidden_channels,
            in_channels=in_channels,
            out_channels=out_channels,
            lifting_channels=lifting_channels,
            projection_channels=projection_channels,
            n_layers=n_layers,
            non_linearity=non_linearity,
            use_mlp=use_mlp,
            mlp=mlp,
            norm=norm,
            skip=skip,
            separable=separable,
            preactivation=preactivation,
            factorization=factorization,
            rank=rank,
            joint_factorization=joint_factorization,
            implementation=implementation,
            domain_padding=domain_padding,
            domain_padding_mode=domain_padding_mode,
            fft_norm=fft_norm,
            patching_levels=patching_levels,
            SpectralConv=SpectralConv,
        )
        self.n_modes_height = n_modes_height
        self.n_modes_width = n_modes_width
        self.n_modes_depth = n_modes_depth

Transformer

Bases: Arch

A Kind of Transformer Model.

Parameters:

Name	Type	Description	Default
input_keys	Tuple[str, ...]	Name of input keys, such as ("x", "y", "z").	required
output_keys	Tuple[str, ...]	Name of output keys, such as ("u", "v", "w").	required
num_var_max	int	Maximum number of variables.	required
vocab_size	int	Size of vocab. Size of unary operators = 1, binary operators = 2.	required
seq_length	int	Length of sequance.	required
d_model	int	The innermost dimension of model. Defaults to 256.	256
heads	int	The number of independent heads for the multi-head attention layers. Defaults to 4.	4
num_layers_enc	int	The number of encoders. Defaults to 4.	4
num_layers_dec	int	The number of decoders. Defaults to 8.	8
dropout	float	Dropout regularization. Defaults to 0.0.	0.0

Examples:

>>> import paddle
>>> import ppsci
>>> model = ppsci.arch.Transformer(
...     input_keys=("input", "target_seq"),
...     output_keys=("output",),
...     num_var_max=7,
...     vocab_size=20,
...     seq_length=30,
... )
>>> input_dict = {"input": paddle.rand([512, 50, 7, 1]),
...               "target_seq": paddle.rand([512, 30])}
>>> output_dict = model(input_dict)
>>> print(output_dict["output"].shape)
[512, 30, 20]

Source code in ppsci/arch/transformer.py

class Transformer(base.Arch):
    """A Kind of Transformer Model.

    Args:
        input_keys (Tuple[str, ...]): Name of input keys, such as ("x", "y", "z").
        output_keys (Tuple[str, ...]): Name of output keys, such as ("u", "v", "w").
        num_var_max (int): Maximum number of variables.
        vocab_size (int): Size of vocab. Size of unary operators = 1, binary operators = 2.
        seq_length (int): Length of sequance.
        d_model (int, optional): The innermost dimension of model. Defaults to 256.
        heads (int, optional): The number of independent heads for the multi-head attention layers. Defaults to 4.
        num_layers_enc (int, optional): The number of encoders. Defaults to 4.
        num_layers_dec (int, optional): The number of decoders. Defaults to 8.
        dropout (float, optional): Dropout regularization. Defaults to 0.0.

    Examples:
        >>> import paddle
        >>> import ppsci
        >>> model = ppsci.arch.Transformer(
        ...     input_keys=("input", "target_seq"),
        ...     output_keys=("output",),
        ...     num_var_max=7,
        ...     vocab_size=20,
        ...     seq_length=30,
        ... )
        >>> input_dict = {"input": paddle.rand([512, 50, 7, 1]),
        ...               "target_seq": paddle.rand([512, 30])}
        >>> output_dict = model(input_dict)
        >>> print(output_dict["output"].shape)
        [512, 30, 20]
    """

    def __init__(
        self,
        input_keys: Tuple[str, ...],
        output_keys: Tuple[str, ...],
        num_var_max: int,
        vocab_size: int,
        seq_length: int,
        d_model: int = 256,
        heads: int = 4,
        num_layers_enc: int = 4,
        num_layers_dec: int = 8,
        act: str = "relu",
        dropout: float = 0.0,
    ):
        super().__init__()
        self.input_keys = input_keys
        self.output_keys = output_keys
        self.num_var_max = num_var_max
        self.vocab_size = vocab_size
        self.seq_length = seq_length
        self.d_model = d_model
        self.heads = heads
        self.num_layers_enc = num_layers_enc
        self.num_layers_dec = num_layers_dec
        self.act = act
        self.dropout = dropout

        self.encoder = Encoder(
            num_layers_enc, num_var_max, d_model, heads, act="relu", dropout=dropout
        )
        self.decoder = Decoder(
            num_layers_dec,
            vocab_size,
            seq_length,
            d_model,
            heads,
            act="relu",
            dropout=dropout,
        )
        self.last_layer = paddle.nn.Linear(in_features=d_model, out_features=vocab_size)

    def get_mask(self, target_seq):
        padding_mask = paddle.equal(target_seq, 0).unsqueeze(axis=1).unsqueeze(axis=1)
        future_mask = paddle.triu(
            paddle.ones(shape=[target_seq.shape[1], target_seq.shape[1]]),
            diagonal=1,
        ).astype(dtype="bool")
        mask = paddle.logical_or(x=padding_mask, y=future_mask)
        return mask

    def forward_tensor(self, x_lst):
        y, target_seq = x_lst[0], x_lst[1]
        mask = self.get_mask(target_seq)
        y_enc = self.encoder(y)
        y = self.decoder(target_seq, y_enc, mask)
        y = self.last_layer(y)
        return y

    def forward(self, x):
        if self._input_transform is not None:
            x = self._input_transform(x)

        x_lst = [x[key] for key in self.input_keys]  # input, target_seq
        y = self.forward_tensor(x_lst)
        y = self.split_to_dict(y, self.output_keys, axis=-1)

        if self._output_transform is not None:
            y = self._output_transform(x, y)
        return y

    @paddle.no_grad()
    def decode_process(
        self, dataset: paddle.Tensor, complete_func: Callable
    ) -> paddle.Tensor:
        """Greedy decode with the Transformer model, decode until the equation tree is completed.

        Args:
            dataset (paddle.Tensor): Tabular dataset.
            complete_func (Callable): Function used to calculate whether inference is complete.
        """
        encoder_output = self.encoder(dataset)
        decoder_output = paddle.zeros(
            shape=(dataset.shape[0], self.seq_length + 1), dtype=paddle.int64
        )
        decoder_output[:, 0] = 1
        is_complete = paddle.zeros(shape=dataset.shape[0], dtype=paddle.bool)
        for n1 in range(self.seq_length):
            padding_mask = (
                paddle.equal(x=decoder_output[:, :-1], y=0)
                .unsqueeze(axis=1)
                .unsqueeze(axis=1)
            )
            future_mask = paddle.triu(
                x=paddle.ones(shape=[self.seq_length, self.seq_length]), diagonal=1
            ).astype(dtype=paddle.bool)
            mask_dec = paddle.logical_or(x=padding_mask, y=future_mask)
            y_dec = self.decoder(
                x_target=decoder_output[:, :-1],
                x_enc=encoder_output,
                mask=mask_dec,
            )
            y_mlp = self.last_layer(y_dec)
            # set value depending on complete condition
            decoder_output[:, n1 + 1] = paddle.where(
                is_complete, 0, paddle.argmax(y_mlp[:, n1], axis=-1)
            )
            # set complete condition
            for n2 in range(dataset.shape[0]):
                if complete_func(decoder_output[n2, 1:]):
                    is_complete[n2] = True
        return decoder_output

decode_process(dataset, complete_func)

Greedy decode with the Transformer model, decode until the equation tree is completed.

Parameters:

Name	Type	Description	Default
dataset	Tensor	Tabular dataset.	required
complete_func	Callable	Function used to calculate whether inference is complete.	required

Source code in ppsci/arch/transformer.py

@paddle.no_grad()
def decode_process(
    self, dataset: paddle.Tensor, complete_func: Callable
) -> paddle.Tensor:
    """Greedy decode with the Transformer model, decode until the equation tree is completed.

    Args:
        dataset (paddle.Tensor): Tabular dataset.
        complete_func (Callable): Function used to calculate whether inference is complete.
    """
    encoder_output = self.encoder(dataset)
    decoder_output = paddle.zeros(
        shape=(dataset.shape[0], self.seq_length + 1), dtype=paddle.int64
    )
    decoder_output[:, 0] = 1
    is_complete = paddle.zeros(shape=dataset.shape[0], dtype=paddle.bool)
    for n1 in range(self.seq_length):
        padding_mask = (
            paddle.equal(x=decoder_output[:, :-1], y=0)
            .unsqueeze(axis=1)
            .unsqueeze(axis=1)
        )
        future_mask = paddle.triu(
            x=paddle.ones(shape=[self.seq_length, self.seq_length]), diagonal=1
        ).astype(dtype=paddle.bool)
        mask_dec = paddle.logical_or(x=padding_mask, y=future_mask)
        y_dec = self.decoder(
            x_target=decoder_output[:, :-1],
            x_enc=encoder_output,
            mask=mask_dec,
        )
        y_mlp = self.last_layer(y_dec)
        # set value depending on complete condition
        decoder_output[:, n1 + 1] = paddle.where(
            is_complete, 0, paddle.argmax(y_mlp[:, n1], axis=-1)
        )
        # set complete condition
        for n2 in range(dataset.shape[0]):
            if complete_func(decoder_output[n2, 1:]):
                is_complete[n2] = True
    return decoder_output

Transolver

Bases: Arch

Source code in ppsci/arch/transolver.py

class Transolver(base.Arch):
    def __init__(
        self,
        input_keys: Tuple[str, ...],
        output_keys: Tuple[str, ...],
        space_dim: int = 1,
        n_layers: int = 5,
        n_hidden: int = 256,
        dropout: int = 0,
        n_head: int = 8,
        act: str = "gelu",
        mlp_ratio: int = 1,
        fun_dim: int = 1,
        out_dim: Union[int, List[int]] = 1,
        slice_num: int = 32,
        ref: int = 8,
        unified_pos: bool = False,
    ):
        super().__init__()
        self.input_keys = input_keys
        self.output_keys = output_keys
        if isinstance(out_dim, int):
            out_dim = [out_dim]
        out_dim = list(out_dim)
        self.out_dim = out_dim
        self.ref = ref
        self.unified_pos = unified_pos
        if self.unified_pos:
            self.preprocess = MLP(
                fun_dim + self.ref * self.ref * self.ref,
                n_hidden * 2,
                n_hidden,
                n_layers=0,
                res=False,
                act=act,
            )
        else:
            self.preprocess = MLP(
                fun_dim + space_dim,
                n_hidden * 2,
                n_hidden,
                n_layers=0,
                res=False,
                act=act,
            )
        self.n_hidden = n_hidden
        self.space_dim = space_dim
        self.blocks = nn.LayerList(
            sublayers=[
                TransolverBlock(
                    num_heads=n_head,
                    hidden_dim=n_hidden,
                    dropout=dropout,
                    act=act,
                    mlp_ratio=mlp_ratio,
                    out_dim=sum(out_dim),
                    slice_num=slice_num,
                    last_layer=(i == n_layers - 1),
                )
                for i in range(n_layers)
            ]
        )
        self.initialize_weights()
        self.placeholder = nn.Parameter(
            1 / n_hidden * paddle.rand(shape=[n_hidden], dtype=paddle.float32)
        )
        x = paddle.linspace(-1.5, 1.5, self.ref)
        y = paddle.linspace(0, 2, self.ref)
        z = paddle.linspace(-4, 4, self.ref)
        gridx, gridy, gridz = paddle.meshgrid(x, y, z, indexing="ij")
        grid_ref = paddle.stack([gridx, gridy, gridz], axis=-1).reshape(
            [1, self.ref**3, 3]
        )
        self.register_buffer("grid_ref", grid_ref)

    def initialize_weights(self):
        self.apply(self._init_weights)

    def _init_weights(self, m):
        if isinstance(m, nn.Linear):
            nn.init.trunc_normal_(m.weight, std=0.02)
            if m.bias is not None:
                nn.init.constant_(m.bias, 0)
        elif isinstance(m, (nn.LayerNorm, nn.BatchNorm1D)):
            nn.init.constant_(m.weight, 1.0)
            nn.init.constant_(m.bias, 0)

    def get_grid(self, my_pos):
        pos = paddle.sqrt(
            # [B, X, 1, 3] - [1, 1, R**3, 3] --> [B, X, R**3, 3]
            paddle.sum(
                (my_pos[:, :, None, :] - self.grid_ref[:, None, :, :]) ** 2, axis=-1
            )
        )

        return pos  # [B, X, R³]

    def forward(self, input_dict):
        # [B, N, C]
        x = input_dict[self.input_keys[0]]

        if self.unified_pos:
            # [B, N, C]
            pos = input_dict[self.input_keys[1]]
            # [B, N, R³]
            new_pos = self.get_grid(pos)
            # [B, N, C+R³]
            x = paddle.concat((x, new_pos), axis=-1)

        # [B, N, C]
        y = self.preprocess(x)
        y = y + self.placeholder[None, None, :]

        for block in self.blocks:
            y: paddle.Tensor = block(y)

        # [B, N, C]
        outputs = y.split(self.out_dim, axis=-1)
        return {k: v for k, v in zip(self.output_keys, outputs)}

UNetEx

Bases: Arch

U-Net Extension for CFD.

Reference: Ribeiro M D, Rehman A, Ahmed S, et al. DeepCFD: Efficient steady-state laminar flow approximation with deep convolutional neural networks[J]. arXiv preprint arXiv:2004.08826, 2020.

Parameters:

Name	Type	Description	Default
input_key	str	Name of function data for input.	required
output_key	str	Name of function data for output.	required
in_channel	int	Number of channels of input.	required
out_channel	int	Number of channels of output.	required
kernel_size	int	Size of kernel of convolution layer. Defaults to 3.	3
filters	Tuple[int, ...]	Number of filters. Defaults to (16, 32, 64).	(16, 32, 64)
layers	int	Number of encoders or decoders. Defaults to 3.	3
weight_norm	bool	Whether use weight normalization layer. Defaults to True.	True
batch_norm	bool	Whether add batch normalization layer. Defaults to True.	True
activation	Type[Layer]	Name of activation function. Defaults to nn.ReLU.	ReLU
final_activation	Optional[Type[Layer]]	Name of final activation function. Defaults to None.	None

Examples:

>>> import ppsci
>>> model = ppsci.arch.UNetEx(
...     input_key="input",
...     output_key="output",
...     in_channel=3,
...     out_channel=3,
...     kernel_size=5,
...     filters=(4, 4, 4, 4),
...     layers=3,
...     weight_norm=False,
...     batch_norm=False,
...     activation=None,
...     final_activation=None,
... )
>>> input_dict = {'input': paddle.rand([4, 3, 4, 4])}
>>> output_dict = model(input_dict)
>>> print(output_dict['output'])
>>> print(output_dict['output'].shape)
[4, 3, 4, 4]

Source code in ppsci/arch/unetex.py

class UNetEx(base.Arch):
    """U-Net Extension for CFD.

    Reference: [Ribeiro M D, Rehman A, Ahmed S, et al. DeepCFD: Efficient steady-state laminar flow approximation with deep convolutional neural networks[J]. arXiv preprint arXiv:2004.08826, 2020.](https://arxiv.org/abs/2004.08826)

    Args:
        input_key (str): Name of function data for input.
        output_key (str): Name of function data for output.
        in_channel (int): Number of channels of input.
        out_channel (int): Number of channels of output.
        kernel_size (int, optional): Size of kernel of convolution layer. Defaults to 3.
        filters (Tuple[int, ...], optional): Number of filters. Defaults to (16, 32, 64).
        layers (int, optional): Number of encoders or decoders. Defaults to 3.
        weight_norm (bool, optional): Whether use weight normalization layer. Defaults to True.
        batch_norm (bool, optional): Whether add batch normalization layer. Defaults to True.
        activation (Type[nn.Layer], optional): Name of activation function. Defaults to nn.ReLU.
        final_activation (Optional[Type[nn.Layer]]): Name of final activation function. Defaults to None.

    Examples:
        >>> import ppsci
        >>> model = ppsci.arch.UNetEx(
        ...     input_key="input",
        ...     output_key="output",
        ...     in_channel=3,
        ...     out_channel=3,
        ...     kernel_size=5,
        ...     filters=(4, 4, 4, 4),
        ...     layers=3,
        ...     weight_norm=False,
        ...     batch_norm=False,
        ...     activation=None,
        ...     final_activation=None,
        ... )
        >>> input_dict = {'input': paddle.rand([4, 3, 4, 4])}
        >>> output_dict = model(input_dict)
        >>> print(output_dict['output']) # doctest: +SKIP
        >>> print(output_dict['output'].shape)
        [4, 3, 4, 4]
    """

    def __init__(
        self,
        input_key: str,
        output_key: str,
        in_channel: int,
        out_channel: int,
        kernel_size: int = 3,
        filters: Tuple[int, ...] = (16, 32, 64),
        layers: int = 3,
        weight_norm: bool = True,
        batch_norm: bool = True,
        activation: Type[nn.Layer] = nn.ReLU,
        final_activation: Optional[Type[nn.Layer]] = None,
    ):
        if len(filters) == 0:
            raise ValueError("The filters shouldn't be empty ")

        super().__init__()
        self.input_keys = (input_key,)
        self.output_keys = (output_key,)
        self.final_activation = final_activation
        self.encoder = create_encoder(
            in_channel,
            filters,
            kernel_size,
            weight_norm,
            batch_norm,
            activation,
            layers,
        )
        decoders = [
            create_decoder(
                1, filters, kernel_size, weight_norm, batch_norm, activation, layers
            )
            for i in range(out_channel)
        ]
        self.decoders = nn.Sequential(*decoders)

    def encode(self, x):
        tensors = []
        indices = []
        sizes = []
        for encoder in self.encoder:
            x = encoder(x)
            sizes.append(x.shape)
            tensors.append(x)
            x, ind = nn.functional.max_pool2d(x, 2, 2, return_mask=True)
            indices.append(ind)
        return x, tensors, indices, sizes

    def decode(self, x, tensors, indices, sizes):
        y = []
        for _decoder in self.decoders:
            _x = x
            _tensors = tensors[:]
            _indices = indices[:]
            _sizes = sizes[:]
            for decoder in _decoder:
                tensor = _tensors.pop()
                size = _sizes.pop()
                indice = _indices.pop()
                # upsample operations
                _x = nn.functional.max_unpool2d(_x, indice, 2, 2, output_size=size)
                _x = paddle.concat([tensor, _x], axis=1)
                _x = decoder(_x)
            y.append(_x)
        return paddle.concat(y, axis=1)

    def forward(self, x):
        x = x[self.input_keys[0]]
        x, tensors, indices, sizes = self.encode(x)
        x = self.decode(x, tensors, indices, sizes)
        if self.final_activation is not None:
            x = self.final_activation(x)
        return {self.output_keys[0]: x}

UNONet

Bases: Arch

N-Dimensional U-Shaped Neural Operator.

Parameters:

Name	Type	Description	Default
input_keys	Tuple[str, ...]	Name of input keys, such as ("input",).	required
output_keys	Tuple[str, ...]	Name of output keys, such as ("output",).	required
in_channels	int	Number of input channels.	required
out_channels	int	Number of output channels.	required
hidden_channels	int	Width of the FNO (i.e. number of channels).	required
lifting_channels	int	Number of hidden channels of the lifting block of the FNO. Defaults to 256.	256
projection_channels	int	Number of hidden channels of the projection block of the FNO. Defaults to 256.	256
n_layers	int	Number of Fourier Layers. Defaults to 4.	4
uno_out_channels	Tuple[int, ...]	Number of output channel of each Fourier Layers. Eaxmple: For a Five layer UNO uno_out_channels can be [32,64,64,64,32].c	None
uno_n_modes	Tuple[Tuple[int, ...], ...]	Number of Fourier Modes to use in integral operation of each Fourier Layers (along each dimension). Example: For a five layer UNO with 2D input the uno_n_modes can be: [[5,5],[5,5],[5,5],[5,5],[5,5]]. Defaults to None.	None
uno_scalings	Tuple[Tuple[int, ...], ...]	Scaling Factors for each Fourier Layers. Example: For a five layer UNO with 2D input, the uno_scalings can be : [[1.0,1.0],[0.5,0.5],[1,1],[1,1],[2,2]].Defaults to None.	None
horizontal_skips_map	Dict	A map {...., b: a, ....} denoting horizontal skip connection from a-th layer to b-th layer. If None default skip connection is applied. Example: For a 5 layer UNO architecture, the skip connections can be horizontal_skips_map ={4:0,3:1}.Defaults to None.	None
incremental_n_modes	(tuple[int], optional)	Incremental number of modes to use in Fourier domain. * If not None, this allows to incrementally increase the number of modes in Fourier domain during training. Has to verify n <= N for (n, m) in zip(incremental_n_modes, n_modes). * If None, all the n_modes are used. This can be updated dynamically during training.Defaults to None.	None
use_mlp	bool	Whether to use an MLP layer after each FNO block. Defaults to False.	False
mlp	Dict[str, float]	Parameters of the MLP. {'expansion': float, 'dropout': float}. Defaults to None.	None
non_linearity	functional	Non-Linearity module to use. Defaults to F.gelu.	gelu
norm	str	Normalization layer to use. Defaults to None.	None
ada_in_features	(Optional[int], optional)	The input channels of the adaptive normalization.Defaults to None.	None
preactivation	bool	Whether to use resnet-style preactivation. Defaults to False.	False
fno_skip	str	Type of skip connection to use for fno_block. Defaults to "linear".	'linear'
horizontal_skip	str	Type of skip connection to use for horizontal skip. Defaults to "linear".	'linear'
mlp_skip	str	Type of skip connection to use for mlp. Defaults to "soft-gating".	'soft-gating'
separable	bool	Whether to use a depthwise separable spectral convolution. Defaults to False.	False
factorization	str	Tensor factorization of the parameters weight to use. * If None, a dense tensor parametrizes the Spectral convolutions. * Otherwise, the specified tensor factorization is used. Defaults to "Tucker".	None
rank	float	Rank of the tensor factorization of the Fourier weights. Defaults to 1.0.	1.0
joint_factorization	bool	Whether all the Fourier Layers should be parametrized by a single tensor (vs one per layer). Defaults to False.	False
implementation	str	{'factorized', 'reconstructed'}, optional. Defaults to "factorized". If factorization is not None, forward mode to use:: * reconstructed : the full weight tensor is reconstructed from the factorization and used for the forward pass. * factorized : the input is directly contracted with the factors of the decomposition.	'factorized'
domain_padding	Optional[Union[list, float, int]]	Whether to use percentage of padding. Defaults to None.	None
domain_padding_mode	str	{'symmetric', 'one-sided'}, optional How to perform domain padding, by default 'one-sided'. Defaults to "one-sided".	'one-sided'
fft_norm	str	The normalization mode for the FFT. Defaults to "forward".	'forward'
patching_levels	int	Number of patching levels to use. Defaults to 0.	0

Source code in ppsci/arch/unonet.py

class UNONet(base.Arch):
    """N-Dimensional U-Shaped Neural Operator.

    Args:
        input_keys (Tuple[str, ...]): Name of input keys, such as ("input",).
        output_keys (Tuple[str, ...]): Name of output keys, such as ("output",).
        in_channels (int, optional): Number of input channels.
        out_channels (int, optional): Number of output channels.
        hidden_channels (int): Width of the FNO (i.e. number of channels).
        lifting_channels (int, optional): Number of hidden channels of the lifting block of the FNO.
            Defaults to 256.
        projection_channels (int, optional): Number of hidden channels of the projection block of the FNO.
            Defaults to 256.
        n_layers (int, optional): Number of Fourier Layers. Defaults to 4.
        uno_out_channels (Tuple[int, ...], optional): Number of output channel of each Fourier Layers.
            Eaxmple: For a Five layer UNO uno_out_channels can be [32,64,64,64,32].c
        uno_n_modes (Tuple[Tuple[int, ...], ...]): Number of Fourier Modes to use in integral operation of each
            Fourier Layers (along each dimension).
            Example: For a five layer UNO with 2D input the uno_n_modes can be: [[5,5],[5,5],[5,5],[5,5],[5,5]]. Defaults to None.
        uno_scalings (Tuple[Tuple[int, ...], ...]): Scaling Factors for each Fourier Layers.
            Example: For a five layer UNO with 2D input, the uno_scalings can be : [[1.0,1.0],[0.5,0.5],[1,1],[1,1],[2,2]].Defaults to None.
        horizontal_skips_map (Dict, optional): A map {...., b: a, ....} denoting horizontal skip connection
            from a-th layer to b-th layer. If None default skip connection is applied.
            Example: For a 5 layer UNO architecture, the skip connections can be horizontal_skips_map ={4:0,3:1}.Defaults to None.
        incremental_n_modes (tuple[int],optional): Incremental number of modes to use in Fourier domain.
            * If not None, this allows to incrementally increase the number of modes in Fourier domain
            during training. Has to verify n <= N for (n, m) in zip(incremental_n_modes, n_modes).
            * If None, all the n_modes are used.
            This can be updated dynamically during training.Defaults to None.
        use_mlp (bool, optional): Whether to use an MLP layer after each FNO block. Defaults to False.
        mlp (Dict[str, float], optional): Parameters of the MLP. {'expansion': float, 'dropout': float}.
            Defaults to None.
        non_linearity (nn.functional, optional): Non-Linearity module to use. Defaults to F.gelu.
        norm (str, optional): Normalization layer to use. Defaults to None.
        ada_in_features (Optional[int],optional): The input channels of the adaptive normalization.Defaults to
            None.
        preactivation (bool, optional): Whether to use resnet-style preactivation. Defaults to False.
        fno_skip (str, optional): Type of skip connection to use for fno_block. Defaults to "linear".
        horizontal_skip (str, optional): Type of skip connection to use for horizontal skip. Defaults to
            "linear".
        mlp_skip (str, optional): Type of skip connection to use for mlp. Defaults to "soft-gating".
        separable (bool, optional): Whether to use a depthwise separable spectral convolution.
            Defaults to  False.
        factorization (str, optional): Tensor factorization of the parameters weight to use.
            * If None, a dense tensor parametrizes the Spectral convolutions.
            * Otherwise, the specified tensor factorization is used. Defaults to "Tucker".
        rank (float, optional): Rank of the tensor factorization of the Fourier weights. Defaults to 1.0.
        joint_factorization (bool, optional): Whether all the Fourier Layers should be parametrized by a
            single tensor (vs one per layer). Defaults to False.
        implementation (str, optional): {'factorized', 'reconstructed'}, optional. Defaults to "factorized".
            If factorization is not None, forward mode to use::
            * `reconstructed` : the full weight tensor is reconstructed from the factorization and used for the forward pass.
            * `factorized` : the input is directly contracted with the factors of the decomposition.
        domain_padding (Optional[Union[list, float, int]], optional): Whether to use percentage of padding.
            Defaults to None.
        domain_padding_mode (str, optional): {'symmetric', 'one-sided'}, optional
            How to perform domain padding, by default 'one-sided'. Defaults to "one-sided".
        fft_norm (str, optional): The normalization mode for the FFT. Defaults to "forward".
        patching_levels (int, optional): Number of patching levels to use. Defaults to 0.
    """

    def __init__(
        self,
        input_keys: Tuple[str, ...],
        output_keys: Tuple[str, ...],
        in_channels: int,
        out_channels: int,
        hidden_channels: int,
        lifting_channels: int = 256,
        projection_channels: int = 256,
        n_layers: int = 4,
        uno_out_channels: Tuple[int, ...] = None,
        uno_n_modes: Tuple[Tuple[int, ...], ...] = None,
        uno_scalings: Tuple[Tuple[int, ...], ...] = None,
        horizontal_skips_map: Dict = None,
        incremental_n_modes: Tuple[int, ...] = None,
        use_mlp: bool = False,
        mlp: Optional[Dict[str, float]] = None,
        non_linearity: nn.functional = F.gelu,
        norm: str = None,
        ada_in_features: Optional[int] = None,
        preactivation: bool = False,
        fno_skip: str = "linear",
        horizontal_skip: str = "linear",
        mlp_skip: str = "soft-gating",
        separable: bool = False,
        factorization: str = None,
        rank: float = 1.0,
        joint_factorization: bool = False,
        implementation: str = "factorized",
        domain_padding: Optional[Union[list, float, int]] = None,
        domain_padding_mode: str = "one-sided",
        fft_norm: str = "forward",
        patching_levels: int = 0,
        **kwargs,
    ):
        super().__init__()
        self.input_keys = input_keys
        self.output_keys = output_keys
        if uno_out_channels is None:
            raise ValueError("uno_out_channels can not be None")
        if uno_n_modes is None:
            raise ValueError("uno_n_modes can not be None")
        if uno_scalings is None:
            raise ValueError("uno_scalings can not be None")

        if len(uno_out_channels) != n_layers:
            raise ValueError("Output channels for all layers are not given")

        if len(uno_n_modes) != n_layers:
            raise ValueError("Number of modes for all layers are not given")

        if len(uno_scalings) != n_layers:
            raise ValueError("Scaling factor for all layers are not given")

        self.n_dim = len(uno_n_modes[0])
        self.uno_out_channels = uno_out_channels
        self.uno_n_modes = uno_n_modes
        self.uno_scalings = uno_scalings

        self.hidden_channels = hidden_channels
        self.lifting_channels = lifting_channels
        self.projection_channels = projection_channels
        self.in_channels = in_channels
        if patching_levels:
            self.in_channels = self.in_channels * patching_levels + 1
        self.out_channels = out_channels
        self.n_layers = n_layers
        self.horizontal_skips_map = horizontal_skips_map
        self.joint_factorization = joint_factorization
        self.non_linearity = non_linearity
        self.rank = rank
        self.factorization = factorization
        self.fno_skip = (fno_skip,)
        self.mlp_skip = (mlp_skip,)
        self.fft_norm = fft_norm
        self.implementation = implementation
        self.separable = separable
        self.preactivation = preactivation
        self._incremental_n_modes = incremental_n_modes
        self.mlp = mlp
        # constructing default skip maps
        if self.horizontal_skips_map is None:
            self.horizontal_skips_map = {}
            for i in range(
                0,
                n_layers // 2,
            ):
                # example, if n_layers = 5, then 4:0, 3:1
                self.horizontal_skips_map[n_layers - i - 1] = i
        # self.uno_scalings may be a 1d list specifying uniform scaling factor at each layer
        # or a 2d list, where each row specifies scaling factors along each dimension.
        # To get the final (end to end) scaling factors we need to multiply
        # the scaling factors (a list) of all layer.

        self.end_to_end_scaling_factor = [1] * len(self.uno_scalings[0])
        # multiplying scaling factors
        for k in self.uno_scalings:
            self.end_to_end_scaling_factor = [
                i * j for (i, j) in zip(self.end_to_end_scaling_factor, k)
            ]

        # list with a single element is replaced by the scaler.
        if len(self.end_to_end_scaling_factor) == 1:
            self.end_to_end_scaling_factor = self.end_to_end_scaling_factor[0]

        if isinstance(self.end_to_end_scaling_factor, (float, int)):
            self.end_to_end_scaling_factor = [
                self.end_to_end_scaling_factor
            ] * self.n_dim

        if domain_padding is not None and (
            (isinstance(domain_padding, list) and sum(domain_padding) > 0)
            or (isinstance(domain_padding, (float, int)) and domain_padding > 0)
        ):
            self.domain_padding = fno_block.DomainPadding(
                domain_padding=domain_padding, padding_mode=domain_padding_mode
            )
        else:
            self.domain_padding = None
        self.domain_padding_mode = domain_padding_mode

        self.lifting = fno_block.MLP(
            in_channels=in_channels,
            out_channels=self.hidden_channels,
            hidden_channels=self.lifting_channels,
            n_layers=2,
            n_dim=self.n_dim,
        )

        self.fno_blocks = nn.LayerList([])
        self.horizontal_skips = nn.LayerDict({})
        prev_out = self.hidden_channels
        for i in range(self.n_layers):
            if i in self.horizontal_skips_map.keys():
                prev_out = (
                    prev_out + self.uno_out_channels[self.horizontal_skips_map[i]]
                )
            self.fno_blocks.append(
                fno_block.FNOBlocks(
                    in_channels=prev_out,
                    out_channels=self.uno_out_channels[i],
                    n_modes=self.uno_n_modes[i],
                    use_mlp=use_mlp,
                    mlp=mlp,
                    output_scaling_factor=[self.uno_scalings[i]],
                    non_linearity=non_linearity,
                    norm=norm,
                    ada_in_features=ada_in_features,
                    preactivation=preactivation,
                    fno_skip=fno_skip,
                    mlp_skip=mlp_skip,
                    separable=separable,
                    incremental_n_modes=incremental_n_modes,
                    factorization=factorization,
                    rank=rank,
                    SpectralConv=fno_block.FactorizedSpectralConv,
                    joint_factorization=joint_factorization,
                    implementation=implementation,
                    fft_norm=fft_norm,
                )
            )

            if i in self.horizontal_skips_map.values():
                self.horizontal_skips[str(i)] = fno_block.skip_connection(
                    self.uno_out_channels[i],
                    self.uno_out_channels[i],
                    type=horizontal_skip,
                    n_dim=self.n_dim,
                )
            prev_out = self.uno_out_channels[i]

        self.projection = fno_block.MLP(
            in_channels=prev_out,
            out_channels=out_channels,
            hidden_channels=self.projection_channels,
            n_layers=2,
            n_dim=self.n_dim,
            non_linearity=non_linearity,
        )

    def forward(self, x, **kwargs):
        x = self.concat_to_tensor(x, self.input_keys)
        x = self.lifting(x)
        if self.domain_padding is not None:
            x = self.domain_padding.pad(x)
        output_shape = [
            int(round(i * j))
            for (i, j) in zip(x.shape[-self.n_dim :], self.end_to_end_scaling_factor)
        ]

        skip_outputs = {}
        cur_output = None
        for layer_idx in range(self.n_layers):
            if layer_idx in self.horizontal_skips_map.keys():
                skip_val = skip_outputs[self.horizontal_skips_map[layer_idx]]
                output_scaling_factors = [
                    m / n for (m, n) in zip(x.shape, skip_val.shape)
                ]
                output_scaling_factors = output_scaling_factors[-1 * self.n_dim :]
                t = fno_block.resample(
                    skip_val, output_scaling_factors, list(range(-self.n_dim, 0))
                )
                x = paddle.concat([x, t], axis=1)

            if layer_idx == self.n_layers - 1:
                cur_output = output_shape
            x = self.fno_blocks[layer_idx](x, output_shape=cur_output)
            if layer_idx in self.horizontal_skips_map.values():
                skip_outputs[layer_idx] = self.horizontal_skips[str(layer_idx)](x)

        if self.domain_padding is not None:
            x = self.domain_padding.unpad(x)

        out = self.projection(x)
        return {self.output_keys[0]: out}

USCNN

Bases: Arch

Physics-informed convolutional neural networks.

Parameters:

Name	Type	Description	Default
input_keys	Tuple[str, ...]	Name of input keys, such as ("coords").	required
output_keys	Tuple[str, ...]	Name of output keys, such as ("outputV").	required
hidden_size	Union[int, Tuple[int, ...]]	The hidden channel for convolutional layers	required
h	float	The spatial step	required
nx	int	the number of grids along x-axis	required
ny	int	The number of grids along y-axis	required
nvar_in	int	input channel. Defaults to 1.	1
nvar_out	int	Output channel. Defaults to 1.	1
pad_singleside	int	Pad for hard boundary constraint. Defaults to 1.	1
k	int	Kernel_size. Defaults to 5.	5
s	int	Stride. Defaults to 1.	1
p	int	Padding. Defaults to 2.	2

Examples:

>>> import ppsci
>>> model = ppsci.arch.USCNN(
...     ["coords"],
...     ["outputV"],
...     [16, 32, 16],
...     h=0.01,
...     ny=19,
...     nx=84,
...     nvar_in=2,
...     nvar_out=1,
...     pad_singleside=1,
... )

Source code in ppsci/arch/uscnn.py

class USCNN(base.Arch):
    """Physics-informed convolutional neural networks.

    Args:
        input_keys (Tuple[str, ...]): Name of input keys, such as ("coords").
        output_keys (Tuple[str, ...]):Name of output keys, such as ("outputV").
        hidden_size (Union[int, Tuple[int, ...]]): The hidden channel for convolutional layers
        h (float): The spatial step
        nx (int):  the number of grids along x-axis
        ny (int): The number of grids along y-axis
        nvar_in (int, optional):  input channel. Defaults to 1.
        nvar_out (int, optional): Output channel. Defaults to 1.
        pad_singleside (int, optional): Pad for hard boundary constraint. Defaults to 1.
        k (int, optional): Kernel_size. Defaults to 5.
        s (int, optional): Stride. Defaults to 1.
        p (int, optional): Padding. Defaults to 2.

    Examples:
        >>> import ppsci
        >>> model = ppsci.arch.USCNN(
        ...     ["coords"],
        ...     ["outputV"],
        ...     [16, 32, 16],
        ...     h=0.01,
        ...     ny=19,
        ...     nx=84,
        ...     nvar_in=2,
        ...     nvar_out=1,
        ...     pad_singleside=1,
        ... )
    """

    def __init__(
        self,
        input_keys: Tuple[str, ...],
        output_keys: Tuple[str, ...],
        hidden_size: Union[int, Tuple[int, ...]],
        h: float,
        nx: int,
        ny: int,
        nvar_in: int = 1,
        nvar_out: int = 1,
        pad_singleside: int = 1,
        k: int = 5,
        s: int = 1,
        p: int = 2,
    ):
        super().__init__()
        self.input_keys = input_keys
        self.output_keys = output_keys
        self.nvar_in = nvar_in
        self.nvar_out = nvar_out
        self.k = k
        self.s = s
        self.p = p
        self.deltaX = h
        self.nx = nx
        self.ny = ny
        self.pad_singleside = pad_singleside
        self.relu = nn.ReLU()
        self.US = nn.Upsample(size=[self.ny - 2, self.nx - 2], mode="bicubic")
        self.conv1 = nn.Conv2D(
            self.nvar_in, hidden_size[0], kernel_size=k, stride=s, padding=p
        )
        self.conv2 = nn.Conv2D(
            hidden_size[0], hidden_size[1], kernel_size=k, stride=s, padding=p
        )
        self.conv3 = nn.Conv2D(
            hidden_size[1], hidden_size[2], kernel_size=k, stride=s, padding=p
        )
        self.conv4 = nn.Conv2D(
            hidden_size[2], self.nvar_out, kernel_size=k, stride=s, padding=p
        )
        self.pixel_shuffle = nn.PixelShuffle(1)
        self.apply(self.init_weights)
        self.udfpad = nn.Pad2D(
            [pad_singleside, pad_singleside, pad_singleside, pad_singleside], value=0
        )

    def init_weights(self, m):
        if isinstance(m, nn.Conv2D):
            bound = 1 / np.sqrt(np.prod(m.weight.shape[1:]))
            ppsci.utils.initializer.uniform_(m.weight, -bound, bound)
            if m.bias is not None:
                ppsci.utils.initializer.uniform_(m.bias, -bound, bound)

    def forward(self, x):
        y = self.concat_to_tensor(x, self.input_keys, axis=-1)
        y = self.US(y)
        y = self.relu(self.conv1(y))
        y = self.relu(self.conv2(y))
        y = self.relu(self.conv3(y))
        y = self.pixel_shuffle(self.conv4(y))

        y = self.udfpad(y)
        y = y[:, 0, :, :].reshape([y.shape[0], 1, y.shape[2], y.shape[3]])
        y = self.split_to_dict(y, self.output_keys)
        if self._output_transform is not None:
            y = self._output_transform(x, y)
        return y