Clifford models

We provide exemplary 2D and 3D Clifford models as used in the paper.

All these modules are available for different algebras.

2D models

The following code snippet initializes a 2D Clifford ResNet.

import torch.nn.functional as F

from cliffordlayers.models.basic.twod import (
    CliffordFluidNet2d,
    CliffordBasicBlock2d,
)

model = CliffordFluidNet2d(
        g = [-1, -1],
        block = CliffordBasicBlock2d,
        num_blocks = [2, 2, 2, 2],
        in_channels = in_channels,
        out_channels = out_channels,
        hidden_channels = 32,
        activation = F.gelu,
        norm = True,
        rotation = False,
    )

The following code snippet initializes a 2D rotational Clifford ResNet.

import torch.nn.functional as F

from cliffordlayers.models.basic.twod import (
    CliffordFluidNet2d,
    CliffordBasicBlock2d,
)

model = CliffordNet2d(
        g = [-1, -1],
        block = CliffordBasicBlock2d,
        num_blocks = [2, 2, 2, 2],
        in_channels = in_channels,
        out_channels = out_channels,
        hidden_channels = 32,
        activation = F.gelu,
        norm = True,
        rotation = True,
    )

The following code snippet initializes a 2D Clifford FNO.

import torch.nn.functional as F

from cliffordlayers.models.utils import partialclass
from cliffordlayers.models.basic.twod import (
    CliffordFluidNet2d,
    CliffordFourierBasicBlock2d,
)

model = CliffordFluidNet2d(
        g = [-1, -1],
        block = partialclass(
                "CliffordFourierBasicBlock2d", CliffordFourierBasicBlock2d, modes1=32, modes2=32
            ),
        num_blocks = [1, 1, 1, 1],
        in_channels = in_channels,
        out_channels = out_channels,
        hidden_channels = 32,
        activation = F.gelu,
        norm = False,
        rotation = False,
    )

CliffordBasicBlock2d

Bases: Module

2D building block for Clifford ResNet architectures.

Parameters:

Name	Type	Description	Default
g	Union[tuple, list, Tensor]	Signature of Clifford algebra.	required
in_channels	int	Number of input channels.	required
out_channels	int	Number of output channels.	required
activation	Callable	Activation function. Defaults to F.gelu.	gelu
kernel_size	int	Kernel size of Clifford convolution. Defaults to 3.	3
stride	int	Stride of Clifford convolution. Defaults to 1.	1
padding	int	Padding of Clifford convolution. Defaults to 1.	1
rotation	bool	Wether to use rotational Clifford convolution. Defaults to False.	False
norm	bool	Wether to use Clifford (group) normalization. Defaults to False.	False
num_groups	int	Number of groups when using Clifford (group) normalization. Defaults to 1.	1

Source code in cliffordlayers/models/basic/twod.py

class CliffordBasicBlock2d(nn.Module):
    """2D building block for Clifford ResNet architectures.

    Args:
        g (Union[tuple, list, torch.Tensor]): Signature of Clifford algebra.
        in_channels (int): Number of input channels.
        out_channels (int): Number of output channels.
        activation (Callable, optional): Activation function. Defaults to F.gelu.
        kernel_size (int, optional): Kernel size of Clifford convolution. Defaults to 3.
        stride (int, optional): Stride of Clifford convolution. Defaults to 1.
        padding (int, optional): Padding of Clifford convolution. Defaults to 1.
        rotation (bool, optional): Wether to use rotational Clifford convolution. Defaults to False.
        norm (bool, optional): Wether to use Clifford (group) normalization. Defaults to False.
        num_groups (int, optional): Number of groups when using Clifford (group) normalization. Defaults to 1.
    """

    expansion: int = 1

    def __init__(
        self,
        g: Union[tuple, list, torch.Tensor],
        in_channels: int,
        out_channels: int,
        activation: Callable = F.gelu,
        kernel_size: int = 3,
        stride: int = 1,
        padding: int = 1,
        rotation: bool = False,
        norm: bool = False,
        num_groups: int = 1,
    ) -> None:
        super().__init__()
        self.conv1 = CliffordConv2d(
            g,
            in_channels,
            out_channels,
            kernel_size=kernel_size,
            stride=stride,
            padding=padding,
            bias=True,
            rotation=rotation,
        )
        self.conv2 = CliffordConv2d(
            g,
            out_channels,
            out_channels,
            kernel_size=kernel_size,
            stride=stride,
            padding=padding,
            bias=True,
            rotation=rotation,
        )
        self.norm1 = CliffordGroupNorm2d(g, num_groups, in_channels) if norm else nn.Identity()
        self.norm2 = CliffordGroupNorm2d(g, num_groups, out_channels) if norm else nn.Identity()
        self.activation = activation

    def __repr__(self):
        return "CliffordBasicBlock2d"

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        out = self.conv1(self.activation(self.norm1(x)))
        out = self.conv2(self.activation(self.norm2(out)))
        return out + x

