Skip to content

Available PDE Surrogate Modules¤

SpectralConv1d ¤

Bases: Module

1D Fourier layer. Does FFT, linear transform, and Inverse FFT. Implemented in a way to allow multi-gpu training. Args: in_channels (int): Number of input channels out_channels (int): Number of output channels modes (int): Number of Fourier modes paper

Source code in pdearena/modules/fourier.py 
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
class SpectralConv1d(nn.Module):
    """1D Fourier layer. Does FFT, linear transform, and Inverse FFT.
    Implemented in a way to allow multi-gpu training.
    Args:
        in_channels (int): Number of input channels
        out_channels (int): Number of output channels
        modes (int): Number of Fourier modes
    [paper](https://arxiv.org/abs/2010.08895)
    """

    def __init__(self, in_channels: int, out_channels: int, modes: int):
        super().__init__()

        self.in_channels = in_channels
        self.out_channels = out_channels
        self.modes = modes

        self.scale = 1 / (in_channels * out_channels)
        self.weights = nn.Parameter(
            self.scale * torch.rand(in_channels, out_channels, self.modes, 2, dtype=torch.float32)
        )

    def forward(self, x):
        batchsize = x.shape[0]
        # Compute Fourier coeffcients up to factor of e^(- something constant)
        x_ft = torch.fft.rfft(x)

        # Multiply relevant Fourier modes
        out_ft = torch.zeros(
            batchsize,
            self.out_channels,
            x.size(-1) // 2 + 1,
            dtype=torch.cfloat,
            device=x.device,
        )
        out_ft[:, :, : self.modes] = batchmul1d(x_ft[:, :, : self.modes], torch.view_as_complex(self.weights))

        # Return to physical space
        x = torch.fft.irfft(out_ft, n=x.size(-1))
        return x

SpectralConv2d ¤

Bases: Module

2D Fourier layer. Does FFT, linear transform, and Inverse FFT. Implemented in a way to allow multi-gpu training. Args: in_channels (int): Number of input channels out_channels (int): Number of output channels modes1 (int): Number of Fourier modes to keep in the first spatial direction modes2 (int): Number of Fourier modes to keep in the second spatial direction paper

Source code in pdearena/modules/fourier.py 
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
class SpectralConv2d(nn.Module):
    """2D Fourier layer. Does FFT, linear transform, and Inverse FFT.
    Implemented in a way to allow multi-gpu training.
    Args:
        in_channels (int): Number of input channels
        out_channels (int): Number of output channels
        modes1 (int): Number of Fourier modes to keep in the first spatial direction
        modes2 (int): Number of Fourier modes to keep in the second spatial direction
    [paper](https://arxiv.org/abs/2010.08895)
    """

    def __init__(self, in_channels: int, out_channels: int, modes1: int, modes2: int):
        super().__init__()

        self.in_channels = in_channels
        self.out_channels = out_channels
        self.modes1 = modes1  # Number of Fourier modes to multiply, at most floor(N/2) + 1
        self.modes2 = modes2

        self.scale = 1 / (in_channels * out_channels)
        self.weights1 = nn.Parameter(
            self.scale * torch.rand(in_channels, out_channels, self.modes1, self.modes2, 2, dtype=torch.float32)
        )
        self.weights2 = nn.Parameter(
            self.scale * torch.rand(in_channels, out_channels, self.modes1, self.modes2, 2, dtype=torch.float32)
        )

    def forward(self, x, x_dim=None, y_dim=None):
        batchsize = x.shape[0]
        # Compute Fourier coeffcients up to factor of e^(- something constant)
        x_ft = torch.fft.rfft2(x)

        # Multiply relevant Fourier modes
        out_ft = torch.zeros(
            batchsize,
            self.out_channels,
            x.size(-2),
            x.size(-1) // 2 + 1,
            dtype=torch.cfloat,
            device=x.device,
        )
        out_ft[:, :, : self.modes1, : self.modes2] = batchmul2d(
            x_ft[:, :, : self.modes1, : self.modes2], torch.view_as_complex(self.weights1)
        )
        out_ft[:, :, -self.modes1 :, : self.modes2] = batchmul2d(
            x_ft[:, :, -self.modes1 :, : self.modes2], torch.view_as_complex(self.weights2)
        )

        # Return to physical space
        x = torch.fft.irfft2(out_ft, s=(x.size(-2), x.size(-1)))
        return x

SpectralConv3d ¤

Bases: Module

3D Fourier layer. Does FFT, linear transform, and Inverse FFT. Implemented in a way to allow multi-gpu training. Args: in_channels (int): Number of input channels out_channels (int): Number of output channels modes1 (int): Number of Fourier modes to keep in the first spatial direction modes2 (int): Number of Fourier modes to keep in the second spatial direction modes3 (int): Number of Fourier modes to keep in the third spatial direction paper

Source code in pdearena/modules/fourier.py 
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
class SpectralConv3d(nn.Module):
    """3D Fourier layer. Does FFT, linear transform, and Inverse FFT.
    Implemented in a way to allow multi-gpu training.
    Args:
        in_channels (int): Number of input channels
        out_channels (int): Number of output channels
        modes1 (int): Number of Fourier modes to keep in the first spatial direction
        modes2 (int): Number of Fourier modes to keep in the second spatial direction
        modes3 (int): Number of Fourier modes to keep in the third spatial direction
    [paper](https://arxiv.org/abs/2010.08895)
    """

    def __init__(self, in_channels: int, out_channels: int, modes1: int, modes2: int, modes3: int):
        super().__init__()

        self.in_channels = in_channels
        self.out_channels = out_channels
        self.modes1 = modes1  # Number of Fourier modes to multiply, at most floor(N/2) + 1
        self.modes2 = modes2
        self.modes3 = modes3

        self.scale = 1 / (in_channels * out_channels)
        self.weights1 = nn.Parameter(
            self.scale
            * torch.rand(in_channels, out_channels, self.modes1, self.modes2, self.modes3, 2, dtype=torch.float32)
        )
        self.weights2 = nn.Parameter(
            self.scale
            * torch.rand(in_channels, out_channels, self.modes1, self.modes2, self.modes3, 2, dtype=torch.float32)
        )
        self.weights3 = nn.Parameter(
            self.scale
            * torch.rand(in_channels, out_channels, self.modes1, self.modes2, self.modes3, 2, dtype=torch.float32)
        )
        self.weights4 = nn.Parameter(
            self.scale
            * torch.rand(in_channels, out_channels, self.modes1, self.modes2, self.modes3, 2, dtype=torch.float32)
        )

    def forward(self, x):
        batchsize = x.shape[0]
        # Compute Fourier coeffcients up to factor of e^(- something constant)
        x_ft = torch.fft.rfftn(x, dim=[-3, -2, -1])

        # Multiply relevant Fourier modes
        out_ft = torch.zeros(
            batchsize,
            self.out_channels,
            x.size(-3),
            x.size(-2),
            x.size(-1) // 2 + 1,
            dtype=torch.cfloat,
            device=x.device,
        )
        out_ft[:, :, : self.modes1, : self.modes2, : self.modes3] = batchmul3d(
            x_ft[:, :, : self.modes1, : self.modes2, : self.modes3], torch.view_as_complex(self.weights1)
        )
        out_ft[:, :, -self.modes1 :, : self.modes2, : self.modes3] = batchmul3d(
            x_ft[:, :, -self.modes1 :, : self.modes2, : self.modes3], torch.view_as_complex(self.weights2)
        )
        out_ft[:, :, : self.modes1, -self.modes2 :, : self.modes3] = batchmul3d(
            x_ft[:, :, : self.modes1, -self.modes2 :, : self.modes3], torch.view_as_complex(self.weights3)
        )
        out_ft[:, :, -self.modes1 :, -self.modes2 :, : self.modes3] = batchmul3d(
            x_ft[:, :, -self.modes1 :, -self.modes2 :, : self.modes3], torch.view_as_complex(self.weights4)
        )

        # Return to physical space
        x = torch.fft.irfftn(out_ft, s=(x.size(-3), x.size(-2), x.size(-1)))
        return x

DilatedBasicBlock ¤

Bases: Module

Basic block for Dilated ResNet

Parameters:

Name Type Description Default
in_planes int

number of input channels

 required
planes int

number of output channels

 required
stride int

stride of the convolution. Defaults to 1.

 1
activation str

activation function. Defaults to "relu".

 'relu'
norm bool

whether to use group normalization. Defaults to True.

 True
num_groups int

number of groups for group normalization. Defaults to 1.

 1
Source code in pdearena/modules/twod_resnet.py 
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
class DilatedBasicBlock(nn.Module):
    """Basic block for Dilated ResNet

    Args:
        in_planes (int): number of input channels
        planes (int): number of output channels
        stride (int, optional): stride of the convolution. Defaults to 1.
        activation (str, optional): activation function. Defaults to "relu".
        norm (bool, optional): whether to use group normalization. Defaults to True.
        num_groups (int, optional): number of groups for group normalization. Defaults to 1.
    """

    expansion = 1

    def __init__(
        self,
        in_planes: int,
        planes: int,
        stride: int = 1,
        activation: str = "relu",
        norm: bool = True,
        num_groups: int = 1,
    ):
        super().__init__()

        self.dilation = [1, 2, 4, 8, 4, 2, 1]
        dilation_layers = []
        for dil in self.dilation:
            dilation_layers.append(
                nn.Conv2d(
                    in_planes,
                    planes,
                    kernel_size=3,
                    stride=stride,
                    dilation=dil,
                    padding=dil,
                    bias=True,
                )
            )
        self.dilation_layers = nn.ModuleList(dilation_layers)
        self.norm_layers = nn.ModuleList(
            nn.GroupNorm(num_groups, num_channels=planes) if norm else nn.Identity() for dil in self.dilation
        )
        self.activation: nn.Module = ACTIVATION_REGISTRY.get(activation, None)
        if self.activation is None:
            raise NotImplementedError(f"Activation {activation} not implemented")

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        out = x
        for layer, norm in zip(self.dilation_layers, self.norm_layers):
            out = self.activation(layer(norm(out)))
        return out + x

FourierBasicBlock ¤

Bases: Module

Basic block for Fourier Neural Operators

Parameters:

Name Type Description Default
in_planes int

number of input channels

 required
planes int

number of output channels

 required
stride int

stride of the convolution. Defaults to 1.

 1