CliffordFluidNet2d

Bases: Module

2D building block for Clifford architectures for fluid mechanics (vector field+scalar field) with ResNet backbone network. The backbone networks follows these three steps: 1. Clifford scalar+vector field encoding. 2. Basic blocks as provided. 3. Clifford scalar+vector field decoding.

Parameters:

Name	Type	Description	Default
g	Union[tuple, list, Tensor]	Signature of Clifford algebra.	required
block	Module	Choice of basic blocks.	required
num_blocks	list	List of basic blocks in each residual block.	required
in_channels	int	Number of input channels.	required
out_channels	int	Number of output channels.	required
activation	Callable	Activation function. Defaults to F.gelu.	required
rotation	bool	Wether to use rotational Clifford convolution. Defaults to False.	required
norm	bool	Wether to use Clifford (group) normalization. Defaults to False.	False
num_groups	int	Number of groups when using Clifford (group) normalization. Defaults to 1.	1

Source code in cliffordlayers/models/basic/twod.py

class CliffordFluidNet2d(nn.Module):
    """2D building block for Clifford architectures for fluid mechanics (vector field+scalar field)
    with ResNet backbone network. The backbone networks follows these three steps:
        1. Clifford scalar+vector field encoding.
        2. Basic blocks as provided.
        3. Clifford scalar+vector field decoding.

    Args:
        g (Union[tuple, list, torch.Tensor]): Signature of Clifford algebra.
        block (nn.Module): Choice of basic blocks.
        num_blocks (list): List of basic blocks in each residual block.
        in_channels (int): Number of input channels.
        out_channels (int): Number of output channels.
        activation (Callable, optional): Activation function. Defaults to F.gelu.
        rotation (bool, optional): Wether to use rotational Clifford convolution. Defaults to False.
        norm (bool, optional): Wether to use Clifford (group) normalization. Defaults to False.
        num_groups (int, optional): Number of groups when using Clifford (group) normalization. Defaults to 1.
    """

    # For periodic boundary conditions, set padding = 0.
    padding = 9

    def __init__(
        self,
        g: Union[tuple, list, torch.Tensor],
        block: nn.Module,
        num_blocks: list,
        in_channels: int,
        out_channels: int,
        hidden_channels: int,
        activation: Callable,
        rotation: False,
        norm: bool = False,
        num_groups: int = 1,
    ):
        super().__init__()

        self.activation = activation
        # Encoding and decoding layers
        self.encoder = CliffordConv2dScalarVectorEncoder(
            g,
            in_channels=in_channels,
            out_channels=hidden_channels,
            kernel_size=1,
            padding=0,
            rotation=rotation,
        )
        self.decoder = CliffordConv2dScalarVectorDecoder(
            g,
            in_channels=hidden_channels,
            out_channels=out_channels,
            kernel_size=1,
            padding=0,
            rotation=rotation,
        )

        # Residual blocks
        self.layers = nn.ModuleList(
            [
                self._make_basic_block(
                    g,
                    block,
                    hidden_channels,
                    num_blocks[i],
                    activation=activation,
                    rotation=rotation,
                    norm=norm,
                    num_groups=num_groups,
                )
                for i in range(len(num_blocks))
            ]
        )

    def _make_basic_block(
        self,
        g,
        block: nn.Module,
        hidden_channels: int,
        num_blocks: int,
        activation: Callable,
        rotation: bool,
        norm: bool,
        num_groups: int,
    ) -> nn.Sequential:
        blocks = []
        for _ in range(num_blocks):
            blocks.append(
                block(
                    g,
                    hidden_channels,
                    hidden_channels,
                    activation=activation,
                    rotation=rotation,
                    norm=norm,
                    num_groups=num_groups,
                )
            )
        return nn.Sequential(*blocks)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        assert x.dim() == 5

        # Encoding layer
        x = self.encoder(self.activation(x))

        # Embed for non-periodic boundaries
        if self.padding > 0:
            B_dim, C_dim, *D_dims, I_dim = range(len(x.shape))
            x = x.permute(B_dim, I_dim, C_dim, *D_dims)
            x = F.pad(x, [0, self.padding, 0, self.padding])
            B_dim, I_dim, C_dim, *D_dims = range(len(x.shape))
            x = x.permute(B_dim, C_dim, *D_dims, I_dim)

        # Apply residual layers
        for layer in self.layers:
            x = layer(x)

        # Decoding layer
        if self.padding > 0:
            B_dim, C_dim, *D_dims, I_dim = range(len(x.shape))
            x = x.permute(B_dim, I_dim, C_dim, *D_dims)
            x = x[..., : -self.padding, : -self.padding]
            B_dim, I_dim, C_dim, *D_dims = range(len(x.shape))
            x = x.permute(B_dim, C_dim, *D_dims, I_dim)

        # Output layer
        x = self.decoder(x)
        return x

CliffordFourierBasicBlock2d

Bases: Module

2D building block for Clifford FNO architectures.

Parameters:

Name	Type	Description	Default
g	Union[tuple, list, Tensor]	Signature of Clifford algebra.	required
in_channels	int	Number of input channels.	required
out_channels	int	Number of output channels.	required
activation	Callable	Activation function. Defaults to F.gelu.	gelu
kernel_size	int	Kernel size of Clifford convolution. Defaults to 3.	1
stride	int	Stride of Clifford convolution. Defaults to 1.	1
padding	int	Padding of Clifford convolution. Defaults to 1.	0
rotation	bool	Wether to use rotational Clifford convolution. Defaults to False.	False
norm	bool	Wether to use Clifford (group) normalization. Defaults to False.	False
num_groups	int	Number of groups when using Clifford (group) normalization. Defaults to 1.	1
modes1	int	Number of Fourier modes in the first dimension. Defaults to 16.	16
modes2	int	Number of Fourier modes in the second dimension. Defaults to 16.	16

Source code in cliffordlayers/models/basic/twod.py

class CliffordFourierBasicBlock2d(nn.Module):
    """2D building block for Clifford FNO architectures.

    Args:
        g (Union[tuple, list, torch.Tensor]): Signature of Clifford algebra.
        in_channels (int): Number of input channels.
        out_channels (int): Number of output channels.
        activation (Callable, optional): Activation function. Defaults to F.gelu.
        kernel_size (int, optional): Kernel size of Clifford convolution. Defaults to 3.
        stride (int, optional): Stride of Clifford convolution. Defaults to 1.
        padding (int, optional): Padding of Clifford convolution. Defaults to 1.
        rotation (bool, optional): Wether to use rotational Clifford convolution. Defaults to False.
        norm (bool, optional): Wether to use Clifford (group) normalization. Defaults to False.
        num_groups (int, optional): Number of groups when using Clifford (group) normalization. Defaults to 1.
        modes1 (int, optional): Number of Fourier modes in the first dimension. Defaults to 16.
        modes2 (int, optional): Number of Fourier modes in the second dimension. Defaults to 16.
    """

    expansion: int = 1

    def __init__(
        self,
        g: Union[tuple, list, torch.Tensor],
        in_channels: int,
        out_channels: int,
        activation: Callable = F.gelu,
        kernel_size: int = 1,
        stride: int = 1,
        padding: int = 0,
        rotation: bool = False,
        norm: bool = False,
        num_groups: int = 1,
        modes1: int = 16,
        modes2: int = 16,
    ):
        super().__init__()
        self.fourier = CliffordSpectralConv2d(
            g,
            in_channels,
            out_channels,
            modes1=modes1,
            modes2=modes2,
        )
        self.conv = CliffordConv2d(
            g,
            in_channels,
            out_channels,
            kernel_size=kernel_size,
            stride=stride,
            padding=padding,
            bias=True,
            rotation=rotation,
        )
        self.norm = CliffordGroupNorm2d(g, num_groups, out_channels) if norm else nn.Identity()
        self.activation = activation

    def __repr__(self):
        return "CliffordFourierBasicBlock2d"

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x1 = self.fourier(x)
        x2 = self.conv(x)
        return self.activation(self.norm(x1 + x2))

CliffordG3BasicBlock2d

Bases: Module

Basic block for G3 convolutions on 2D grids, comprising two G3 Clifford convolutional layers.

Parameters:

Name	Type	Description	Default
in_channels	int	Number of input channels.	required
out_channels	int	Number of output channels.	required
kernel_size	int	Size of the convolutional kernel. Defaults to 3.	3
stride	int	Stride of the convolution operation. Defaults to 1.	1
padding	int	Padding added to both sides of the input. Defaults to 1.	1
activation	str	Type of activation function. Defaults to "vlin".	'vlin'
norm	bool	If True, normalization is applied. Defaults to True.	True
num_groups	int	Number of groups for the group normalization. Defaults to 1.	1
prenorm	bool	If True, normalization is applied before activation, otherwise after. Defaults to True.	True

Source code in cliffordlayers/models/gca/twod.py