modes1 int

number of modes for the first spatial dimension. Defaults to 16.

 16
modes2 int

number of modes for the second spatial dimension. Defaults to 16.

 16
activation str

activation function. Defaults to "relu".

 'gelu'
norm bool

whether to use group normalization. Defaults to False.

 False
Source code in pdearena/modules/twod_resnet.py 
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
class FourierBasicBlock(nn.Module):
    """Basic block for Fourier Neural Operators

    Args:
        in_planes (int): number of input channels
        planes (int): number of output channels
        stride (int, optional): stride of the convolution. Defaults to 1.
        modes1 (int, optional): number of modes for the first spatial dimension. Defaults to 16.
        modes2 (int, optional): number of modes for the second spatial dimension. Defaults to 16.
        activation (str, optional): activation function. Defaults to "relu".
        norm (bool, optional): whether to use group normalization. Defaults to False.

    """

    expansion: int = 1

    def __init__(
        self,
        in_planes: int,
        planes: int,
        stride: int = 1,
        modes1: int = 16,
        modes2: int = 16,
        activation: str = "gelu",
        norm: bool = False,
    ):
        super().__init__()
        self.modes1 = modes1
        self.modes2 = modes2
        assert not norm
        self.fourier1 = SpectralConv2d(in_planes, planes, modes1=self.modes1, modes2=self.modes2)
        self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=1, padding=0, padding_mode="zeros", bias=True)
        self.fourier2 = SpectralConv2d(planes, planes, modes1=self.modes1, modes2=self.modes2)
        self.conv2 = nn.Conv2d(planes, planes, kernel_size=1, padding=0, padding_mode="zeros", bias=True)

        # So far shortcut connections are not helping
        # self.shortcut = nn.Sequential()
        # if stride != 1 or in_planes != self.expansion * planes:
        #     self.shortcut = nn.Sequential(
        #         nn.Conv2d(in_planes, self.expansion * planes, kernel_size=1)
        #     )

        self.activation: nn.Module = ACTIVATION_REGISTRY.get(activation, None)
        if self.activation is None:
            raise NotImplementedError(f"Activation {activation} not implemented")

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x1 = self.fourier1(x)
        x2 = self.conv1(x)
        out = self.activation(x1 + x2)

        x1 = self.fourier2(out)
        x2 = self.conv2(out)
        out = x1 + x2
        # out += self.shortcut(x)
        out = self.activation(out)
        return out

ResNet ¤

Bases: Module

Class to support ResNet like feedforward architectures

Parameters:

Name Type Description Default
n_input_scalar_components int

Number of input scalar components in the model

 required
n_input_vector_components int

Number of input vector components in the model

 required
n_output_scalar_components int

Number of output scalar components in the model

 required
n_output_vector_components int

Number of output vector components in the model

 required
block Callable

BasicBlock or DilatedBasicBlock or FourierBasicBlock

 required
num_blocks List[int]

Number of blocks in each layer

 required
time_history int

Number of time steps to use in the input

 required
time_future int

Number of time steps to predict in the output

 required
hidden_channels int

Number of channels in the hidden layers

 64
activation str

Activation function to use

 'gelu'
norm bool

Whether to use normalization

 True
Source code in pdearena/modules/twod_resnet.py 
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
class ResNet(nn.Module):
    """Class to support ResNet like feedforward architectures

    Args:
        n_input_scalar_components (int): Number of input scalar components in the model
        n_input_vector_components (int): Number of input vector components in the model
        n_output_scalar_components (int): Number of output scalar components in the model
        n_output_vector_components (int): Number of output vector components in the model
        block (Callable): BasicBlock or DilatedBasicBlock or FourierBasicBlock
        num_blocks (List[int]): Number of blocks in each layer
        time_history (int): Number of time steps to use in the input
        time_future (int): Number of time steps to predict in the output
        hidden_channels (int): Number of channels in the hidden layers
        activation (str): Activation function to use
        norm (bool): Whether to use normalization
    """

    padding = 9

    def __init__(
        self,
        n_input_scalar_components: int,
        n_input_vector_components: int,
        n_output_scalar_components: int,
        n_output_vector_components: int,
        block,
        num_blocks: list,
        time_history: int,
        time_future: int,
        hidden_channels: int = 64,
        activation: str = "gelu",
        norm: bool = True,
        diffmode: bool = False,
        usegrid: bool = False,
    ):
        super().__init__()
        self.n_input_scalar_components = n_input_scalar_components
        self.n_input_vector_components = n_input_vector_components
        self.n_output_scalar_components = n_output_scalar_components
        self.n_output_vector_components = n_output_vector_components
        self.diffmode = diffmode
        self.usegrid = usegrid
        self.in_planes = hidden_channels
        insize = time_history * (self.n_input_scalar_components + self.n_input_vector_components * 2)
        if self.usegrid:
            insize += 2
        self.conv_in1 = nn.Conv2d(
            insize,
            self.in_planes,
            kernel_size=1,
            bias=True,
        )
        self.conv_in2 = nn.Conv2d(
            self.in_planes,
            self.in_planes,
            kernel_size=1,
            bias=True,
        )
        self.conv_out1 = nn.Conv2d(
            self.in_planes,
            self.in_planes,
            kernel_size=1,
            bias=True,
        )
        self.conv_out2 = nn.Conv2d(
            self.in_planes,
            time_future * (self.n_output_scalar_components + self.n_output_vector_components * 2),
            kernel_size=1,
            bias=True,
        )

        self.layers = nn.ModuleList(
            [
                self._make_layer(
                    block,
                    self.in_planes,
                    num_blocks[i],
                    stride=1,
                    activation=activation,
                    norm=norm,
                )
                for i in range(len(num_blocks))
            ]
        )
        self.activation: nn.Module = ACTIVATION_REGISTRY.get(activation, None)
        if self.activation is None:
            raise NotImplementedError(f"Activation {activation} not implemented")

    def _make_layer(
        self,
        block: Callable,
        planes: int,
        num_blocks: int,
        stride: int,
        activation: str,
        norm: bool = True,
    ) -> nn.Sequential:
        strides = [stride] + [1] * (num_blocks - 1)
        layers = []
        for stride in strides:
            layers.append(
                block(
                    self.in_planes,
                    planes,
                    stride,
                    activation=activation,
                    norm=norm,
                )
            )
            self.in_planes = planes * block.expansion
        return nn.Sequential(*layers)

    def __repr__(self):
        return "ResNet"

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        assert x.dim() == 5
        orig_shape = x.shape
        x = x.reshape(x.size(0), -1, *x.shape[3:])  # collapse T,C
        # prev = x.float()
        x = self.activation(self.conv_in1(x.float()))
        x = self.activation(self.conv_in2(x.float()))

        if self.padding > 0:
            x = F.pad(x, [0, self.padding, 0, self.padding])

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

        if self.padding > 0:
            x = x[..., : -self.padding, : -self.padding]

        x = self.activation(self.conv_out1(x))
        x = self.conv_out2(x)

        if self.diffmode:
            raise NotImplementedError("diffmode")
            # x = x + prev[:, -1:, ...].detach()
        return x.reshape(
            orig_shape[0], -1, (self.n_output_scalar_components + self.n_output_vector_components * 2), *orig_shape[3:]
        )

OperatorBlock_2D ¤

Bases: Module

Parameters:

Name Type Description Default
in_codim int

Input co-domian dimension

 required
out_codim int

output co-domain dimension

 required
dim1 int

Default output grid size along x (or 1st dimension)

 required
dim2 int

Default output grid size along y ( or 2nd dimension)

 required
modes1 int

Number of fourier modes to consider along 1st dimension

 required
modes2 int

Number of fourier modes to consider along 2nd dimension

 required
norm bool

Whether to use normalization (torch.nn.InstanceNorm2d)

 True
nonlin bool

Whether to use non-linearity (torch.nn.GELU)

 True

All variables are consistent with the SpectralConv2d_Uno.

Source code in pdearena/modules/twod_uno.py 
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
class OperatorBlock_2D(nn.Module):
    """

    Args:
        in_codim (int): Input co-domian dimension
        out_codim (int): output co-domain dimension
        dim1 (int):  Default output grid size along x (or 1st dimension)
        dim2 (int): Default output grid size along y ( or 2nd dimension)
        modes1 (int): Number of fourier modes to consider along 1st dimension
        modes2 (int): Number of fourier modes to consider along 2nd dimension
        norm (bool): Whether to use normalization ([torch.nn.InstanceNorm2d][])
        nonlin (bool): Whether to use non-linearity ([torch.nn.GELU][])


    All variables are consistent with the [`SpectralConv2d_Uno`][pdearena.modules.twod_uno.SpectralConv2d_Uno].
    """

    def __init__(self, in_codim, out_codim, dim1, dim2, modes1, modes2, norm=True, nonlin=True):
        super().__init__()
        self.conv = SpectralConv2d_Uno(in_codim, out_codim, dim1, dim2, modes1, modes2)
        self.w = Pointwise_op_2D(in_codim, out_codim, dim1, dim2)
        self.norm = norm
        self.non_lin = nonlin
        if norm:
            self.normalize_layer = nn.InstanceNorm2d(int(out_codim), affine=True)

    def forward(self, x, dim1=None, dim2=None):
        #
        # input shape = (batch, in_codim, input_dim1,input_dim2)
        # output shape = (batch, out_codim, dim1,dim2)
        x1_out = self.conv(x, dim1, dim2)
        x2_out = self.w(x, dim1, dim2)
        x_out = x1_out + x2_out
        if self.norm:
            x_out = self.normalize_layer(x_out)
        if self.non_lin:
            x_out = F.gelu(x_out)
        return x_out

Pointwise_op_2D ¤

Bases: Module

Parameters:

Name Type Description Default
in_codim int

Input co-domian dimension

 required
out_codim int

output co-domain dimension

 required
dim1 int

Default output grid size along x (or 1st dimension)

 required
dim2 int

Default output grid size along y ( or 2nd dimension)

 required
Source code in pdearena/modules/twod_uno.py 
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
class Pointwise_op_2D(nn.Module):
    """

    Args:
        in_codim (int): Input co-domian dimension
        out_codim (int): output co-domain dimension

        dim1 (int):  Default output grid size along x (or 1st dimension)
        dim2 (int): Default output grid size along y ( or 2nd dimension)
    """

    def __init__(self, in_codim: int, out_codim: int, dim1: int, dim2: int):
        super().__init__()
        self.conv = nn.Conv2d(int(in_codim), int(out_codim), 1)
        self.dim1 = int(dim1)
        self.dim2 = int(dim2)

    def forward(self, x, dim1=None, dim2=None):
        #
        # input shape = (batch, in_codim, input_dim1,input_dim2)
        # output shape = (batch, out_codim, dim1,dim2)

        if dim1 is None:
            dim1 = self.dim1
            dim2 = self.dim2
        x_out = self.conv(x)

        x_out = F.interpolate(x_out, size=(dim1, dim2), mode="bicubic", align_corners=True, antialias=True)
        return x_out

SpectralConv2d_Uno ¤

Bases: Module

2D Fourier layer. It does FFT, linear transform, and Inverse FFT.

Modified to support multi-gpu training.

Parameters:

Name Type Description Default
in_codim int

Input co-domian dimension

 required
out_codim int

output co-domain dimension

 required
dim1 int

Default output grid size along x (or 1st dimension of output domain)

 required
dim2 int

Default output grid size along y ( or 2nd dimension of output domain) Ratio of grid size of the input and the output implecitely set the expansion or contraction farctor along each dimension.

 required