class CliffordG3BasicBlock2d(nn.Module):
    """
    Basic block for G3 convolutions on 2D grids, comprising two G3 Clifford convolutional layers.

    Args:
        in_channels (int): Number of input channels.
        out_channels (int): Number of output channels.
        kernel_size (int, optional): Size of the convolutional kernel. Defaults to 3.
        stride (int, optional): Stride of the convolution operation. Defaults to 1.
        padding (int, optional): Padding added to both sides of the input. Defaults to 1.
        activation (str, optional): Type of activation function. Defaults to "vlin".
        norm (bool, optional): If True, normalization is applied. Defaults to True.
        num_groups (int, optional): Number of groups for the group normalization. Defaults to 1.
        prenorm (bool, optional): If True, normalization is applied before activation, otherwise after. Defaults to True.
    """

    expansion: int = 1

    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: int = 3,
        stride: int = 1,
        padding: int = 1,
        activation: str = "vlin",
        norm: bool = True,
        num_groups: int = 1,
        prenorm: bool = True,
    ):
        super().__init__()
        self.conv1 = CliffordG3Conv2d(
            in_channels,
            out_channels,
            kernel_size=kernel_size,
            stride=stride,
            padding=padding,
            bias=True,
        )
        self.conv2 = CliffordG3Conv2d(
            out_channels,
            out_channels,
            kernel_size=kernel_size,
            stride=stride,
            padding=padding,
            bias=True,
        )

        self.norm1 = CliffordG3GroupNorm(num_groups, in_channels, 3) if norm else nn.Identity()
        self.norm2 = CliffordG3GroupNorm(num_groups, out_channels, 3) if norm else nn.Identity()

        if in_channels != out_channels:
            self.shortcut = CliffordG3Conv2d(
                in_channels,
                out_channels,
                kernel_size=1,
            )
        else:
            self.shortcut = nn.Identity()

        self.act1 = get_activation(activation, in_channels)
        self.act2 = get_activation(activation, out_channels)

        self.prenorm = prenorm

    def forward(self, x):
        if self.prenorm:
            out = self.conv1(self.act1(self.norm1(x)))
            out = self.conv2(self.act2(self.norm2(out)))
        else:
            out = self.conv1(self.norm1(self.act1(x)))
            out = self.conv2(self.norm2(self.act2(out)))

        return out + self.shortcut(x)

CliffordG3DownBlock

Bases: Module

UNet encoder block for G3 Clifford convolutions on 2D grids.

Parameters:

Name	Type	Description	Default
in_channels	int	Number of input channels.	required
out_channels	int	Number of output channels.	required
activation	str	Type of activation function.	required
norm	bool	If True, normalization is applied. Defaults to False.	False
prenorm	bool	If True, normalization is applied before activation, otherwise after. Defaults to True.	True
num_groups	int	Number of groups for the group normalization. Defaults to 1.	1

Source code in cliffordlayers/models/gca/twod.py

class CliffordG3DownBlock(nn.Module):
    """
    UNet encoder block for G3 Clifford convolutions on 2D grids.

    Args:
        in_channels (int): Number of input channels.
        out_channels (int): Number of output channels.
        activation (str): Type of activation function.
        norm (bool, optional): If True, normalization is applied. Defaults to False.
        prenorm (bool, optional): If True, normalization is applied before activation, otherwise after. Defaults to True.
        num_groups (int, optional): Number of groups for the group normalization. Defaults to 1.
    """

    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        activation: str,
        norm: bool = False,
        prenorm: bool = True,
        num_groups: int = 1,
    ):
        super().__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.block = CliffordG3BasicBlock2d(
            in_channels, out_channels, activation=activation, norm=norm, prenorm=prenorm, num_groups=num_groups
        )

    def forward(self, x):
        return self.block(x)

CliffordG3Downsample

Bases: Module

Scale down the two-dimensional G3 Clifford feature map by a half.

Parameters:

Name	Type	Description	Default
n_channels	int	Number of channels.	required

Source code in cliffordlayers/models/gca/twod.py

class CliffordG3Downsample(nn.Module):
    """
    Scale down the two-dimensional G3 Clifford feature map by a half.

    Args:
        n_channels (int): Number of channels.
    """

    def __init__(self, n_channels):
        super().__init__()
        self.n_channels = n_channels
        self.conv = CliffordG3Conv2d(
            n_channels,
            n_channels,
            kernel_size=3,
            stride=2,
            padding=1,
        )

    def forward(self, x):
        return self.conv(x)

CliffordG3MiddleBlock

Bases: Module

UNet middle block for G3 Clifford convolutions on 2D grids.

Parameters:

Name	Type	Description	Default
n_channels	int	Number of channels.	required
activation	str	Type of activation function.	required
norm	bool	If True, normalization is applied. Defaults to False.	False
prenorm	bool	If True, normalization is applied before activation, otherwise after. Defaults to True.	True
num_groups	int	Number of groups for the group normalization. Defaults to 1.	1

Source code in cliffordlayers/models/gca/twod.py

class CliffordG3MiddleBlock(nn.Module):
    """
    UNet middle block for G3 Clifford convolutions on 2D grids.

    Args:
        n_channels (int): Number of channels.
        activation (str): Type of activation function.
        norm (bool, optional): If True, normalization is applied. Defaults to False.
        prenorm (bool, optional): If True, normalization is applied before activation, otherwise after. Defaults to True.
        num_groups (int, optional): Number of groups for the group normalization. Defaults to 1.
    """

    def __init__(
        self,
        n_channels: int,
        activation: str,
        norm: bool = False,
        prenorm: bool = True,
        num_groups: int = 1,
    ):
        super().__init__()
        self.res1 = CliffordG3BasicBlock2d(
            n_channels,
            n_channels,
            activation=activation,
            norm=norm,
            prenorm=prenorm,
            num_groups=num_groups,
        )
        self.res2 = CliffordG3BasicBlock2d(
            n_channels,
            n_channels,
            activation=activation,
            norm=norm,
            prenorm=prenorm,
            num_groups=num_groups,
        )

    def forward(self, x):
        x = self.res1(x)
        x = self.res2(x)
        return x