modes1 int), modes2 (int

Number of fourier modes to consider for the ontegral operator Number of modes must be compatibale with the input grid size and desired output grid size. i.e., modes1 <= min( dim1/2, input_dim1/2). Here "input_dim1" is the grid size along x axis (or first dimension) of the input domain. Other modes also the have same constrain.

 None
Source code in pdearena/modules/twod_uno.py 
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
class SpectralConv2d_Uno(nn.Module):
    """2D Fourier layer. It does FFT, linear transform, and Inverse FFT.

    Modified to support multi-gpu training.

    Args:
        in_codim (int): Input co-domian dimension
        out_codim (int): output co-domain dimension

        dim1 (int): Default output grid size along x (or 1st dimension of output domain)
        dim2 (int): Default output grid size along y ( or 2nd dimension of output domain)
                    Ratio of grid size of the input and the output implecitely
                    set the expansion or contraction farctor along each dimension.
        modes1 (int), modes2 (int):  Number of fourier modes to consider for the ontegral operator
                    Number of modes must be compatibale with the input grid size
                    and desired output grid size.
                    i.e., modes1 <= min( dim1/2, input_dim1/2).
                    Here "input_dim1" is the grid size along x axis (or first dimension) of the input domain.
                    Other modes also the have same constrain.
    """

    def __init__(self, in_codim, out_codim, dim1, dim2, modes1=None, modes2=None):
        super().__init__()

        in_codim = int(in_codim)
        out_codim = int(out_codim)
        self.in_channels = in_codim
        self.out_channels = out_codim
        self.dim1 = dim1
        self.dim2 = dim2
        if modes1 is not None:
            self.modes1 = modes1
            self.modes2 = modes2
        else:
            self.modes1 = dim1 // 2 - 1
            self.modes2 = dim2 // 2
        self.scale = (1 / (2 * in_codim)) ** (1.0 / 2.0)
        self.weights1 = nn.Parameter(
            self.scale * (torch.randn(in_codim, out_codim, self.modes1, self.modes2, 2, dtype=torch.float32))
        )
        self.weights2 = nn.Parameter(
            self.scale * (torch.randn(in_codim, out_codim, self.modes1, self.modes2, 2, dtype=torch.float32))
        )

    # Complex multiplication
    def compl_mul2d(self, input, weights):
        return torch.einsum("bixy,ioxy->boxy", input, weights)

    def forward(self, x, dim1=None, dim2=None):
        if dim1 is not None:
            self.dim1 = dim1
            self.dim2 = dim2
        batchsize = x.shape[0]
        # Compute Fourier coeffcients up to factor of e^(- something constant)
        x_ft = torch.fft.rfft2(x, norm="forward")

        # Multiply relevant Fourier modes
        out_ft = torch.zeros(
            batchsize,
            self.out_channels,
            self.dim1,
            self.dim2 // 2 + 1,
            dtype=torch.cfloat,
            device=x.device,
        )
        out_ft[:, :, : self.modes1, : self.modes2] = self.compl_mul2d(
            x_ft[:, :, : self.modes1, : self.modes2], torch.view_as_complex(self.weights1)
        )
        out_ft[:, :, -self.modes1 :, : self.modes2] = self.compl_mul2d(
            x_ft[:, :, -self.modes1 :, : self.modes2], torch.view_as_complex(self.weights2)
        )

        # Return to physical space
        x = torch.fft.irfft2(out_ft, s=(self.dim1, self.dim2), norm="forward")
        return x

UNO ¤

Bases: Module

UNO model

Parameters:

Name Type Description Default
n_input_scalar_components int

Number of scalar components in the model

 required
n_input_vector_components int

Number of vector components in the model

 required
n_output_scalar_components int

Number of output scalar components in the model

 required
n_output_vector_components int

Number of output vector components in the model

 required
time_history int

Number of time steps to include in the model

 required
time_future int

Number of time steps to predict in the model

 required
hidden_channels int

Number of hidden channels in the model

 required
pad int

Padding to use in the model

 0
factor int

Scaling factor to use in the model

 3 / 4
activation str

Activation function to use in the model

 'gelu'
Source code in pdearena/modules/twod_uno.py 
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
class UNO(nn.Module):
    """UNO model

    Args:
        n_input_scalar_components (int): Number of scalar components in the model
        n_input_vector_components (int): Number of vector components in the model
        n_output_scalar_components (int): Number of output scalar components in the model
        n_output_vector_components (int): Number of output vector components in the model
        time_history (int): Number of time steps to include in the model
        time_future (int): Number of time steps to predict in the model
        hidden_channels (int): Number of hidden channels in the model
        pad (int): Padding to use in the model
        factor (int): Scaling factor to use in the model
        activation (str): Activation function to use in the model
    """

    def __init__(
        self,
        n_input_scalar_components: int,
        n_input_vector_components: int,
        n_output_scalar_components: int,
        n_output_vector_components: int,
        time_history: int,
        time_future: int,
        hidden_channels: int,
        pad=0,
        factor=3 / 4,
        activation="gelu",
    ):
        super().__init__()

        self.n_input_scalar_components = n_input_scalar_components
        self.n_input_vector_components = n_input_vector_components
        self.n_output_scalar_components = n_output_scalar_components
        self.n_output_vector_components = n_output_vector_components

        self.width = hidden_channels
        self.factor = factor
        self.padding = pad
        self.activation: nn.Module = ACTIVATION_REGISTRY.get(activation, None)
        if self.activation is None:
            raise NotImplementedError(f"Activation {activation} not implemented")

        in_width = time_history * (self.n_input_scalar_components + self.n_input_vector_components * 2)
        out_width = time_future * (self.n_output_scalar_components + self.n_output_vector_components * 2)
        self.fc = nn.Linear(in_width, self.width // 2)

        self.fc0 = nn.Linear(self.width // 2, self.width)  # input channel is 3: (a(x, y), x, y)

        self.L0 = OperatorBlock_2D(self.width, 2 * factor * self.width, 48, 48, 18, 18)

        self.L1 = OperatorBlock_2D(2 * factor * self.width, 4 * factor * self.width, 32, 32, 14, 14)

        self.L2 = OperatorBlock_2D(4 * factor * self.width, 8 * factor * self.width, 16, 16, 6, 6)

        self.L3 = OperatorBlock_2D(8 * factor * self.width, 8 * factor * self.width, 16, 16, 6, 6)

        self.L4 = OperatorBlock_2D(8 * factor * self.width, 4 * factor * self.width, 32, 32, 6, 6)

        self.L5 = OperatorBlock_2D(8 * factor * self.width, 2 * factor * self.width, 48, 48, 14, 14)

        self.L6 = OperatorBlock_2D(4 * factor * self.width, self.width, 64, 64, 18, 18)  # will be reshaped

        self.fc1 = nn.Linear(2 * self.width, 4 * self.width)
        self.fc2 = nn.Linear(4 * self.width, out_width)

    def forward(self, x):
        assert x.dim() == 5
        orig_shape = x.shape
        x = x.reshape(x.size(0), -1, *x.shape[3:])  # collapse T,C

        x = x.permute(0, 2, 3, 1)
        x_fc = self.fc(x)
        x_fc = self.activation(x_fc)

        x_fc0 = self.fc0(x_fc)
        x_fc0 = self.activation(x_fc0)

        x_fc0 = x_fc0.permute(0, 3, 1, 2)

        x_fc0 = F.pad(x_fc0, [self.padding, self.padding, self.padding, self.padding])

        D1, D2 = x_fc0.shape[-2], x_fc0.shape[-1]

        x_c0 = self.L0(x_fc0, int(D1 * self.factor), int(D2 * self.factor))
        x_c1 = self.L1(x_c0, D1 // 2, D2 // 2)

        x_c2 = self.L2(x_c1, D1 // 4, D2 // 4)
        x_c3 = self.L3(x_c2, D1 // 4, D2 // 4)
        x_c4 = self.L4(x_c3, D1 // 2, D2 // 2)
        x_c4 = torch.cat([x_c4, x_c1], dim=1)
        x_c5 = self.L5(x_c4, int(D1 * self.factor), int(D2 * self.factor))
        x_c5 = torch.cat([x_c5, x_c0], dim=1)
        x_c6 = self.L6(x_c5, D1, D2)
        x_c6 = torch.cat([x_c6, x_fc0], dim=1)

        if self.padding != 0:
            x_c6 = x_c6[..., : -self.padding, : -self.padding]

        x_c6 = x_c6.permute(0, 2, 3, 1)

        x_fc1 = self.fc1(x_c6)
        x_fc1 = self.activation(x_fc1)

        x_out = self.fc2(x_fc1)
        x_out = x_out.permute(0, 3, 1, 2)

        return x_out.reshape(
            orig_shape[0], -1, (self.n_output_scalar_components + self.n_output_vector_components * 2), *orig_shape[3:]
        )

Unet2015 ¤

Bases: Module

Two-dimensional UNet based on original architecture.

Parameters:

Name Type Description Default
n_input_scalar_components int

Number of scalar components in the model

 required
n_input_vector_components int

Number of vector components in the model

 required
n_output_scalar_components int

Number of output scalar components in the model

 required
n_output_vector_components int

Number of output vector components in the model

 required
time_history int

Number of time steps in the input.

 required
time_future int

Number of time steps in the output.

 required
hidden_channels int

Number of hidden channels.

 required
activation str

Activation function.

 required
Source code in pdearena/modules/twod_unet2015.py 
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
class Unet2015(nn.Module):
    """Two-dimensional UNet based on original architecture.

    Args:
        n_input_scalar_components (int): Number of scalar components in the model
        n_input_vector_components (int): Number of vector components in the model
        n_output_scalar_components (int): Number of output scalar components in the model
        n_output_vector_components (int): Number of output vector components in the model
        time_history (int): Number of time steps in the input.
        time_future (int): Number of time steps in the output.
        hidden_channels (int): Number of hidden channels.
        activation (str): Activation function.
    """

    def __init__(
        self,
        n_input_scalar_components: int,
        n_input_vector_components: int,
        n_output_scalar_components: int,
        n_output_vector_components: int,
        time_history: int,
        time_future: int,
        hidden_channels: int,
        activation: str,
    ) -> None:
        super().__init__()
        self.n_input_scalar_components = n_input_scalar_components
        self.n_input_vector_components = n_input_vector_components
        self.n_output_scalar_components = n_output_scalar_components
        self.n_output_vector_components = n_output_vector_components
        self.time_history = time_history
        self.time_future = time_future
        self.hidden_channels = hidden_channels
        self.activation = ACTIVATION_REGISTRY.get(activation, None)
        if self.activation is None:
            raise NotImplementedError(f"Activation {activation} not implemented")

        in_channels = time_history * (self.n_input_scalar_components + self.n_input_vector_components * 2)
        out_channels = time_future * (self.n_output_scalar_components + self.n_output_vector_components * 2)

        features = hidden_channels
        self.encoder1 = Unet2015._block(in_channels, features, name="enc1", activation=self.activation)
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.encoder2 = Unet2015._block(features, features * 2, name="enc2", activation=self.activation)
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.encoder3 = Unet2015._block(features * 2, features * 4, name="enc3", activation=self.activation)
        self.pool3 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.encoder4 = Unet2015._block(features * 4, features * 8, name="enc4", activation=self.activation)
        self.pool4 = nn.MaxPool2d(kernel_size=2, stride=2)

        self.bottleneck = Unet2015._block(features * 8, features * 16, name="bottleneck", activation=self.activation)

        self.upconv4 = nn.ConvTranspose2d(features * 16, features * 8, kernel_size=2, stride=2)
        self.decoder4 = Unet2015._block((features * 8) * 2, features * 8, name="dec4", activation=self.activation)
        self.upconv3 = nn.ConvTranspose2d(features * 8, features * 4, kernel_size=2, stride=2)
        self.decoder3 = Unet2015._block((features * 4) * 2, features * 4, name="dec3", activation=self.activation)
        self.upconv2 = nn.ConvTranspose2d(features * 4, features * 2, kernel_size=2, stride=2)
        self.decoder2 = Unet2015._block((features * 2) * 2, features * 2, name="dec2", activation=self.activation)
        self.upconv1 = nn.ConvTranspose2d(features * 2, features, kernel_size=2, stride=2)
        self.decoder1 = Unet2015._block(features * 2, features, name="dec1", activation=self.activation)

        self.conv = nn.Conv2d(in_channels=features, out_channels=out_channels, kernel_size=1)

    def forward(self, x):
        assert x.dim() == 5
        orig_shape = x.shape
        x = x.reshape(x.size(0), -1, *x.shape[3:])
        enc1 = self.encoder1(x)
        enc2 = self.encoder2(self.pool1(enc1))
        enc3 = self.encoder3(self.pool2(enc2))
        enc4 = self.encoder4(self.pool3(enc3))

        bottleneck = self.bottleneck(self.pool4(enc4))

        dec4 = self.upconv4(bottleneck)
        dec4 = torch.cat((dec4, enc4), dim=1)
        dec4 = self.decoder4(dec4)
        dec3 = self.upconv3(dec4)
        dec3 = torch.cat((dec3, enc3), dim=1)
        dec3 = self.decoder3(dec3)
        dec2 = self.upconv2(dec3)
        dec2 = torch.cat((dec2, enc2), dim=1)
        dec2 = self.decoder2(dec2)
        dec1 = self.upconv1(dec2)
        dec1 = torch.cat((dec1, enc1), dim=1)
        dec1 = self.decoder1(dec1)
        out = self.conv(dec1)
        return out.reshape(orig_shape[0], -1, *orig_shape[2:])

    @staticmethod
    def _block(in_channels, features, name, activation=nn.Tanh()):
        return nn.Sequential(
            OrderedDict(
                [
                    (
                        name + "conv1",
                        nn.Conv2d(
                            in_channels=in_channels,
                            out_channels=features,
                            kernel_size=3,
                            padding=1,
                            bias=False,
                        ),
                    ),
                    (name + "norm1", nn.BatchNorm2d(num_features=features)),
                    (name + "act1", activation),
                    (
                        name + "conv2",
                        nn.Conv2d(
                            in_channels=features,
                            out_channels=features,
                            kernel_size=3,
                            padding=1,
                            bias=False,
                        ),
                    ),
                    (name + "norm2", nn.BatchNorm2d(num_features=features)),
                    (name + "act2", activation),
                ]
            )
        )

Unetbase ¤

Bases: Module

Our interpretation of the original U-Net architecture.

Uses torch.nn.GroupNorm instead of torch.nn.BatchNorm2d. Also there is no BottleNeck block.

Parameters:

Name Type Description Default
n_input_scalar_components int

Number of scalar components in the model

 required
n_input_vector_components int

Number of vector components in the model

 required
n_output_scalar_components int

Number of output scalar components in the model

 required
n_output_vector_components int

Number of output vector components in the model

 required
time_history int

Number of time steps in the input.

 required
time_future int

Number of time steps in the output.

 required
hidden_channels int

Number of channels in the hidden layers.

 required
activation str

Activation function to use. One of ["gelu", "relu", "silu"].

 'gelu'
Source code in pdearena/modules/twod_unetbase.py 
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
class Unetbase(nn.Module):
    """Our interpretation of the original U-Net architecture.

    Uses [torch.nn.GroupNorm][] instead of [torch.nn.BatchNorm2d][]. Also there is no `BottleNeck` block.

    Args:
        n_input_scalar_components (int): Number of scalar components in the model
        n_input_vector_components (int): Number of vector components in the model
        n_output_scalar_components (int): Number of output scalar components in the model
        n_output_vector_components (int): Number of output vector components in the model
        time_history (int): Number of time steps in the input.
        time_future (int): Number of time steps in the output.
        hidden_channels (int): Number of channels in the hidden layers.
        activation (str): Activation function to use. One of ["gelu", "relu", "silu"].
    """

    def __init__(
        self,
        n_input_scalar_components: int,
        n_input_vector_components: int,
        n_output_scalar_components: int,
        n_output_vector_components: int,
        time_history: int,
        time_future: int,
        hidden_channels: int,
        activation="gelu",
    ) -> None:
        super().__init__()
        self.n_input_scalar_components = n_input_scalar_components
        self.n_input_vector_components = n_input_vector_components
        self.n_output_scalar_components = n_output_scalar_components
        self.n_output_vector_components = n_output_vector_components
        self.time_history = time_history
        self.time_future = time_future
        self.hidden_channels = hidden_channels
        self.activation = ACTIVATION_REGISTRY.get(activation, None)
        if self.activation is None:
            raise NotImplementedError(f"Activation {activation} not implemented")

        insize = time_history * (self.n_input_scalar_components + self.n_input_vector_components * 2)
        n_channels = hidden_channels
        self.image_proj = ConvBlock(insize, n_channels, activation=activation)

        self.down = nn.ModuleList(
            [
                Down(n_channels, n_channels * 2, activation=activation),
                Down(n_channels * 2, n_channels * 4, activation=activation),
                Down(n_channels * 4, n_channels * 8, activation=activation),
                Down(n_channels * 8, n_channels * 16, activation=activation),
            ]
        )
        self.up = nn.ModuleList(
            [
                Up(n_channels * 16, n_channels * 8, activation=activation),
                Up(n_channels * 8, n_channels * 4, activation=activation),
                Up(n_channels * 4, n_channels * 2, activation=activation),
                Up(n_channels * 2, n_channels, activation=activation),
            ]
        )
        out_channels = time_future * (self.n_output_scalar_components + self.n_output_vector_components * 2)
        # should there be a final norm too? but we aren't doing "prenorm" in the original
        self.final = nn.Conv2d(n_channels, out_channels, kernel_size=(3, 3), padding=(1, 1))

    def forward(self, x):
        assert x.dim() == 5
        orig_shape = x.shape
        x = x.reshape(x.size(0), -1, *x.shape[3:])
        h = self.image_proj(x)

        x1 = self.down[0](h)
        x2 = self.down[1](x1)
        x3 = self.down[2](x2)
        x4 = self.down[3](x3)
        x = self.up[0](x4, x3)
        x = self.up[1](x, x2)
        x = self.up[2](x, x1)
        x = self.up[3](x, h)

        x = self.final(x)
        return x.reshape(
            orig_shape[0], -1, (self.n_output_scalar_components + self.n_output_vector_components * 2), *orig_shape[3:]
        )

AttentionBlock ¤

Bases: Module

Attention block This is similar to [transformer multi-head attention]

Parameters:

Name Type Description Default
n_channels int

the number of channels in the input

 required
n_heads int

the number of heads in multi-head attention

 1
d_k Optional[int]

the number of dimensions in each head

 None
n_groups int

the number of groups for group normalization.

 1
Source code in pdearena/modules/twod_unet.py 
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
class AttentionBlock(nn.Module):
    """Attention block This is similar to [transformer multi-head
    attention]

    Args:
        n_channels (int): the number of channels in the input
        n_heads (int): the number of heads in multi-head attention
        d_k: the number of dimensions in each head
        n_groups (int): the number of groups for [group normalization][torch.nn.GroupNorm].

    """

    def __init__(self, n_channels: int, n_heads: int = 1, d_k: Optional[int] = None, n_groups: int = 1):
        super().__init__()

        # Default `d_k`
        if d_k is None:
            d_k = n_channels
        # Normalization layer
        self.norm = nn.GroupNorm(n_groups, n_channels)
        # Projections for query, key and values
        self.projection = nn.Linear(n_channels, n_heads * d_k * 3)
        # Linear layer for final transformation
        self.output = nn.Linear(n_heads * d_k, n_channels)
        # Scale for dot-product attention
        self.scale = d_k**-0.5
        #
        self.n_heads = n_heads
        self.d_k = d_k

    def forward(self, x: torch.Tensor):
        # Get shape
        batch_size, n_channels, height, width = x.shape
        # Change `x` to shape `[batch_size, seq, n_channels]`
        x = x.view(batch_size, n_channels, -1).permute(0, 2, 1)
        # Get query, key, and values (concatenated) and shape it to `[batch_size, seq, n_heads, 3 * d_k]`
        qkv = self.projection(x).view(batch_size, -1, self.n_heads, 3 * self.d_k)
        # Split query, key, and values. Each of them will have shape `[batch_size, seq, n_heads, d_k]`
        q, k, v = torch.chunk(qkv, 3, dim=-1)
        # Calculate scaled dot-product $\frac{Q K^\top}{\sqrt{d_k}}$
        attn = torch.einsum("bihd,bjhd->bijh", q, k) * self.scale
        # Softmax along the sequence dimension $\underset{seq}{softmax}\Bigg(\frac{Q K^\top}{\sqrt{d_k}}\Bigg)$
        attn = attn.softmax(dim=1)
        # Multiply by values
        res = torch.einsum("bijh,bjhd->bihd", attn, v)
        # Reshape to `[batch_size, seq, n_heads * d_k]`
        res = res.view(batch_size, -1, self.n_heads * self.d_k)
        # Transform to `[batch_size, seq, n_channels]`
        res = self.output(res)

        # Add skip connection
        res += x

        # Change to shape `[batch_size, in_channels, height, width]`
        res = res.permute(0, 2, 1).view(batch_size, n_channels, height, width)
        return res

DownBlock ¤

Bases: Module

Down block This combines ResidualBlock and AttentionBlock.

These are used in the first half of U-Net at each resolution.

Parameters:

Name Type Description Default
in_channels int

Number of input channels

 required
out_channels int

Number of output channels

 required
has_attn bool

Whether to use attention block

 False
activation Module

Activation function

 'gelu'
norm bool

Whether to use normalization

 False
Source code in pdearena/modules/twod_unet.py 
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
class DownBlock(nn.Module):
    """Down block This combines [`ResidualBlock`][pdearena.modules.twod_unet.ResidualBlock] and [`AttentionBlock`][pdearena.modules.twod_unet.AttentionBlock].

    These are used in the first half of U-Net at each resolution.

    Args:
        in_channels (int): Number of input channels
        out_channels (int): Number of output channels
        has_attn (bool): Whether to use attention block
        activation (nn.Module): Activation function
        norm (bool): Whether to use normalization
    """

    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        has_attn: bool = False,
        activation: str = "gelu",
        norm: bool = False,
    ):
        super().__init__()
        self.res = ResidualBlock(in_channels, out_channels, activation=activation, norm=norm)
        if has_attn:
            self.attn = AttentionBlock(out_channels)
        else:
            self.attn = nn.Identity()

    def forward(self, x: torch.Tensor):
        x = self.res(x)
        x = self.attn(x)
        return x

Downsample ¤

Bases: Module

Scale down the feature map by \(\frac{1}{2} \times\)

Parameters:

Name Type Description Default
n_channels int

Number of channels in the input and output.

 required
Source code in pdearena/modules/twod_unet.py 
374
375
376
377
378
379
380
381
382
383
384
385
386
class Downsample(nn.Module):
    r"""Scale down the feature map by $\frac{1}{2} \times$

    Args:
        n_channels (int): Number of channels in the input and output.
    """

    def __init__(self, n_channels):
        super().__init__()
        self.conv = nn.Conv2d(n_channels, n_channels, (3, 3), (2, 2), (1, 1))

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

FourierDownBlock ¤

Bases: Module

Down block This combines FourierResidualBlock and AttentionBlock.

These are used in the first half of U-Net at each resolution.

Source code in pdearena/modules/twod_unet.py 
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
class FourierDownBlock(nn.Module):
    """Down block This combines [`FourierResidualBlock`][pdearena.modules.twod_unet.FourierResidualBlock] and [`AttentionBlock`][pdearena.modules.twod_unet.AttentionBlock].

    These are used in the first half of U-Net at each resolution.
    """

    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        modes1: int = 16,
        modes2: int = 16,
        has_attn: bool = False,
        activation: str = "gelu",
        norm: bool = False,
    ):
        super().__init__()
        self.res = FourierResidualBlock(
            in_channels,
            out_channels,
            modes1=modes1,
            modes2=modes2,
            activation=activation,
            norm=norm,
        )
        if has_attn:
            self.attn = AttentionBlock(out_channels)
        else:
            self.attn = nn.Identity()

    def forward(self, x: torch.Tensor):
        x = self.res(x)
        x = self.attn(x)
        return x

FourierResidualBlock ¤

Bases: Module

Fourier Residual Block to be used in modern Unet architectures.

Parameters:

Name Type Description Default
in_channels int

Number of input channels.

 required
out_channels int

Number of output channels.

 required
modes1 int

Number of modes in the first dimension.

 16
modes2 int

Number of modes in the second dimension.

 16
activation str

Activation function to use.

 'gelu'
norm bool

Whether to use normalization.

 False
n_groups int

Number of groups for group normalization.

 1
Source code in pdearena/modules/twod_unet.py 
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
class FourierResidualBlock(nn.Module):
    """Fourier Residual Block to be used in modern Unet architectures.

    Args:
        in_channels (int): Number of input channels.
        out_channels (int): Number of output channels.
        modes1 (int): Number of modes in the first dimension.
        modes2 (int): Number of modes in the second dimension.
        activation (str): Activation function to use.
        norm (bool): Whether to use normalization.
        n_groups (int): Number of groups for group normalization.
    """

    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        modes1: int = 16,
        modes2: int = 16,
        activation: str = "gelu",
        norm: bool = False,
        n_groups: int = 1,
    ):
        super().__init__()
        self.activation: nn.Module = ACTIVATION_REGISTRY.get(activation, None)
        if self.activation is None:
            raise NotImplementedError(f"Activation {activation} not implemented")

        self.modes1 = modes1
        self.modes2 = modes2

        self.fourier1 = SpectralConv2d(in_channels, out_channels, modes1=self.modes1, modes2=self.modes2)
        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=1, padding=0, padding_mode="zeros")
        self.fourier2 = SpectralConv2d(out_channels, out_channels, modes1=self.modes1, modes2=self.modes2)
        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=1, padding=0, padding_mode="zeros")
        # If the number of input channels is not equal to the number of output channels we have to
        # project the shortcut connection
        if in_channels != out_channels:
            self.shortcut = nn.Conv2d(in_channels, out_channels, kernel_size=(1, 1))
        else:
            self.shortcut = nn.Identity()

        if norm:
            self.norm1 = nn.GroupNorm(n_groups, in_channels)
            self.norm2 = nn.GroupNorm(n_groups, out_channels)
        else:
            self.norm1 = nn.Identity()
            self.norm2 = nn.Identity()

    def forward(self, x: torch.Tensor):
        # using pre-norms
        h = self.activation(self.norm1(x))
        x1 = self.fourier1(h)
        x2 = self.conv1(h)
        out = x1 + x2
        out = self.activation(self.norm2(out))
        x1 = self.fourier2(out)
        x2 = self.conv2(out)
        out = x1 + x2 + self.shortcut(x)
        return out

FourierUnet ¤

Bases: Module

Unet with Fourier layers in early downsampling blocks.

Parameters:

Name Type Description Default
n_input_scalar_components int

Number of scalar components in the model

 required
n_input_vector_components int

Number of vector components in the model

 required
n_output_scalar_components int

Number of output scalar components in the model

 required
n_output_vector_components int

Number of output vector components in the model

 required
time_history int

Number of time steps in the input.

 required
time_future int

Number of time steps in the output.

 required
hidden_channels int

Number of channels in the first layer.

 required
activation str

Activation function to use.

 required
modes1 int

Number of Fourier modes to use in the first spatial dimension.

 12
modes2 int

Number of Fourier modes to use in the second spatial dimension.

 12
norm bool

Whether to use normalization.

 False
ch_mults list

List of integers to multiply the number of channels by at each resolution.

 (1, 2, 2, 4)
is_attn list

List of booleans indicating whether to use attention at each resolution.

 (False, False, False, False)
mid_attn bool

Whether to use attention in the middle block.

 False
n_blocks int

Number of blocks to use at each resolution.

 2
n_fourier_layers int

Number of early downsampling layers to use Fourier layers in.

 2
mode_scaling bool

Whether to scale the number of modes with resolution.

 True
use1x1 bool

Whether to use 1x1 convolutions in the initial and final layer.

 False
Source code in pdearena/modules/twod_unet.py 
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
class FourierUnet(nn.Module):
    """Unet with Fourier layers in early downsampling blocks.

    Args:
        n_input_scalar_components (int): Number of scalar components in the model
        n_input_vector_components (int): Number of vector components in the model
        n_output_scalar_components (int): Number of output scalar components in the model
        n_output_vector_components (int): Number of output vector components in the model
        time_history (int): Number of time steps in the input.
        time_future (int): Number of time steps in the output.
        hidden_channels (int): Number of channels in the first layer.
        activation (str): Activation function to use.
        modes1 (int): Number of Fourier modes to use in the first spatial dimension.
        modes2 (int): Number of Fourier modes to use in the second spatial dimension.
        norm (bool): Whether to use normalization.
        ch_mults (list): List of integers to multiply the number of channels by at each resolution.
        is_attn (list): List of booleans indicating whether to use attention at each resolution.
        mid_attn (bool): Whether to use attention in the middle block.
        n_blocks (int): Number of blocks to use at each resolution.
        n_fourier_layers (int): Number of early downsampling layers to use Fourier layers in.
        mode_scaling (bool): Whether to scale the number of modes with resolution.
        use1x1 (bool): Whether to use 1x1 convolutions in the initial and final layer.
    """

    def __init__(
        self,
        n_input_scalar_components: int,
        n_input_vector_components: int,
        n_output_scalar_components: int,
        n_output_vector_components: int,
        time_history: int,
        time_future: int,
        hidden_channels: int,
        activation: str,
        modes1: int = 12,
        modes2: int = 12,
        norm: bool = False,
        ch_mults: Union[Tuple[int, ...], List[int]] = (1, 2, 2, 4),
        is_attn: Union[Tuple[bool, ...], List[bool]] = (False, False, False, False),
        mid_attn: bool = False,
        n_blocks: int = 2,
        n_fourier_layers: int = 2,
        mode_scaling: bool = True,
        use1x1: bool = False,
    ) -> None:
        super().__init__()
        self.n_input_scalar_components = n_input_scalar_components
        self.n_input_vector_components = n_input_vector_components
        self.n_output_scalar_components = n_output_scalar_components
        self.n_output_vector_components = n_output_vector_components
        self.time_history = time_history
        self.time_future = time_future
        self.hidden_channels = hidden_channels
        self.activation: nn.Module = ACTIVATION_REGISTRY.get(activation, None)
        if self.activation is None:
            raise NotImplementedError(f"Activation {activation} not implemented")
        # Number of resolutions
        n_resolutions = len(ch_mults)

        insize = time_history * (self.n_input_scalar_components + self.n_input_vector_components * 2)
        n_channels = hidden_channels
        # Project image into feature map
        if use1x1:
            self.image_proj = nn.Conv2d(insize, n_channels, kernel_size=1)
        else:
            self.image_proj = nn.Conv2d(insize, n_channels, kernel_size=(3, 3), padding=(1, 1))

        # #### First half of U-Net - decreasing resolution
        down = []
        # Number of channels
        out_channels = in_channels = n_channels
        # For each resolution
        for i in range(n_resolutions):
            # Number of output channels at this resolution
            out_channels = in_channels * ch_mults[i]
            if i < n_fourier_layers:
                for _ in range(n_blocks):
                    down.append(
                        FourierDownBlock(
                            in_channels,
                            out_channels,
                            modes1=max(modes1 // 2**i, 4) if mode_scaling else modes1,
                            modes2=max(modes2 // 2**i, 4) if mode_scaling else modes2,
                            has_attn=is_attn[i],
                            activation=activation,
                            norm=norm,
                        )
                    )
                    in_channels = out_channels
            else:
                # Add `n_blocks`
                for _ in range(n_blocks):
                    down.append(
                        DownBlock(
                            in_channels,
                            out_channels,
                            has_attn=is_attn[i],
                            activation=activation,
                            norm=norm,
                        )
                    )
                    in_channels = out_channels
            # Down sample at all resolutions except the last
            if i < n_resolutions - 1:
                down.append(Downsample(in_channels))

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

        # Middle block
        self.middle = MiddleBlock(out_channels, has_attn=mid_attn, activation=activation, norm=norm)

        # #### Second half of U-Net - 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(
                    UpBlock(
                        in_channels,
                        out_channels,
                        has_attn=is_attn[i],
                        activation=activation,
                        norm=norm,
                    )
                )
            # Final block to reduce the number of channels
            out_channels = in_channels // ch_mults[i]
            up.append(UpBlock(in_channels, out_channels, has_attn=is_attn[i], activation=activation, norm=norm))
            in_channels = out_channels
            # Up sample at all resolutions except last
            if i > 0:
                up.append(Upsample(in_channels))

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

        if norm:
            self.norm = nn.GroupNorm(8, n_channels)
        else:
            self.norm = nn.Identity()
        out_channels = time_future * (self.n_output_scalar_components + self.n_output_vector_components * 2)
        if use1x1:
            self.final = nn.Conv2d(n_channels, out_channels, kernel_size=1)
        else:
            self.final = nn.Conv2d(in_channels, out_channels, kernel_size=(3, 3), padding=(1, 1))

    def forward(self, x: torch.Tensor):
        assert x.dim() == 5
        orig_shape = x.shape
        x = x.reshape(x.size(0), -1, *x.shape[3:])  # collapse T,C
        x = self.image_proj(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, Upsample):
                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.final(self.activation(self.norm(x)))
        return x.reshape(
            orig_shape[0], -1, (self.n_output_scalar_components + self.n_output_vector_components * 2), *orig_shape[3:]
        )

FourierUpBlock ¤

Bases: Module

Up block that combines FourierResidualBlock and AttentionBlock.

These are used in the second half of U-Net at each resolution.

Note

We currently don't recommend using this block.

Source code in pdearena/modules/twod_unet.py 
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
class FourierUpBlock(nn.Module):
    """Up block that combines [`FourierResidualBlock`][pdearena.modules.twod_unet.FourierResidualBlock] and [`AttentionBlock`][pdearena.modules.twod_unet.AttentionBlock].

    These are used in the second half of U-Net at each resolution.

    Note:
        We currently don't recommend using this block.
    """

    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        modes1: int = 16,
        modes2: int = 16,
        has_attn: bool = False,
        activation: str = "gelu",
        norm: bool = False,
    ):
        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 = FourierResidualBlock(
            in_channels + out_channels,
            out_channels,
            modes1=modes1,
            modes2=modes2,
            activation=activation,
            norm=norm,
        )
        if has_attn:
            self.attn = AttentionBlock(out_channels)
        else:
            self.attn = nn.Identity()

    def forward(self, x: torch.Tensor):
        x = self.res(x)
        x = self.attn(x)
        return x

MiddleBlock ¤

Bases: Module

Middle block

It combines a ResidualBlock, AttentionBlock, followed by another ResidualBlock.

This block is applied at the lowest resolution of the U-Net.

Parameters:

Name Type Description Default
n_channels int

Number of channels in the input and output.

 required
has_attn bool

Whether to use attention block. Defaults to False.

 False
activation str

Activation function to use. Defaults to "gelu".

 'gelu'
norm bool

Whether to use normalization. Defaults to False.

 False
Source code in pdearena/modules/twod_unet.py 
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
class MiddleBlock(nn.Module):
    """Middle block

    It combines a `ResidualBlock`, `AttentionBlock`, followed by another
    `ResidualBlock`.

    This block is applied at the lowest resolution of the U-Net.

    Args:
        n_channels (int): Number of channels in the input and output.
        has_attn (bool, optional): Whether to use attention block. Defaults to False.
        activation (str): Activation function to use. Defaults to "gelu".
        norm (bool, optional): Whether to use normalization. Defaults to False.
    """

    def __init__(self, n_channels: int, has_attn: bool = False, activation: str = "gelu", norm: bool = False):
        super().__init__()
        self.res1 = ResidualBlock(n_channels, n_channels, activation=activation, norm=norm)
        self.attn = AttentionBlock(n_channels) if has_attn else nn.Identity()
        self.res2 = ResidualBlock(n_channels, n_channels, activation=activation, norm=norm)

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

ResidualBlock ¤

Bases: Module

Wide Residual Blocks used in modern Unet architectures.

Parameters:

Name Type Description Default
in_channels int

Number of input channels.

 required
out_channels int

Number of output channels.

 required
activation str

Activation function to use.

 'gelu'
norm bool

Whether to use normalization.

 False
n_groups int

Number of groups for group normalization.

 1
Source code in pdearena/modules/twod_unet.py 
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
class ResidualBlock(nn.Module):
    """Wide Residual Blocks used in modern Unet architectures.

    Args:
        in_channels (int): Number of input channels.
        out_channels (int): Number of output channels.
        activation (str): Activation function to use.
        norm (bool): Whether to use normalization.
        n_groups (int): Number of groups for group normalization.
    """

    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        activation: str = "gelu",
        norm: bool = False,
        n_groups: int = 1,
    ):
        super().__init__()
        self.activation: nn.Module = ACTIVATION_REGISTRY.get(activation, None)
        if self.activation is None:
            raise NotImplementedError(f"Activation {activation} not implemented")
        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=(3, 3), padding=(1, 1))
        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=(3, 3), padding=(1, 1))
        # If the number of input channels is not equal to the number of output channels we have to
        # project the shortcut connection
        if in_channels != out_channels:
            self.shortcut = nn.Conv2d(in_channels, out_channels, kernel_size=(1, 1))
        else:
            self.shortcut = nn.Identity()

        if norm:
            self.norm1 = nn.GroupNorm(n_groups, in_channels)
            self.norm2 = nn.GroupNorm(n_groups, out_channels)
        else:
            self.norm1 = nn.Identity()
            self.norm2 = nn.Identity()

    def forward(self, x: torch.Tensor):
        # First convolution layer
        h = self.conv1(self.activation(self.norm1(x)))
        # Second convolution layer
        h = self.conv2(self.activation(self.norm2(h)))
        # Add the shortcut connection and return
        return h + self.shortcut(x)