CliffordG3ResNet2d

Bases: Module

ResNet for G3 Clifford convolutions on 2D grids.

Parameters:

Name	Type	Description	Default
num_blocks	list	Number of blocks at each resolution.	required
in_channels	int	Number of input channels.	required
out_channels	int	Number of output channels.	required
hidden_channels	int	Number of hidden channels.	required
activation	str	Type of activation function. Defaults to "vlin".	'vlin'
block	Module	Type of block. Defaults to CliffordG3BasicBlock2d.	CliffordG3BasicBlock2d
norm	bool	If True, normalization is applied. Defaults to True.	False
num_groups	int	Number of groups for the group normalization. Defaults to 1.	1
prenorm	bool	If True, normalization is applied before activation, otherwise after. Defaults to True.	True

Source code in cliffordlayers/models/gca/twod.py

class CliffordG3ResNet2d(nn.Module):
    """
    ResNet for G3 Clifford convolutions on 2D grids.

    Args:
        num_blocks (list): Number of blocks at each resolution.
        in_channels (int): Number of input channels.
        out_channels (int): Number of output channels.
        hidden_channels (int): Number of hidden channels.
        activation (str, optional): Type of activation function. Defaults to "vlin".
        block (nn.Module, optional): Type of block. Defaults to CliffordG3BasicBlock2d.
        norm (bool, optional): If True, normalization is applied. Defaults to True.
        num_groups (int, optional): Number of groups for the group normalization. Defaults to 1.
        prenorm (bool, optional): If True, normalization is applied before activation, otherwise after. Defaults to True.
    """

    padding = 9

    def __init__(
        self,
        num_blocks: list,
        in_channels: int,
        out_channels: int,
        hidden_channels: int,
        activation: str = "vlin",
        block: nn.Module = CliffordG3BasicBlock2d,
        norm: bool = False,
        num_groups: int = 1,
        prenorm=True,
    ):
        super().__init__()

        # Embedding layers
        self.conv_in1 = CliffordG3Conv2d(
            in_channels,
            hidden_channels,
            kernel_size=1,
            padding=0,
        )

        self.conv_in2 = CliffordG3Conv2d(
            hidden_channels,
            hidden_channels,
        )

        # Output layers
        self.conv_out1 = CliffordG3Conv2d(
            hidden_channels,
            hidden_channels,
        )
        self.conv_out2 = CliffordG3Conv2d(
            hidden_channels,
            out_channels,
            kernel_size=1,
            padding=0,
        )

        # ResNet blocks
        self.layers = nn.ModuleList(
            [
                self._make_layer(
                    block,
                    hidden_channels,
                    num_blocks[i],
                    activation=activation,
                    norm=norm,
                    num_groups=num_groups,
                    prenorm=prenorm,
                )
                for i in range(len(num_blocks))
            ]
        )

        self.act1 = get_activation(activation, hidden_channels)
        self.act2 = get_activation(activation, hidden_channels)

    def _make_layer(
        self,
        block: nn.Module,
        channels: int,
        num_blocks: int,
        activation: str,
        num_groups: int,
        norm: bool = True,
        prenorm: bool = True,
    ) -> nn.Sequential:
        layers = []
        for _ in range(num_blocks):
            layers.append(
                block(channels, channels, activation=activation, norm=norm, prenorm=prenorm, num_groups=num_groups)
            )
        return nn.Sequential(*layers)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        assert x.dim() == 5

        h = self.conv_in1(x)
        h = self.act1(h)

        # Second embedding layer
        h = self.conv_in2(h)

        for layer in self.layers:
            h = layer(h)

        # Output layers
        h = self.conv_out1(h)
        h = self.act2(h)
        h = self.conv_out2(h)

        # return output
        return h

CliffordG3UNet2d

Bases: Module

U-Net architecture with Clifford G3 convolutions for 2D grids.

Parameters:

Name	Type	Description	Default
in_channels	int	Number of input channels.	required
out_channels	int	Number of output channels.	required
hidden_channels	int	Number of channels in the first hidden convolutional layer.	required
activation	str	Type of activation function. Defaults to "vlin".	'vlin'
norm	bool	If True, normalization is applied. Defaults to False.	False
ch_mults	Union[Tuple[int, ...], List[int]]	Multipliers for the number of channels at each depth. Defaults to (1, 2, 2, 2).	(1, 2, 2, 2)
n_blocks	int	Number of convolutional blocks at each resolution. Defaults to 2.	2
prenorm	bool	If True, normalization is applied before activation, otherwise after. Defaults to True.	True
num_groups	int	Number of groups for the group normalization. Defaults to 1.	1

Source code in cliffordlayers/models/gca/twod.py

class CliffordG3UNet2d(nn.Module):
    """
    U-Net architecture with Clifford G3 convolutions for 2D grids.

    Args:
        in_channels (int): Number of input channels.
        out_channels (int): Number of output channels.
        hidden_channels (int): Number of channels in the first hidden convolutional layer.
        activation (str, optional): Type of activation function. Defaults to "vlin".
        norm (bool, optional): If True, normalization is applied. Defaults to False.
        ch_mults (Union[Tuple[int, ...], List[int]], optional): Multipliers for the number of channels at each depth.
                                                            Defaults to (1, 2, 2, 2).
        n_blocks (int, optional): Number of convolutional blocks at each resolution. Defaults to 2.
        prenorm (bool, optional): If True, normalization is applied before activation, otherwise after. Defaults to True.
        num_groups (int, optional): Number of groups for the group normalization. Defaults to 1.
    """

    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        hidden_channels: int,
        activation: str = "vlin",
        norm: bool = False,
        ch_mults: Union[Tuple[int, ...], List[int]] = (1, 2, 2, 2),
        n_blocks: int = 2,
        prenorm: bool = True,
        num_groups: int = 1,
    ) -> None:
        super().__init__()

        self.out_channels = out_channels

        # Number of resolutions
        n_resolutions = len(ch_mults)

        self.conv1 = CliffordG3Conv2d(
            in_channels,
            hidden_channels,
            kernel_size=3,
            padding=1,
        )

        # Decreasing resolution
        down = []
        # Number of channels
        out_channels = in_channels = hidden_channels
        # For each resolution
        for i in range(n_resolutions):
            # Number of output channels at this resolution
            out_channels = in_channels * ch_mults[i]
            for _ in range(n_blocks):
                down.append(
                    CliffordG3DownBlock(
                        in_channels,
                        out_channels,
                        activation=activation,
                        norm=norm,
                        prenorm=prenorm,
                        num_groups=num_groups,
                    )
                )
                in_channels = out_channels
            # Down sample at all resolutions except the last
            if i < n_resolutions - 1:
                down.append(
                    CliffordG3Downsample(
                        in_channels,
                    )
                )

        # Combine the set of modules
        self.down = nn.ModuleList(down)

        # Middle block
        self.middle = CliffordG3MiddleBlock(out_channels, activation=activation, norm=norm, prenorm=prenorm)

        # Increasing resolution
        up = []
        # Number of channels
        in_channels = out_channels
        # For each resolution
        for i in reversed(range(n_resolutions)):
            # `n_blocks` at the same resolution
            out_channels = in_channels
            for _ in range(n_blocks):
                up.append(
                    CliffordG3UpBlock(
                        in_channels,
                        out_channels,
                        activation=activation,
                        norm=norm,
                        prenorm=prenorm,
                        num_groups=num_groups,
                    )
                )
            # Final block to reduce the number of channels
            out_channels = in_channels // ch_mults[i]
            up.append(
                CliffordG3UpBlock(
                    in_channels,
                    out_channels,
                    activation=activation,
                    norm=norm,
                    prenorm=prenorm,
                )
            )
            in_channels = out_channels
            # Up sample at all resolutions except last
            if i > 0:
                up.append(
                    CliffordUpsample(
                        in_channels,
                    )
                )

        # Combine the set of modules
        self.up = nn.ModuleList(up)

        self.activation = get_activation(activation, out_channels)

        if norm:
            self.norm = CliffordG3GroupNorm(num_groups, out_channels, 3)
        else:
            self.norm = nn.Identity()

        # Output layers
        self.conv2 = CliffordG3Conv2d(
            in_channels,
            self.out_channels,
            kernel_size=3,
            padding=1,
        )

    def forward(self, x: torch.Tensor):
        assert x.dim() == 5

        x = self.conv1(x)

        h = [x]
        for m in self.down:
            x = m(x)
            h.append(x)

        x = self.middle(x)

        for m in self.up:
            if isinstance(m, CliffordUpsample):
                x = m(x)
            else:
                # Get the skip connection from first half of U-Net and concatenate
                s = h.pop()
                x = torch.cat((x, s), dim=1)
                x = m(x)

        x = self.activation(self.norm(x))
        x = self.conv2(x)
        return x

CliffordG3UpBlock

Bases: Module

UNet decoder block for G3 Clifford convolutions on 2D grids.

Parameters:

Name	Type	Description	Default
in_channels	int	Number of input channels.	required
out_channels	int	Number of output channels.	required
activation	str	Type of activation function.	required
norm	bool	If True, normalization is applied. Defaults to False.	False
prenorm	bool	If True, normalization is applied before activation, otherwise after. Defaults to True.	True
num_groups	int	Number of groups for the group normalization. Defaults to 1.	1

Source code in cliffordlayers/models/gca/twod.py