Unet ¤

Bases: Module

Modern U-Net architecture

This is a modern U-Net architecture with wide-residual blocks and spatial attention blocks

Parameters:

Name Type Description Default
n_input_scalar_components int

Number of scalar components in the model

 required
n_input_vector_components int

Number of vector components in the model

 required
n_output_scalar_components int

Number of output scalar components in the model

 required
n_output_vector_components int

Number of output vector components in the model

 required
time_history int

Number of time steps in the input

 required
time_future int

Number of time steps in the output

 required
hidden_channels int

Number of channels in the hidden layers

 required
activation str

Activation function to use

 required
norm bool

Whether to use normalization

 False
ch_mults list

List of channel multipliers for each resolution

 (1, 2, 2, 4)
is_attn list

List of booleans indicating whether to use attention blocks

 (False, False, False, False)
mid_attn bool

Whether to use attention block in the middle block

 False
n_blocks int

Number of residual blocks in each resolution

 2
use1x1 bool

Whether to use 1x1 convolutions in the initial and final layers

 False
Source code in pdearena/modules/twod_unet.py 
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
class Unet(nn.Module):
    """Modern U-Net architecture

    This is a modern U-Net architecture with wide-residual blocks and spatial attention blocks

    Args:
        n_input_scalar_components (int): Number of scalar components in the model
        n_input_vector_components (int): Number of vector components in the model
        n_output_scalar_components (int): Number of output scalar components in the model
        n_output_vector_components (int): Number of output vector components in the model
        time_history (int): Number of time steps in the input
        time_future (int): Number of time steps in the output
        hidden_channels (int): Number of channels in the hidden layers
        activation (str): Activation function to use
        norm (bool): Whether to use normalization
        ch_mults (list): List of channel multipliers for each resolution
        is_attn (list): List of booleans indicating whether to use attention blocks
        mid_attn (bool): Whether to use attention block in the middle block
        n_blocks (int): Number of residual blocks in each resolution
        use1x1 (bool): Whether to use 1x1 convolutions in the initial and final layers
    """

    def __init__(
        self,
        n_input_scalar_components: int,
        n_input_vector_components: int,
        n_output_scalar_components: int,
        n_output_vector_components: int,
        time_history: int,
        time_future: int,
        hidden_channels: int,
        activation: str,
        norm: bool = False,
        ch_mults: Union[Tuple[int, ...], List[int]] = (1, 2, 2, 4),
        is_attn: Union[Tuple[bool, ...], List[bool]] = (False, False, False, False),
        mid_attn: bool = False,
        n_blocks: int = 2,
        use1x1: bool = False,
    ) -> None:
        super().__init__()
        self.n_input_scalar_components = n_input_scalar_components
        self.n_input_vector_components = n_input_vector_components
        self.n_output_scalar_components = n_output_scalar_components
        self.n_output_vector_components = n_output_vector_components
        self.time_history = time_history
        self.time_future = time_future
        self.hidden_channels = hidden_channels

        self.activation: nn.Module = ACTIVATION_REGISTRY.get(activation, None)
        if self.activation is None:
            raise NotImplementedError(f"Activation {activation} not implemented")
        # Number of resolutions
        n_resolutions = len(ch_mults)

        insize = time_history * (self.n_input_scalar_components + self.n_input_vector_components * 2)
        n_channels = hidden_channels
        # Project image into feature map
        if use1x1:
            self.image_proj = nn.Conv2d(insize, n_channels, kernel_size=1)
        else:
            self.image_proj = nn.Conv2d(insize, n_channels, kernel_size=(3, 3), padding=(1, 1))

        # #### First half of U-Net - decreasing resolution
        down = []
        # Number of channels
        out_channels = in_channels = n_channels
        # For each resolution
        for i in range(n_resolutions):
            # Number of output channels at this resolution
            out_channels = in_channels * ch_mults[i]
            # Add `n_blocks`
            for _ in range(n_blocks):
                down.append(
                    DownBlock(
                        in_channels,
                        out_channels,
                        has_attn=is_attn[i],
                        activation=activation,
                        norm=norm,
                    )
                )
                in_channels = out_channels
            # Down sample at all resolutions except the last
            if i < n_resolutions - 1:
                down.append(Downsample(in_channels))

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

        # Middle block
        self.middle = MiddleBlock(out_channels, has_attn=mid_attn, activation=activation, norm=norm)

        # #### Second half of U-Net - 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(
                    UpBlock(
                        in_channels,
                        out_channels,
                        has_attn=is_attn[i],
                        activation=activation,
                        norm=norm,
                    )
                )
            # Final block to reduce the number of channels
            out_channels = in_channels // ch_mults[i]
            up.append(UpBlock(in_channels, out_channels, has_attn=is_attn[i], activation=activation, norm=norm))
            in_channels = out_channels
            # Up sample at all resolutions except last
            if i > 0:
                up.append(Upsample(in_channels))

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

        if norm:
            self.norm = nn.GroupNorm(8, n_channels)
        else:
            self.norm = nn.Identity()
        out_channels = time_future * (self.n_output_scalar_components + self.n_output_vector_components * 2)
        #
        if use1x1:
            self.final = nn.Conv2d(in_channels, out_channels, kernel_size=1)
        else:
            self.final = nn.Conv2d(in_channels, out_channels, kernel_size=(3, 3), padding=(1, 1))

    def forward(self, x: torch.Tensor):
        assert x.dim() == 5
        orig_shape = x.shape
        x = x.reshape(x.size(0), -1, *x.shape[3:])  # collapse T,C
        x = self.image_proj(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, Upsample):
                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.final(self.activation(self.norm(x)))
        x = x.reshape(
            orig_shape[0], -1, (self.n_output_scalar_components + self.n_output_vector_components * 2), *orig_shape[3:]
        )
        return x