class CliffordG3UpBlock(nn.Module):
    """
    UNet decoder block for G3 Clifford convolutions on 2D grids.

    Args:
        in_channels (int): Number of input channels.
        out_channels (int): Number of output channels.
        activation (str): Type of activation function.
        norm (bool, optional): If True, normalization is applied. Defaults to False.
        prenorm (bool, optional): If True, normalization is applied before activation, otherwise after. Defaults to True.
        num_groups (int, optional): Number of groups for the group normalization. Defaults to 1.
    """

    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        activation: str,
        norm: bool = False,
        prenorm: bool = True,
        num_groups: int = 1,
    ):
        super().__init__()
        # The input has `in_channels + out_channels` because we concatenate the output of the same resolution
        # from the first half of the U-Net
        self.res = CliffordG3BasicBlock2d(
            in_channels + out_channels,
            out_channels,
            activation=activation,
            norm=norm,
            prenorm=prenorm,
            num_groups=num_groups,
        )

    def forward(self, x):
        return self.res(x)

CliffordUpsample

Bases: Module

Scale up the two-dimensional G3 Clifford feature map by a factor of two.

Parameters:

Name	Type	Description	Default
n_channels	int	Number of channels.	required

Source code in cliffordlayers/models/gca/twod.py

class CliffordUpsample(nn.Module):
    """
    Scale up the two-dimensional G3 Clifford feature map by a factor of two.

    Args:
        n_channels (int): Number of channels.
    """

    def __init__(self, n_channels: int):
        super().__init__()
        self.conv = CliffordG3ConvTranspose2d(
            n_channels,
            n_channels,
            4,
            2,
            1,
        )

    def forward(self, x):
        return self.conv(x)

3D models

The following code snippet initializes a 3D Clifford FNO.

import torch.nn.functional as F

from cliffordlayers.models.models_3d import (
    CliffordMaxwellNet3d,
    CliffordFourierBasicBlock3d,
)
model = CliffordMaxwellNet3d(
        g = [1, 1, 1],
        block = CliffordFourierBasicBlock3d,
        num_blocks = [1, 1, 1, 1],
        in_channels = 4,
        out_channels = 1,
        hidden_channels = 16,
        activation = F.gelu,
        norm = False,
    )

CliffordFourierBasicBlock3d

Bases: Module

3D building block for Clifford FNO architectures.

Parameters:

Name	Type	Description	Default
g	Union[tuple, list, Tensor]	Signature of Clifford algebra.	required
in_channels	int	Number of input channels.	required
out_channels	int	Number of output channels.	required
activation	Callable	Activation function. Defaults to F.gelu.	gelu
kernel_size	int	Kernel size of Clifford convolution. Defaults to 3.	1
stride	int	Stride of Clifford convolution. Defaults to 1.	1
padding	int	Padding of Clifford convolution. Defaults to 1.	0
norm	bool	Wether to use Clifford (group) normalization. Defaults to False.	False
num_groups	int	Number of groups when using Clifford (group) normalization. Defaults to 1.	1
modes1	int	Number of Fourier modes in the first dimension. Defaults to 8.	8
modes2	int	Number of Fourier modes in the second dimension. Defaults to 8.	8
modes3	int	Number of Fourier modes in the third dimension. Defaults to 8.	8

Source code in cliffordlayers/models/basic/threed.py

class CliffordFourierBasicBlock3d(nn.Module):
    """3D building block for Clifford FNO architectures.

    Args:
        g (Union[tuple, list, torch.Tensor]): Signature of Clifford algebra.
        in_channels (int): Number of input channels.
        out_channels (int): Number of output channels.
        activation (Callable, optional): Activation function. Defaults to F.gelu.
        kernel_size (int, optional): Kernel size of Clifford convolution. Defaults to 3.
        stride (int, optional): Stride of Clifford convolution. Defaults to 1.
        padding (int, optional): Padding of Clifford convolution. Defaults to 1.
        norm (bool, optional): Wether to use Clifford (group) normalization. Defaults to False.
        num_groups (int, optional): Number of groups when using Clifford (group) normalization. Defaults to 1.
        modes1 (int, optional): Number of Fourier modes in the first dimension. Defaults to 8.
        modes2 (int, optional): Number of Fourier modes in the second dimension. Defaults to 8.
        modes3 (int, optional): Number of Fourier modes in the third dimension. Defaults to 8.
    """

    expansion: int = 1

    def __init__(
        self,
        g: Union[tuple, list, torch.Tensor],
        in_channels: int,
        out_channels: int,
        activation: Callable = F.gelu,
        kernel_size: int = 1,
        stride: int = 1,
        padding: int = 0,
        norm: bool = False,
        num_groups: int = 1,
        modes1: int = 8,
        modes2: int = 8,
        modes3: int = 8,
    ):
        super().__init__()
        self.fourier = CliffordSpectralConv3d(
            g,
            in_channels,
            out_channels,
            modes1=modes1,
            modes2=modes2,
            modes3=modes3,
        )
        self.conv = CliffordConv3d(
            g,
            in_channels,
            out_channels,
            kernel_size=kernel_size,
            stride=stride,
            padding=padding,
            bias=True,
        )
        self.norm = CliffordGroupNorm3d(g, num_groups, in_channels) if norm else nn.Identity()
        self.activation = activation

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x1 = self.fourier(x)
        x2 = self.conv(x)
        return self.activation(self.norm(x1 + x2))

CliffordMaxwellNet3d

Bases: Module

3D building block for Clifford architectures with ResNet backbone network. The backbone networks follows these three steps: 1. Clifford vector+bivector encoding. 2. Basic blocks as provided. 3. Clifford vector+bivector decoding.

Parameters:

Name	Type	Description	Default
g	Union[tuple, list, Tensor]	Signature of Clifford algebra.	required
block	Module	Choice of basic blocks.	required
num_blocks	list	List of basic blocks in each residual block.	required
in_channels	int	Number of input channels.	required
out_channels	int	Number of output channels.	required
activation	Callable	Activation function. Defaults to F.gelu.	required
norm	bool	Wether to use Clifford (group) normalization. Defaults to False.	False
num_groups	int	Number of groups when using Clifford (group) normalization. Defaults to 1.	1

Source code in cliffordlayers/models/basic/threed.py

class CliffordMaxwellNet3d(nn.Module):
    """3D building block for Clifford architectures with ResNet backbone network.
    The backbone networks follows these three steps:
        1. Clifford vector+bivector encoding.
        2. Basic blocks as provided.
        3. Clifford vector+bivector decoding.

    Args:
        g (Union[tuple, list, torch.Tensor]): Signature of Clifford algebra.
        block (nn.Module): Choice of basic blocks.
        num_blocks (list): List of basic blocks in each residual block.
        in_channels (int): Number of input channels.
        out_channels (int): Number of output channels.
        activation (Callable, optional): Activation function. Defaults to F.gelu.
        norm (bool, optional): Wether to use Clifford (group) normalization. Defaults to False.
        num_groups (int, optional): Number of groups when using Clifford (group) normalization. Defaults to 1.
    """

    # For periodic boundary conditions, set padding = 0.
    padding = 2

    def __init__(
        self,
        g: Union[tuple, list, torch.Tensor],
        block: nn.Module,
        num_blocks: list,
        in_channels: int,
        out_channels: int,
        hidden_channels: int,
        activation: Callable,
        norm: bool = False,
        num_groups: int = 1,
    ):
        super().__init__()

        self.activation = activation
        # Encoding and decoding layers.
        self.encoder = CliffordConv3dMaxwellEncoder(
            g,
            in_channels=in_channels,
            out_channels=hidden_channels,
            kernel_size=1,
            padding=0,
        )
        self.decoder = CliffordConv3dMaxwellDecoder(
            g,
            in_channels=hidden_channels,
            out_channels=out_channels,
            kernel_size=1,
            padding=0,
        )

        # Residual blocks.
        self.layers = nn.ModuleList(
            [
                self._make_basic_block(
                    g,
                    block,
                    hidden_channels,
                    num_blocks[i],
                    activation=activation,
                    norm=norm,
                    num_groups=num_groups,
                )
                for i in range(len(num_blocks))
            ]
        )

    def _make_basic_block(
        self,
        g,
        block: nn.Module,
        hidden_channels: int,
        num_blocks: int,
        activation: Callable,
        norm: bool,
        num_groups: int,
    ) -> nn.Sequential:
        blocks = []
        for _ in range(num_blocks):
            blocks.append(
                block(
                    g,
                    hidden_channels,
                    hidden_channels,
                    activation=activation,
                    norm=norm,
                    num_groups=num_groups,
                )
            )
        return nn.Sequential(*blocks)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        assert x.dim() == 6

        # Encoding layer.
        x = self.encoder(self.activation(x))

        # Embed for non-periodic boundaries.
        if self.padding > 0:
            B_dim, C_dim, *D_dims, I_dim = range(len(x.shape))
            x = x.permute(B_dim, I_dim, C_dim, *D_dims)
            x = F.pad(x, [0, self.padding, 0, self.padding, 0, self.padding])
            B_dim, I_dim, C_dim, *D_dims = range(len(x.shape))
            x = x.permute(B_dim, C_dim, *D_dims, I_dim)

        # Apply residual layers.
        for layer in self.layers:
            x = layer(x)

        # Decoding layer.
        if self.padding > 0:
            B_dim, C_dim, *D_dims, I_dim = range(len(x.shape))
            x = x.permute(B_dim, I_dim, C_dim, *D_dims)
            x = x[..., : -self.padding, : -self.padding, : -self.padding]
            B_dim, I_dim, C_dim, *D_dims = range(len(x.shape))
            x = x.permute(B_dim, C_dim, *D_dims, I_dim)

        # Output layer.
        x = self.decoder(x)
        return x