UpBlock ¤

Bases: Module

Up block that combines ResidualBlock and AttentionBlock.

These are used in the second half of U-Net at each resolution.

Parameters:

Name Type Description Default
in_channels int

Number of input channels

 required
out_channels int

Number of output channels

 required
has_attn bool

Whether to use attention block

 False
activation str

Activation function

 'gelu'
norm bool

Whether to use normalization

 False
Source code in pdearena/modules/twod_unet.py 
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
class UpBlock(nn.Module):
    """Up block that combines [`ResidualBlock`][pdearena.modules.twod_unet.ResidualBlock] and [`AttentionBlock`][pdearena.modules.twod_unet.AttentionBlock].

    These are used in the second half of U-Net at each resolution.

    Args:
        in_channels (int): Number of input channels
        out_channels (int): Number of output channels
        has_attn (bool): Whether to use attention block
        activation (str): Activation function
        norm (bool): Whether to use normalization
    """

    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        has_attn: bool = False,
        activation: str = "gelu",
        norm: bool = False,
    ):
        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 = ResidualBlock(in_channels + out_channels, out_channels, activation=activation, norm=norm)
        if has_attn:
            self.attn = AttentionBlock(out_channels)
        else:
            self.attn = nn.Identity()

    def forward(self, x: torch.Tensor):
        x = self.res(x)
        x = self.attn(x)
        return x

Upsample ¤

Bases: Module

Scale up the feature map by \(2 \times\)

Parameters:

Name Type Description Default
n_channels int

Number of channels in the input and output.

 required
Source code in pdearena/modules/twod_unet.py 
359
360
361
362
363
364
365
366
367
368
369
370
371
class Upsample(nn.Module):
    r"""Scale up the feature map by $2 \times$

    Args:
        n_channels (int): Number of channels in the input and output.
    """

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

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

CliffordMaxwellResNet3D ¤

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
time_history int

Number of input timesteps.

 required
time_future int

Number of output timesteps.

 required
hidden_channels int

Number of hidden channels.

 required
activation Callable

Activation function. Defaults to F.gelu.

 'gelu'
norm bool

Whether 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 pdearena/modules/threed.py 
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
class CliffordMaxwellResNet3D(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.
        time_history (int): Number of input timesteps.
        time_future (int): Number of output timesteps.
        hidden_channels (int): Number of hidden channels.
        activation (Callable, optional): Activation function. Defaults to F.gelu.
        norm (bool, optional): Whether 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,
        time_history: int,
        time_future: int,
        hidden_channels: int,
        activation: str = "gelu",
        norm: bool = False,
        num_groups: int = 1,
        diffmode: bool = False,
    ):
        super().__init__()

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

        self.activation: nn.Module = ACTIVATION_REGISTRY.get(activation, None)
        if self.activation is None:
            raise NotImplementedError(f"Activation {activation} not implemented")

        # Residual blocks.
        self.layers = nn.ModuleList(
            [
                self._make_basic_block(
                    g,
                    block,
                    hidden_channels,
                    num_blocks[i],
                    activation=self.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

        # Get data into shape where I dimension is last.
        B_dim, C_dim, I_dim, *D_dims = range(len(x.shape))
        x = x.permute(B_dim, C_dim, *D_dims, I_dim)

        # 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)

        # Get data back to normal shape.
        B_dim, C_dim, *D_dims, I_dim = range(len(x.shape))
        x = x.permute(B_dim, C_dim, I_dim, *D_dims)

        return x

FourierBasicBlock3D ¤

Bases: Module

Basic 3d FNO ResNet building block consisting of two 3d convolutional layers, two 3d SpectralConv layers and skip connections.

Source code in pdearena/modules/threed.py 
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
class FourierBasicBlock3D(nn.Module):
    """Basic 3d FNO ResNet building block consisting of two 3d convolutional layers,
    two 3d SpectralConv layers and skip connections."""

    expansion: int = 1

    def __init__(
        self,
        in_planes: int,
        planes: int,
        modes1: int,
        modes2: int,
        modes3: int,
        stride: int = 1,
        activation: str = "gelu",
        norm: bool = False,
    ):
        """Initialize basic 3d FNO ResNet building block
        Args:
            in_planes (int): Input channels
            planes (int): Output channels
            modes1 (int): Fourier modes for x direction.
            modes2 (int): Fourier modes for y direction.
            modes3 (int): Fourier modes for z direction.
            stride (int, optional): stride of 2d convolution. Defaults to 1.
            norm (bool): Wether to use normalization. Defaults to False.
        """
        super().__init__()

        self.fourier1 = SpectralConv3d(in_planes, planes, modes1=modes1, modes2=modes2, modes3=modes3)
        self.conv1 = nn.Conv3d(in_planes, planes, kernel_size=1, padding=0, padding_mode="zeros")
        self.fourier2 = SpectralConv3d(planes, planes, modes1=modes1, modes2=modes2, modes3=modes3)
        self.conv2 = nn.Conv3d(planes, planes, kernel_size=1, padding=0, padding_mode="zeros")

        # Shortcut connection, batchnorm removed
        # So far shortcut connections are not helping
        self.shortcut = nn.Sequential()
        if stride != 1 or in_planes != self.expansion * planes:
            self.shortcut = nn.Sequential(nn.Conv2d(in_planes, self.expansion * planes, kernel_size=1))
        if activation == "gelu":
            self.activation = F.gelu
        elif activation == "relu":
            self.activation = F.relu
        else:
            raise NotImplementedError(f"Activation {activation} not implemented")

        if norm:
            raise NotImplementedError(f"Normalization for FourierBasicBlock3D not implemented")

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Forward pass of basic 3d Fourier ResNet building block.
        Args:
            x (torch.Tensor): input of shape [batch, in_planes, x, y, z]
        Returns:
            torch.Tensor: output of shape [batch, planes, x, y, z]
        """
        x1 = self.fourier1(x)
        x2 = self.conv1(x)
        out = self.activation(x1 + x2)

        x1 = self.fourier2(out)
        x2 = self.conv2(out)
        out = x1 + x2
        # out += self.shortcut(x)
        out = self.activation(out)
        return out

__init__(in_planes, planes, modes1, modes2, modes3, stride=1, activation='gelu', norm=False) ¤

Initialize basic 3d FNO ResNet building block Args: in_planes (int): Input channels planes (int): Output channels modes1 (int): Fourier modes for x direction. modes2 (int): Fourier modes for y direction. modes3 (int): Fourier modes for z direction. stride (int, optional): stride of 2d convolution. Defaults to 1. norm (bool): Wether to use normalization. Defaults to False.

Source code in pdearena/modules/threed.py 
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
def __init__(
    self,
    in_planes: int,
    planes: int,
    modes1: int,
    modes2: int,
    modes3: int,
    stride: int = 1,
    activation: str = "gelu",
    norm: bool = False,
):
    """Initialize basic 3d FNO ResNet building block
    Args:
        in_planes (int): Input channels
        planes (int): Output channels
        modes1 (int): Fourier modes for x direction.
        modes2 (int): Fourier modes for y direction.
        modes3 (int): Fourier modes for z direction.
        stride (int, optional): stride of 2d convolution. Defaults to 1.
        norm (bool): Wether to use normalization. Defaults to False.
    """
    super().__init__()

    self.fourier1 = SpectralConv3d(in_planes, planes, modes1=modes1, modes2=modes2, modes3=modes3)
    self.conv1 = nn.Conv3d(in_planes, planes, kernel_size=1, padding=0, padding_mode="zeros")
    self.fourier2 = SpectralConv3d(planes, planes, modes1=modes1, modes2=modes2, modes3=modes3)
    self.conv2 = nn.Conv3d(planes, planes, kernel_size=1, padding=0, padding_mode="zeros")

    # Shortcut connection, batchnorm removed
    # So far shortcut connections are not helping
    self.shortcut = nn.Sequential()
    if stride != 1 or in_planes != self.expansion * planes:
        self.shortcut = nn.Sequential(nn.Conv2d(in_planes, self.expansion * planes, kernel_size=1))
    if activation == "gelu":
        self.activation = F.gelu
    elif activation == "relu":
        self.activation = F.relu
    else:
        raise NotImplementedError(f"Activation {activation} not implemented")

    if norm:
        raise NotImplementedError(f"Normalization for FourierBasicBlock3D not implemented")

forward(x) ¤

Forward pass of basic 3d Fourier ResNet building block. Args: x (torch.Tensor): input of shape [batch, in_planes, x, y, z] Returns: torch.Tensor: output of shape [batch, planes, x, y, z]

Source code in pdearena/modules/threed.py 
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
def forward(self, x: torch.Tensor) -> torch.Tensor:
    """Forward pass of basic 3d Fourier ResNet building block.
    Args:
        x (torch.Tensor): input of shape [batch, in_planes, x, y, z]
    Returns:
        torch.Tensor: output of shape [batch, planes, x, y, z]
    """
    x1 = self.fourier1(x)
    x2 = self.conv1(x)
    out = self.activation(x1 + x2)

    x1 = self.fourier2(out)
    x2 = self.conv2(out)
    out = x1 + x2
    # out += self.shortcut(x)
    out = self.activation(out)
    return out

MaxwellResNet3D ¤

Bases: Module

3d ResNet model for Maxwell equations, difference to default ResNet architectures is that spatial resolution and channels (in_planes) stay constant throughout the network.

Source code in pdearena/modules/threed.py 
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
class MaxwellResNet3D(nn.Module):
    """3d ResNet model for Maxwell equations, difference to default ResNet architectures is that
    spatial resolution and channels (in_planes) stay constant throughout the network."""

    padding = 2  # no periodic

    def __init__(
        self,
        block: nn.Module,
        num_blocks: list,
        time_history: int,
        time_future: int,
        hidden_channels: int = 64,
        activation: str = "gelu",
        diffmode: bool = False,
    ):
        """Initialize 3d ResNet model

        Args:
            block (nn.Module): basic 3d ResNet building block
            num_blocks (list): list of basic building blocks per layer
            time_history (int): input timesteps
            time_future (int): prediction timesteps
            hidden_channels (int): hidden channels in the ResNet blocks
        """
        super().__init__()

        self.diffmode = diffmode
        self.in_planes = hidden_channels
        self.conv_in1 = nn.Conv3d(
            time_history * 6,
            self.in_planes,
            kernel_size=1,
            bias=True,
        )
        self.conv_in2 = nn.Conv3d(
            self.in_planes,
            self.in_planes,
            kernel_size=1,
            bias=True,
        )
        self.conv_out1 = nn.Conv3d(
            self.in_planes,
            self.in_planes,
            kernel_size=1,
            bias=True,
        )
        self.conv_out2 = nn.Conv3d(
            self.in_planes,
            time_future * 6,
            kernel_size=1,
            bias=True,
        )

        self.layers = nn.ModuleList(
            [
                self._make_layer(
                    block,
                    self.in_planes,
                    num_blocks[i],
                    stride=1,
                    activation=activation,
                )
                for i in range(len(num_blocks))
            ]
        )
        self.activation: nn.Module = ACTIVATION_REGISTRY.get(activation, None)
        if self.activation is None:
            raise NotImplementedError(f"Activation {activation} not implemented")

    def _make_layer(
        self,
        block: Callable,
        planes: int,
        num_blocks: int,
        stride: int,
        activation: str,
    ) -> nn.Sequential:
        """Build 3d ResNet layers out of basic building blocks.

        Args:
            block (nn.Module): basic 3d ResNet building block
            planes (int): input channels
            num_blocks (int): number of basic 3d ResNet building blocks in one layer
            stride (int): stride

        Returns:
            nn.Sequential: 3d ResNet layer as nn.Sequential
        """
        strides = [stride] + [1] * (num_blocks - 1)
        layers = []
        for stride in strides:
            layers.append(block(in_planes=self.in_planes, planes=planes, stride=stride, activation=activation))
            self.in_planes = planes * block.expansion
        return nn.Sequential(*layers)

    def __repr__(self):
        return "ResNet"

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Forward pass of the ResNet in 3 spatial dimensions
        consisting of embedding layer(s), ResNet building blogs and output layer.
        Args:
            x (torch.Tensor): input tensor of shape [b, time_history, 6, x, y, z]
        Returns:
            torch.Tensor: output tensor of shape [b, time_future, 6, x, y, z]
        """
        assert x.dim() == 6
        orig_shape = x.shape
        x = x.reshape(x.size(0), -1, *x.shape[3:])
        prev = x
        x = self.activation(self.conv_in1(x.float()))
        x = self.activation(self.conv_in2(x.float()))

        if self.padding > 0:
            x = F.pad(x, [0, self.padding, 0, self.padding])
        for layer in self.layers:
            x = layer(x)
        if self.padding > 0:
            x = x[..., : -self.padding, : -self.padding]

        x = self.conv_out1(x)
        x = self.conv_out2(x)

        if self.diffmode:
            x = x + prev[:, -1:, ...].detach()
        return x.reshape(orig_shape[0], -1, 6, *orig_shape[3:])

__init__(block, num_blocks, time_history, time_future, hidden_channels=64, activation='gelu', diffmode=False) ¤

Initialize 3d ResNet model

Parameters:

Name Type Description Default
block Module

basic 3d ResNet building block

 required
num_blocks list

list of basic building blocks per layer

 required
time_history int

input timesteps

 required
time_future int

prediction timesteps

 required
hidden_channels int

hidden channels in the ResNet blocks

 64
Source code in pdearena/modules/threed.py 
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
def __init__(
    self,
    block: nn.Module,
    num_blocks: list,
    time_history: int,
    time_future: int,
    hidden_channels: int = 64,
    activation: str = "gelu",
    diffmode: bool = False,
):
    """Initialize 3d ResNet model

    Args:
        block (nn.Module): basic 3d ResNet building block
        num_blocks (list): list of basic building blocks per layer
        time_history (int): input timesteps
        time_future (int): prediction timesteps
        hidden_channels (int): hidden channels in the ResNet blocks
    """
    super().__init__()

    self.diffmode = diffmode
    self.in_planes = hidden_channels
    self.conv_in1 = nn.Conv3d(
        time_history * 6,
        self.in_planes,
        kernel_size=1,
        bias=True,
    )
    self.conv_in2 = nn.Conv3d(
        self.in_planes,
        self.in_planes,
        kernel_size=1,
        bias=True,
    )
    self.conv_out1 = nn.Conv3d(
        self.in_planes,
        self.in_planes,
        kernel_size=1,
        bias=True,
    )
    self.conv_out2 = nn.Conv3d(
        self.in_planes,
        time_future * 6,
        kernel_size=1,
        bias=True,
    )

    self.layers = nn.ModuleList(
        [
            self._make_layer(
                block,
                self.in_planes,
                num_blocks[i],
                stride=1,
                activation=activation,
            )
            for i in range(len(num_blocks))
        ]
    )
    self.activation: nn.Module = ACTIVATION_REGISTRY.get(activation, None)
    if self.activation is None:
        raise NotImplementedError(f"Activation {activation} not implemented")

forward(x) ¤

Forward pass of the ResNet in 3 spatial dimensions consisting of embedding layer(s), ResNet building blogs and output layer. Args: x (torch.Tensor): input tensor of shape [b, time_history, 6, x, y, z] Returns: torch.Tensor: output tensor of shape [b, time_future, 6, x, y, z]

Source code in pdearena/modules/threed.py 
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
def forward(self, x: torch.Tensor) -> torch.Tensor:
    """Forward pass of the ResNet in 3 spatial dimensions
    consisting of embedding layer(s), ResNet building blogs and output layer.
    Args:
        x (torch.Tensor): input tensor of shape [b, time_history, 6, x, y, z]
    Returns:
        torch.Tensor: output tensor of shape [b, time_future, 6, x, y, z]
    """
    assert x.dim() == 6
    orig_shape = x.shape
    x = x.reshape(x.size(0), -1, *x.shape[3:])
    prev = x
    x = self.activation(self.conv_in1(x.float()))
    x = self.activation(self.conv_in2(x.float()))

    if self.padding > 0:
        x = F.pad(x, [0, self.padding, 0, self.padding])
    for layer in self.layers:
        x = layer(x)
    if self.padding > 0:
        x = x[..., : -self.padding, : -self.padding]

    x = self.conv_out1(x)
    x = self.conv_out2(x)

    if self.diffmode:
        x = x + prev[:, -1:, ...].detach()
    return x.reshape(orig_shape[0], -1, 6, *orig_shape[3:])