navier_stokes_uno3d.py

# Codes for section: Results on Navier Stocks Equation (3D)
import torch
import numpy as np
import torch.nn as nn
import torch.nn.functional as F
from integral_operators import *
import matplotlib.pyplot as plt
from utilities3 import *

import operator
from functools import reduce
from functools import partial

from timeit import default_timer

from Adam import Adam

torch.manual_seed(0)
np.random.seed(0)


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

    input: the solution of the first 10 timesteps (u(1), ..., u(10)).
    input shape: (batchsize, x=S, y=S, t=T, c=1)
    output: the solution of the next 40 timesteps
    output shape: (batchsize, x=S, y=S, t=4*T, c=1)
    S,S,T = grid size along x,y and t (of input function)
    S,S,4*T = grid size along x,y and t (of output function)
    (Note that model is dicretization invarient in both spatial (x,y) and time (t) domain)

    in_width = 4; [a(x,y,x),x,y,z]
    with = uplifting dimesion
    pad = padding amount
    pad_both = boolean, if true pad both size of the domain
    factor = scaling factor of the co-domain dimesions
    """

    def __init__(self, in_width, width, pad=2, factor=1, pad_both=False):
        super(Uno3D_T40, self).__init__()

        self.in_width = in_width  # input channel
        self.width = width

        self.pad = pad
        self.pad_both = pad_both

        self.fc = nn.Linear(self.in_width, self.width // 2)

        self.fc0 = nn.Linear(self.width // 2, self.width)

        self.conv0 = OperatorBlock_3D(
            self.width, 2 * factor * self.width, 48, 48, 10, 20, 20, 4, Normalize=True
        )

        self.conv1 = OperatorBlock_3D(
            2 * factor * self.width, 4 * factor * self.width, 32, 32, 10, 14, 14, 4
        )

        self.conv2 = OperatorBlock_3D(
            4 * factor * self.width, 8 * factor * self.width, 16, 16, 16, 6, 6, 4
        )

        self.conv3 = OperatorBlock_3D(
            8 * factor * self.width,
            16 * factor * self.width,
            16,
            16,
            16,
            6,
            6,
            7,
            Normalize=True,
        )

        self.conv6 = OperatorBlock_3D(
            16 * factor * self.width, 4 * factor * self.width, 32, 32, 24, 6, 6, 7
        )

        self.conv7 = OperatorBlock_3D(
            8 * factor * self.width,
            2 * factor * self.width,
            48,
            48,
            32,
            14,
            14,
            10,
            Normalize=True,
        )

        self.conv8 = OperatorBlock_3D(
            4 * factor * self.width, 2 * self.width, 64, 64, 40, 20, 20, 14
        )  # will be reshaped

        self.fc1 = nn.Linear(3 * self.width, 4 * self.width)
        self.fc2 = nn.Linear(4 * self.width, 1)

    def forward(self, x):
        grid = self.get_grid(x.shape, x.device)
        x = torch.cat((x, grid), dim=-1)
        x_fc_1 = self.fc(x)
        x_fc_1 = F.gelu(x_fc_1)
        x_fc0 = self.fc0(x_fc_1)

        x_fc0 = F.gelu(x_fc0)

        x_fc0 = x_fc0.permute(0, 4, 1, 2, 3)
        self.padding = int(self.pad * 0.1 * x_fc0.shape[-1])
        if self.pad_both:
            x_fc0 = F.pad(
                x_fc0, [self.padding, self.padding, 0, 0, 0, 0], mode="constant"
            )
        else:
            x_fc0 = F.pad(x_fc0, [0, self.padding, 0, 0, 0, 0], mode="constant")

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

        x_c0 = self.conv0(x_fc0, int(3 * D1 / 4), int(3 * D2 / 4), D3)
        x_c1 = self.conv1(x_c0, D1 // 2, D2 // 2, D3)
        x_c2 = self.conv2(x_c1, D1 // 4, D2 // 4, int(D3 * 1.6))

        x_c3 = self.conv3(x_c2, D1 // 8, D2 // 8, int(D3 * 1.6))

        x_c6 = self.conv6(x_c3, D1 // 2, D2 // 2, int(D3 * 2.4))
        x_c6 = torch.cat(
            [
                x_c6,
                torch.nn.functional.interpolate(
                    x_c1,
                    size=(x_c6.shape[2], x_c6.shape[3], x_c6.shape[4]),
                    mode="trilinear",
                    align_corners=True,
                ),
            ],
            dim=1,
        )

        x_c7 = self.conv7(x_c6, int(3 * D1 / 4), int(3 * D2 / 4), int(3.2 * D3))
        x_c7 = torch.cat(
            [
                x_c7,
                torch.nn.functional.interpolate(
                    x_c0,
                    size=(x_c7.shape[2], x_c7.shape[3], x_c7.shape[4]),
                    mode="trilinear",
                    align_corners=True,
                ),
            ],
            dim=1,
        )

        x_c8 = self.conv8(x_c7, D1, D2, 4 * D3)

        x_c8 = torch.cat(
            [
                x_c8,
                torch.nn.functional.interpolate(
                    x_fc0,
                    size=(x_c8.shape[2], x_c8.shape[3], x_c8.shape[4]),
                    mode="trilinear",
                    align_corners=True,
                ),
            ],
            dim=1,
        )

        if self.padding != 0:
            if self.pad_both:
                x_c8 = x_c8[..., 4 * self.padding : -4 * self.padding]
            else:
                x_c8 = x_c8[..., : -4 * self.padding]

        x_c8 = x_c8.permute(0, 2, 3, 4, 1)

        x_fc1 = self.fc1(x_c8)

        x_fc1 = F.gelu(x_fc1)
        x_out = self.fc2(x_fc1)

        return x_out

    def get_grid(self, shape, device):
        batchsize, size_x, size_y, size_z = shape[0], shape[1], shape[2], shape[3]
        gridx = torch.tensor(np.linspace(0, 2 * np.pi, size_x), dtype=torch.float)
        gridx = gridx.reshape(1, size_x, 1, 1, 1).repeat(
            [batchsize, 1, size_y, size_z, 1]
        )
        gridy = torch.tensor(np.linspace(0, 2 * np.pi, size_y), dtype=torch.float)
        gridy = gridy.reshape(1, 1, size_y, 1, 1).repeat(
            [batchsize, size_x, 1, size_z, 1]
        )
        gridz = torch.tensor(np.linspace(0, 1, size_z), dtype=torch.float)
        gridz = gridz.reshape(1, 1, 1, size_z, 1).repeat(
            [batchsize, size_x, size_y, 1, 1]
        )
        return torch.cat(
            (
                torch.sin(gridx),
                torch.sin(gridy),
                torch.cos(gridx),
                torch.cos(gridy),
                gridz,
            ),
            dim=-1,
        ).to(device)


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

    input: the solution of the first 10 timesteps (u(1), ..., u(10)).
    input shape: (batchsize, x=S, y=S, t=T, c=1)
    output: the solution of the next 20 timesteps
    output shape: (batchsize, x=S, y=S, t=2*T, c=1)

    S,S,T = grid size along x,y and t (input function)
    S,S,2*T = grid size along x,y and t (output function)

    in_width = 4; [a(x,y,x),x,y,z]
    with = projection dimesion
    pad = padding amount
    pad_both = boolean, if true pad both size of the domain
    factor = scaling factor of the co-domain dimesions
    """

    def __init__(self, in_width, width, pad=2, factor=1, pad_both=False):
        super(Uno3D_T20, self).__init__()

        self.in_width = in_width  # input channel
        self.width = width

        self.pad = pad
        self.pad_both = pad_both

        self.fc = nn.Linear(self.in_width, self.in_width * 2)

        self.fc0 = nn.Linear(self.in_width * 2, self.width)

        self.conv0 = OperatorBlock_3D(
            self.width, 2 * factor * self.width, 48, 48, 10, 22, 22, 5, Normalize=True
        )

        self.conv1 = OperatorBlock_3D(
            2 * factor * self.width, 4 * factor * self.width, 32, 32, 10, 14, 14, 5
        )

        self.conv2 = OperatorBlock_3D(
            4 * factor * self.width, 8 * factor * self.width, 16, 16, 12, 6, 6, 5
        )

        self.conv3 = OperatorBlock_3D(
            8 * factor * self.width,
            16 * factor * self.width,
            16,
            16,
            12,
            6,
            6,
            6,
            Normalize=True,
        )

        self.conv6 = OperatorBlock_3D(
            16 * factor * self.width, 4 * factor * self.width, 32, 32, 18, 6, 6, 6
        )

        self.conv7 = OperatorBlock_3D(
            8 * factor * self.width,
            2 * factor * self.width,
            48,
            48,
            20,
            14,
            14,
            8,
            Normalize=True,
        )

        self.conv8 = OperatorBlock_3D(
            4 * factor * self.width, 2 * self.width, 64, 64, 20, 22, 22, 8
        )  # will be reshaped

        self.fc1 = nn.Linear(3 * self.width, 4 * self.width)
        self.fc2 = nn.Linear(4 * self.width, 1)

    def forward(self, x):
        grid = self.get_grid(x.shape, x.device)
        x = torch.cat((x, grid), dim=-1)
        x_fc_1 = self.fc(x)
        x_fc_1 = F.gelu(x_fc_1)
        x_fc0 = self.fc0(x_fc_1)

        x_fc0 = F.gelu(x_fc0)

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

        self.padding = int(self.pad * 0.1 * x_fc0.shape[-1])
        if self.pad_both:
            x_fc0 = F.pad(
                x_fc0, [self.padding, self.padding, 0, 0, 0, 0], mode="constant"
            )
        else:
            x_fc0 = F.pad(x_fc0, [0, self.padding, 0, 0, 0, 0], mode="constant")

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

        x_c0 = self.conv0(x_fc0, int(3 * D1 / 4), int(3 * D2 / 4), D3)
        x_c1 = self.conv1(x_c0, D1 // 2, D2 // 2, D3)
        x_c2 = self.conv2(x_c1, D1 // 4, D2 // 4, int(D3 * 1.2))

        x_c3 = self.conv3(x_c2, D1 // 4, D2 // 4, int(D3 * 1.2))

        x_c6 = self.conv6(x_c3, D1 // 2, D2 // 2, int(D3 * 1.8))
        x_c6 = torch.cat(
            [
                x_c6,
                torch.nn.functional.interpolate(
                    x_c1,
                    size=(x_c6.shape[2], x_c6.shape[3], x_c6.shape[4]),
                    mode="trilinear",
                    align_corners=True,
                ),
            ],
            dim=1,
        )

        x_c7 = self.conv7(x_c6, int(3 * D1 / 4), int(3 * D2 / 4), int(2.0 * D3))
        x_c7 = torch.cat(
            [
                x_c7,
                torch.nn.functional.interpolate(
                    x_c0,
                    size=(x_c7.shape[2], x_c7.shape[3], x_c7.shape[4]),
                    mode="trilinear",
                    align_corners=True,
                ),
            ],
            dim=1,
        )

        x_c8 = self.conv8(x_c7, D1, D2, 2 * D3)

        x_c8 = torch.cat(
            [
                x_c8,
                torch.nn.functional.interpolate(
                    x_fc0,
                    size=(x_c8.shape[2], x_c8.shape[3], x_c8.shape[4]),
                    mode="trilinear",
                    align_corners=True,
                ),
            ],
            dim=1,
        )

        if self.padding != 0:
            if self.pad_both:
                x_c8 = x_c8[..., 2 * self.padding : -2 * self.padding]
            else:
                x_c8 = x_c8[..., : -2 * self.padding]

        x_c8 = x_c8.permute(0, 2, 3, 4, 1)

        x_fc1 = self.fc1(x_c8)

        x_fc1 = F.gelu(x_fc1)
        x_out = self.fc2(x_fc1)

        return x_out

    def get_grid(self, shape, device):
        batchsize, size_x, size_y, size_z = shape[0], shape[1], shape[2], shape[3]
        gridx = torch.tensor(np.linspace(0, 2 * np.pi, size_x), dtype=torch.float)
        gridx = gridx.reshape(1, size_x, 1, 1, 1).repeat(
            [batchsize, 1, size_y, size_z, 1]
        )
        gridy = torch.tensor(np.linspace(0, 2 * np.pi, size_y), dtype=torch.float)
        gridy = gridy.reshape(1, 1, size_y, 1, 1).repeat(
            [batchsize, size_x, 1, size_z, 1]
        )
        gridz = torch.tensor(np.linspace(0, 1, size_z), dtype=torch.float)
        gridz = gridz.reshape(1, 1, 1, size_z, 1).repeat(
            [batchsize, size_x, size_y, 1, 1]
        )
        return torch.cat(
            (
                torch.sin(gridx),
                torch.sin(gridy),
                torch.cos(gridx),
                torch.cos(gridy),
                gridz,
            ),
            dim=-1,
        ).to(device)


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

    input: the solution of the first 10 timesteps (u(1), ..., u(10)).
    input shape: (batchsize, x=S, y=S, t=T, c=1)
    output: the solution of the next 10 timesteps
    output shape: (batchsize, x=S, y=S, t=T, c=1)

    S,S,T = grid size along x,y and t (input function)
    S,S,T = grid size along x,y and t (output function)

    in_width = 4; [a(x,y,x),x,y,z]
    with = projection dimesion
    pad = padding amount
    pad_both = boolean, if true pad both size of the domain
    factor = scaling factor of the co-domain dimesions
    """

    def __init__(self, in_width, width, pad=2, factor=1, pad_both=False):
        super(Uno3D_T10, self).__init__()

        self.in_width = in_width  # input channel
        self.width = width

        self.pad = pad
        self.pad_both = pad_both

        self.fc = nn.Linear(self.in_width, self.in_width * 2)

        self.fc0 = nn.Linear(self.in_width * 2, self.width)

        self.conv0 = OperatorBlock_3D(
            self.width, 2 * factor * self.width, 48, 48, 10, 22, 22, 5, Normalize=True
        )

        self.conv1 = OperatorBlock_3D(
            2 * factor * self.width, 4 * factor * self.width, 32, 32, 10, 14, 14, 5
        )

        self.conv2 = OperatorBlock_3D(
            4 * factor * self.width, 8 * factor * self.width, 16, 16, 10, 6, 6, 5
        )

        self.conv3 = OperatorBlock_3D(
            8 * factor * self.width,
            16 * factor * self.width,
            16,
            16,
            10,
            6,
            6,
            5,
            Normalize=True,
        )

        self.conv6 = OperatorBlock_3D(
            16 * factor * self.width, 4 * factor * self.width, 32, 32, 10, 6, 6, 5
        )

        self.conv7 = OperatorBlock_3D(
            8 * factor * self.width,
            2 * factor * self.width,
            48,
            48,
            10,
            14,
            14,
            5,
            Normalize=True,
        )

        self.conv8 = OperatorBlock_3D(
            4 * factor * self.width, 2 * self.width, 64, 64, 10, 22, 22, 5
        )  # will be reshaped

        self.fc1 = nn.Linear(3 * self.width, 4 * self.width)
        self.fc2 = nn.Linear(4 * self.width, 1)

    def forward(self, x):
        grid = self.get_grid(x.shape, x.device)
        x = torch.cat((x, grid), dim=-1)
        x_fc_1 = self.fc(x)
        x_fc_1 = F.gelu(x_fc_1)
        x_fc0 = self.fc0(x_fc_1)

        x_fc0 = F.gelu(x_fc0)

        x_fc0 = x_fc0.permute(0, 4, 1, 2, 3)
        self.padding = int(self.pad * 0.1 * x_fc0.shape[-1])
        if self.pad_both:
            x_fc0 = F.pad(
                x_fc0, [self.padding, self.padding, 0, 0, 0, 0], mode="constant"
            )
        else:
            x_fc0 = F.pad(x_fc0, [0, self.padding, 0, 0, 0, 0], mode="constant")

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

        x_c0 = self.conv0(x_fc0, int(3 * D1 / 4), int(3 * D2 / 4), D3)
        x_c1 = self.conv1(x_c0, D1 // 2, D2 // 2, D3)
        x_c2 = self.conv2(x_c1, D1 // 4, D2 // 4, int(1.0 * D3))

        x_c3 = self.conv3(x_c2, D1 // 4, D2 // 4, int(1.0 * D3))

        x_c6 = self.conv6(x_c3, D1 // 2, D2 // 2, int(1.0 * D3))
        x_c6 = torch.cat(
            [
                x_c6,
                torch.nn.functional.interpolate(
                    x_c1,
                    size=(x_c6.shape[2], x_c6.shape[3], x_c6.shape[4]),
                    mode="trilinear",
                    align_corners=True,
                ),
            ],
            dim=1,
        )

        x_c7 = self.conv7(x_c6, int(3 * D1 / 4), int(3 * D2 / 4), D3)
        x_c7 = torch.cat(
            [
                x_c7,
                torch.nn.functional.interpolate(
                    x_c0,
                    size=(x_c7.shape[2], x_c7.shape[3], x_c7.shape[4]),
                    mode="trilinear",
                    align_corners=True,
                ),
            ],
            dim=1,
        )

        x_c8 = self.conv8(x_c7, D1, D2, D3)

        x_c8 = torch.cat(
            [
                x_c8,
                torch.nn.functional.interpolate(
                    x_fc0,
                    size=(x_c8.shape[2], x_c8.shape[3], x_c8.shape[4]),
                    mode="trilinear",
                    align_corners=True,
                ),
            ],
            dim=1,
        )

        if self.padding != 0:
            if self.pad_both:
                x_c8 = x_c8[..., self.padding : -self.padding]
            else:
                x_c8 = x_c8[..., : -self.padding]

        x_c8 = x_c8.permute(0, 2, 3, 4, 1)

        x_fc1 = self.fc1(x_c8)

        x_fc1 = F.gelu(x_fc1)
        x_out = self.fc2(x_fc1)

        return x_out

    def get_grid(self, shape, device):
        batchsize, size_x, size_y, size_z = shape[0], shape[1], shape[2], shape[3]
        gridx = torch.tensor(np.linspace(0, 2 * np.pi, size_x), dtype=torch.float)
        gridx = gridx.reshape(1, size_x, 1, 1, 1).repeat(
            [batchsize, 1, size_y, size_z, 1]
        )
        gridy = torch.tensor(np.linspace(0, 2 * np.pi, size_y), dtype=torch.float)
        gridy = gridy.reshape(1, 1, size_y, 1, 1).repeat(
            [batchsize, size_x, 1, size_z, 1]
        )
        gridz = torch.tensor(np.linspace(0, 1, size_z), dtype=torch.float)
        gridz = gridz.reshape(1, 1, 1, size_z, 1).repeat(
            [batchsize, size_x, size_y, 1, 1]
        )
        return torch.cat(
            (
                torch.sin(gridx),
                torch.sin(gridy),
                torch.cos(gridx),
                torch.cos(gridy),
                gridz,
            ),
            dim=-1,
        ).to(device)


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

    input: the solution of the first 6 timesteps (u(1), ..., u(10)).
    input shape: (batchsize, x=S, y=S, t=T, c=1)
    output: the solution of the next 9 timesteps
    output shape: (batchsize, x=S, y=S, t=int(1.5*T), c=1)

    S,S,T = grid size along x,y and t (input function)
    S,S,int(1.5*T) = grid size along x,y and t (output function)

    in_width = 4; [a(x,y,x),x,y,z]
    with = projection dimesion
    pad = padding amount
    pad_both = boolean, if true pad both size of the domain
    factor = scaling factor of the co-domain dimesions
    """

    def __init__(self, in_width, width, pad=2, factor=1, pad_both=False):
        super(Uno3D_T9, self).__init__()

        self.in_width = in_width  # input channel
        self.width = width

        self.pad = pad
        self.pad_both = pad_both

        self.fc = nn.Linear(self.in_width, self.in_width * 2)

        self.fc0 = nn.Linear(self.in_width * 2, self.width)

        self.conv0 = OperatorBlock_3D(
            self.width, 2 * factor * self.width, 48, 48, 6, 20, 20, 3, Normalize=True
        )

        self.conv1 = OperatorBlock_3D(
            2 * factor * self.width, 4 * factor * self.width, 32, 32, 6, 18, 18, 3
        )

        self.conv2 = OperatorBlock_3D(
            4 * factor * self.width, 8 * factor * self.width, 16, 16, 6, 6, 6, 3
        )

        self.conv3 = OperatorBlock_3D(
            8 * factor * self.width,
            16 * factor * self.width,
            16,
            16,
            8,
            6,
            6,
            3,
            Normalize=True,
        )

        self.conv6 = OperatorBlock_3D(
            16 * factor * self.width, 4 * factor * self.width, 32, 32, 8, 6, 6, 3
        )

        self.conv7 = OperatorBlock_3D(
            8 * factor * self.width,
            2 * factor * self.width,
            48,
            48,
            9,
            14,
            14,
            3,
            Normalize=True,
        )

        self.conv8 = OperatorBlock_3D(
            4 * factor * self.width, 2 * self.width, 64, 64, 9, 20, 20, 4
        )  # will be reshaped

        self.fc1 = nn.Linear(3 * self.width, 4 * self.width)
        self.fc2 = nn.Linear(4 * self.width, 1)

    def forward(self, x):
        grid = self.get_grid(x.shape, x.device)
        x = torch.cat((x, grid), dim=-1)
        x_fc_1 = self.fc(x)
        x_fc_1 = F.gelu(x_fc_1)
        x_fc0 = self.fc0(x_fc_1)

        x_fc0 = F.gelu(x_fc0)

        x_fc0 = x_fc0.permute(0, 4, 1, 2, 3)
        self.padding = int(self.pad * 0.1 * x_fc0.shape[-1])
        if self.pad_both:
            x_fc0 = F.pad(
                x_fc0, [self.padding, self.padding, 0, 0, 0, 0], mode="constant"
            )
        else:
            x_fc0 = F.pad(x_fc0, [0, self.padding, 0, 0, 0, 0], mode="constant")

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

        x_c0 = self.conv0(x_fc0, int(3 * D1 / 4), int(3 * D2 / 4), D3)
        x_c1 = self.conv1(x_c0, D1 // 2, D2 // 2, D3)
        x_c2 = self.conv2(x_c1, D1 // 4, D2 // 4, D3)

        x_c3 = self.conv3(x_c2, D1 // 4, D2 // 4, int(8 * D3 / 6))

        x_c6 = self.conv6(x_c3, D1 // 2, D2 // 2, int(8 * D3 / 6))
        x_c6 = torch.cat(
            [
                x_c6,
                torch.nn.functional.interpolate(
                    x_c1,
                    size=(x_c6.shape[2], x_c6.shape[3], x_c6.shape[4]),
                    mode="trilinear",
                    align_corners=True,
                ),
            ],
            dim=1,
        )

        x_c7 = self.conv7(x_c6, int(3 * D1 / 4), int(3 * D2 / 4), int(9 * D3 / 6))
        x_c7 = torch.cat(
            [
                x_c7,
                torch.nn.functional.interpolate(
                    x_c0,
                    size=(x_c7.shape[2], x_c7.shape[3], x_c7.shape[4]),
                    mode="trilinear",
                    align_corners=True,
                ),
            ],
            dim=1,
        )

        x_c8 = self.conv8(x_c7, D1, D2, int(9 * D3 / 6))

        x_c8 = torch.cat(
            [
                x_c8,
                torch.nn.functional.interpolate(
                    x_fc0,
                    size=(x_c8.shape[2], x_c8.shape[3], x_c8.shape[4]),
                    mode="trilinear",
                    align_corners=True,
                ),
            ],
            dim=1,
        )

        if self.padding != 0:
            if self.pad_both:
                x_c8 = x_c8[..., int(9 * self.padding / 6) : -int(9 * self.padding / 6)]
            else:
                x_c8 = x_c8[..., : -int(9 * self.padding / 6)]

        x_c8 = x_c8.permute(0, 2, 3, 4, 1)

        x_fc1 = self.fc1(x_c8)

        x_fc1 = F.gelu(x_fc1)
        x_out = self.fc2(x_fc1)

        return x_out

    def get_grid(self, shape, device):
        batchsize, size_x, size_y, size_z = shape[0], shape[1], shape[2], shape[3]
        gridx = torch.tensor(np.linspace(0, 2 * np.pi, size_x), dtype=torch.float)
        gridx = gridx.reshape(1, size_x, 1, 1, 1).repeat(
            [batchsize, 1, size_y, size_z, 1]
        )
        gridy = torch.tensor(np.linspace(0, 2 * np.pi, size_y), dtype=torch.float)
        gridy = gridy.reshape(1, 1, size_y, 1, 1).repeat(
            [batchsize, size_x, 1, size_z, 1]
        )
        gridz = torch.tensor(np.linspace(0, 1, size_z), dtype=torch.float)
        gridz = gridz.reshape(1, 1, 1, size_z, 1).repeat(
            [batchsize, size_x, size_y, 1, 1]
        )
        return torch.cat(
            (
                torch.sin(gridx),
                torch.sin(gridy),
                torch.cos(gridx),
                torch.cos(gridy),
                gridz,
            ),
            dim=-1,
        ).to(device)


######
# Codes for Higher resolution simulation of Navier Stockes equations. Functionality of each
# of the following classes are as the above classes.
######


class Uno3D_T40_256(nn.Module):
    def __init__(self, in_width, width, pad=1, factor=1):
        super(Uno3D_T40_256, self).__init__()

        self.in_width = in_width  # input channel
        self.width = width

        self.pad = pad

        self.fc_n1 = nn.Linear(self.in_width, self.width // 2)

        self.fc0 = nn.Linear(self.width // 2, self.width)

        self.conv0 = OperatorBlock_3D(
            self.width, 2 * factor * self.width, 64, 64, 10, 32, 32, 5, Normalize=True
        )

        self.conv1 = OperatorBlock_3D(
            2 * factor * self.width, 4 * factor * self.width, 16, 16, 10, 8, 8, 5
        )

        self.conv2 = OperatorBlock_3D(
            4 * factor * self.width, 8 * factor * self.width, 8, 8, 16, 4, 4, 5
        )

        self.conv3 = OperatorBlock_3D(
            8 * factor * self.width,
            16 * factor * self.width,
            8,
            8,
            16,
            4,
            4,
            8,
            Normalize=True,
        )

        self.conv4 = OperatorBlock_3D(
            16 * factor * self.width, 16 * factor * self.width, 8, 8, 16, 4, 4, 8
        )

        self.conv5 = OperatorBlock_3D(
            16 * factor * self.width,
            8 * factor * self.width,
            8,
            8,
            16,
            4,
            4,
            8,
            Normalize=True,
        )

        self.conv6 = OperatorBlock_3D(
            8 * factor * self.width, 4 * factor * self.width, 16, 16, 24, 4, 4, 8
        )

        self.conv7 = OperatorBlock_3D(
            8 * factor * self.width,
            2 * factor * self.width,
            64,
            64,
            32,
            8,
            8,
            12,
            Normalize=True,
        )

        self.conv8 = OperatorBlock_3D(
            4 * factor * self.width, 2 * self.width, 256, 256, 40, 32, 32, 16
        )  # will be reshaped

        self.fc1 = nn.Linear(3 * self.width, 4 * self.width)
        self.fc2 = nn.Linear(4 * self.width, 1)

    def forward(self, x):
        grid = self.get_grid(x.shape, x.device)
        x = torch.cat((x, grid), dim=-1)
        x_fc_1 = self.fc(x)
        x_fc_1 = F.gelu(x_fc_1)
        x_fc0 = self.fc0(x_fc_1)

        x_fc0 = F.gelu(x_fc0)

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

        self.padding = int(self.pad * 0.1 * x_fc0.shape[-1])
        if self.pad_both:
            x_fc0 = F.pad(
                x_fc0, [self.padding, self.padding, 0, 0, 0, 0], mode="constant"
            )
        else:
            x_fc0 = F.pad(x_fc0, [0, self.padding, 0, 0, 0, 0], mode="constant")

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

        x_c0 = self.conv0(x_fc0, D1 // 4, D2 // 4, D3)
        x_c1 = self.conv1(x_c0, D1 // 16, D2 // 16, D3)
        x_c2 = self.conv2(x_c1, D1 // 32, D2 // 32, int(D3 * 1.6))

        x_c3 = self.conv3(x_c2, D1 // 32, D2 // 32, int(D3 * 1.6))
        x_c4 = self.conv4(x_c3, D1 // 32, D2 // 32, int(D3 * 1.6))
        x_c5 = self.conv5(x_c4, D1 // 32, D2 // 32, int(D3 * 1.6))

        x_c6 = self.conv6(x_c5, D1 // 16, D2 // 16, int(D3 * 2.4))
        x_c6 = torch.cat(
            [
                x_c6,
                torch.nn.functional.interpolate(
                    x_c1,
                    size=(x_c6.shape[2], x_c6.shape[3], x_c6.shape[4]),
                    mode="trilinear",
                    align_corners=True,
                ),
            ],
            dim=1,
        )

        x_c7 = self.conv7(x_c6, D1 // 4, D2 // 4, int(3.2 * D3))
        x_c7 = torch.cat(
            [
                x_c7,
                torch.nn.functional.interpolate(
                    x_c0,
                    size=(x_c7.shape[2], x_c7.shape[3], x_c7.shape[4]),
                    mode="trilinear",
                    align_corners=True,
                ),
            ],
            dim=1,
        )

        x_c8 = self.conv8(x_c7, D1, D2, 4 * D3)

        x_c8 = torch.cat(
            [
                x_c8,
                torch.nn.functional.interpolate(
                    x_fc0,
                    size=(x_c8.shape[2], x_c8.shape[3], x_c8.shape[4]),
                    mode="trilinear",
                    align_corners=True,
                ),
            ],
            dim=1,
        )

        if self.padding != 0:
            if self.pad_both:
                x_c8 = x_c8[..., 4 * self.padding : -4 * self.padding]
            else:
                x_c8 = x_c8[..., : -4 * self.padding]

        x_c8 = x_c8.permute(0, 2, 3, 4, 1)

        x_fc1 = self.fc1(x_c8)

        x_fc1 = F.gelu(x_fc1)
        x_out = self.fc2(x_fc1)

        return x_out

    def get_grid(self, shape, device):
        batchsize, size_x, size_y, size_z = shape[0], shape[1], shape[2], shape[3]
        gridx = torch.tensor(np.linspace(0, 2 * np.pi, size_x), dtype=torch.float)
        gridx = gridx.reshape(1, size_x, 1, 1, 1).repeat(
            [batchsize, 1, size_y, size_z, 1]
        )
        gridy = torch.tensor(np.linspace(0, 2 * np.pi, size_y), dtype=torch.float)
        gridy = gridy.reshape(1, 1, size_y, 1, 1).repeat(
            [batchsize, size_x, 1, size_z, 1]
        )
        gridz = torch.tensor(np.linspace(0, 1, size_z), dtype=torch.float)
        gridz = gridz.reshape(1, 1, 1, size_z, 1).repeat(
            [batchsize, size_x, size_y, 1, 1]
        )
        return torch.cat(
            (
                torch.sin(gridx),
                torch.sin(gridy),
                torch.cos(gridx),
                torch.cos(gridy),
                gridz,
            ),
            dim=-1,
        ).to(device)


class Uno3D_T20_256(nn.Module):
    def __init__(self, in_width, width, pad=2, factor=1, pad_both=False):
        super(Uno3D_T20_256, self).__init__()

        self.in_width = in_width  # input channel
        self.width = width

        self.pad = pad
        self.pad_both = pad_both

        self.fc = nn.Linear(self.in_width, self.width // 2)

        self.fc0 = nn.Linear(self.width // 2, self.width)

        self.conv0 = OperatorBlock_3D(
            self.width, 2 * factor * self.width, 64, 64, 10, 32, 32, 5, Normalize=True
        )

        self.conv1 = OperatorBlock_3D(
            2 * factor * self.width, 4 * factor * self.width, 16, 16, 10, 8, 8, 5
        )

        self.conv2 = OperatorBlock_3D(
            4 * factor * self.width, 8 * factor * self.width, 8, 8, 12, 4, 4, 5
        )

        self.conv3 = OperatorBlock_3D(
            8 * factor * self.width,
            16 * factor * self.width,
            8,
            8,
            12,
            4,
            4,
            6,
            Normalize=True,
        )

        self.conv4 = OperatorBlock_3D(
            16 * factor * self.width, 16 * factor * self.width, 8, 8, 16, 4, 4, 6
        )

        self.conv5 = OperatorBlock_3D(
            16 * factor * self.width,
            8 * factor * self.width,
            8,
            8,
            16,
            4,
            4,
            8,
            Normalize=True,
        )

        self.conv6 = OperatorBlock_3D(
            8 * factor * self.width, 4 * factor * self.width, 16, 16, 18, 4, 4, 8
        )

        self.conv7 = OperatorBlock_3D(
            8 * factor * self.width,
            2 * factor * self.width,
            64,
            64,
            20,
            8,
            8,
            8,
            Normalize=True,
        )

        self.conv8 = OperatorBlock_3D(
            4 * factor * self.width, 2 * self.width, 256, 256, 20, 32, 32, 8
        )  # will be reshaped

        self.fc1 = nn.Linear(3 * self.width, 4 * self.width)
        self.fc2 = nn.Linear(4 * self.width, 1)

    def forward(self, x):
        grid = self.get_grid(x.shape, x.device)
        x = torch.cat((x, grid), dim=-1)
        x_fc_1 = self.fc(x)
        x_fc_1 = F.gelu(x_fc_1)
        x_fc0 = self.fc0(x_fc_1)

        x_fc0 = F.gelu(x_fc0)

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

        self.padding = int(self.pad * 0.1 * x_fc0.shape[-1])

        if self.pad_both:
            x_fc0 = F.pad(
                x_fc0, [self.padding, self.padding, 0, 0, 0, 0], mode="constant"
            )
        else:
            x_fc0 = F.pad(x_fc0, [0, self.padding, 0, 0, 0, 0], mode="constant")

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

        x_c0 = self.conv0(x_fc0, D1 // 4, D2 // 4, D3)
        x_c1 = self.conv1(x_c0, D1 // 16, D2 // 16, D3)
        x_c2 = self.conv2(x_c1, D1 // 32, D2 // 32, int(D3 * 1.2))

        x_c3 = self.conv3(x_c2, D1 // 32, D2 // 32, int(D3 * 1.2))
        x_c4 = self.conv4(x_c3, D1 // 32, D2 // 32, int(D3 * 1.6))
        x_c5 = self.conv5(x_c4, D1 // 32, D2 // 32, int(D3 * 1.6))

        x_c6 = self.conv6(x_c5, D1 // 16, D2 // 16, int(D3 * 1.8))
        x_c6 = torch.cat(
            [
                x_c6,
                torch.nn.functional.interpolate(
                    x_c1,
                    size=(x_c6.shape[2], x_c6.shape[3], x_c6.shape[4]),
                    mode="trilinear",
                    align_corners=True,
                ),
            ],
            dim=1,
        )

        x_c7 = self.conv7(x_c6, D1 // 4, D2 // 4, int(2.0 * D3))
        x_c7 = torch.cat(
            [
                x_c7,
                torch.nn.functional.interpolate(
                    x_c0,
                    size=(x_c7.shape[2], x_c7.shape[3], x_c7.shape[4]),
                    mode="trilinear",
                    align_corners=True,
                ),
            ],
            dim=1,
        )

        x_c8 = self.conv8(x_c7, D1, D2, 2 * D3)

        x_c8 = torch.cat(
            [
                x_c8,
                torch.nn.functional.interpolate(
                    x_fc0,
                    size=(x_c8.shape[2], x_c8.shape[3], x_c8.shape[4]),
                    mode="trilinear",
                    align_corners=True,
                ),
            ],
            dim=1,
        )

        if self.padding != 0:
            if self.pad_both:
                x_c8 = x_c8[..., 2 * self.padding : -2 * self.padding]
            else:
                x_c8 = x_c8[..., : -2 * self.padding]

        x_c8 = x_c8.permute(0, 2, 3, 4, 1)

        x_fc1 = self.fc1(x_c8)

        x_fc1 = F.gelu(x_fc1)
        x_out = self.fc2(x_fc1)

        return x_out

    def get_grid(self, shape, device):
        batchsize, size_x, size_y, size_z = shape[0], shape[1], shape[2], shape[3]
        gridx = torch.tensor(np.linspace(0, 2 * np.pi, size_x), dtype=torch.float)
        gridx = gridx.reshape(1, size_x, 1, 1, 1).repeat(
            [batchsize, 1, size_y, size_z, 1]
        )
        gridy = torch.tensor(np.linspace(0, 2 * np.pi, size_y), dtype=torch.float)
        gridy = gridy.reshape(1, 1, size_y, 1, 1).repeat(
            [batchsize, size_x, 1, size_z, 1]
        )
        gridz = torch.tensor(np.linspace(0, 1, size_z), dtype=torch.float)
        gridz = gridz.reshape(1, 1, 1, size_z, 1).repeat(
            [batchsize, size_x, size_y, 1, 1]
        )
        return torch.cat(
            (
                torch.sin(gridx),
                torch.sin(gridy),
                torch.cos(gridx),
                torch.cos(gridy),
                gridz,
            ),
            dim=-1,
        ).to(device)


class Uno3D_T10_256(nn.Module):
    def __init__(self, in_width, width, pad=2, factor=1, pad_both=False):
        super(Uno3D_T10_256, self).__init__()

        self.in_width = in_width  # input channel
        self.width = width

        self.pad = pad
        self.pad_both = pad_both

        self.fc = nn.Linear(self.in_width, self.width // 2)

        self.fc0 = nn.Linear(self.width // 2, self.width)

        self.conv0 = OperatorBlock_3D(
            self.width, 2 * factor * self.width, 64, 64, 10, 32, 32, 5, Normalize=True
        )

        self.conv1 = OperatorBlock_3D(
            2 * factor * self.width, 4 * factor * self.width, 16, 16, 10, 8, 8, 4
        )

        self.conv2 = OperatorBlock_3D(
            4 * factor * self.width, 8 * factor * self.width, 8, 8, 8, 4, 4, 4
        )

        self.conv3 = OperatorBlock_3D(
            8 * factor * self.width,
            16 * factor * self.width,
            8,
            8,
            8,
            4,
            4,
            4,
            Normalize=True,
        )

        self.conv4 = OperatorBlock_3D(
            16 * factor * self.width, 16 * factor * self.width, 8, 8, 8, 4, 4, 4
        )

        self.conv5 = OperatorBlock_3D(
            16 * factor * self.width,
            8 * factor * self.width,
            8,
            8,
            8,
            4,
            4,
            4,
            Normalize=True,
        )

        self.conv6 = OperatorBlock_3D(
            8 * factor * self.width, 4 * factor * self.width, 16, 16, 8, 4, 4, 4
        )

        self.conv7 = OperatorBlock_3D(
            8 * factor * self.width,
            2 * factor * self.width,
            64,
            64,
            10,
            8,
            8,
            4,
            Normalize=True,
        )

        self.conv8 = OperatorBlock_3D(
            4 * factor * self.width, 2 * self.width, 256, 256, 10, 32, 32, 5
        )  # will be reshaped

        self.fc1 = nn.Linear(3 * self.width, 4 * self.width)
        self.fc2 = nn.Linear(4 * self.width, 1)

    def forward(self, x):
        grid = self.get_grid(x.shape, x.device)
        x = torch.cat((x, grid), dim=-1)
        x_fc_1 = self.fc(x)
        x_fc_1 = F.gelu(x_fc_1)
        x_fc0 = self.fc0(x_fc_1)

        x_fc0 = F.gelu(x_fc0)

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

        self.padding = int(self.pad * 0.1 * x_fc0.shape[-1])

        if self.pad_both:
            x_fc0 = F.pad(
                x_fc0, [self.padding, self.padding, 0, 0, 0, 0], mode="constant"
            )
        else:
            x_fc0 = F.pad(x_fc0, [0, self.padding, 0, 0, 0, 0], mode="constant")

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

        x_c0 = self.conv0(x_fc0, D1 // 4, D2 // 4, D3)
        x_c1 = self.conv1(x_c0, D1 // 16, D2 // 16, D3)
        x_c2 = self.conv2(x_c1, D1 // 32, D2 // 32, int(0.8 * D3))

        x_c3 = self.conv3(x_c2, D1 // 32, D2 // 32, int(0.8 * D3))
        x_c4 = self.conv4(x_c3, D1 // 32, D2 // 32, int(0.8 * D3))
        x_c5 = self.conv5(x_c4, D1 // 32, D2 // 32, int(0.8 * D3))

        x_c6 = self.conv6(x_c5, D1 // 16, D2 // 16, int(0.8 * D3))
        x_c6 = torch.cat(
            [
                x_c6,
                torch.nn.functional.interpolate(
                    x_c1,
                    size=(x_c6.shape[2], x_c6.shape[3], x_c6.shape[4]),
                    mode="trilinear",
                    align_corners=True,
                ),
            ],
            dim=1,
        )

        x_c7 = self.conv7(x_c6, D1 // 4, D2 // 4, D3)
        x_c7 = torch.cat(
            [
                x_c7,
                torch.nn.functional.interpolate(
                    x_c0,
                    size=(x_c7.shape[2], x_c7.shape[3], x_c7.shape[4]),
                    mode="trilinear",
                    align_corners=True,
                ),
            ],
            dim=1,
        )

        x_c8 = self.conv8(x_c7, D1, D2, D3)

        x_c8 = torch.cat(
            [
                x_c8,
                torch.nn.functional.interpolate(
                    x_fc0,
                    size=(x_c8.shape[2], x_c8.shape[3], x_c8.shape[4]),
                    mode="trilinear",
                    align_corners=True,
                ),
            ],
            dim=1,
        )

        if self.padding != 0:
            if self.pad_both:
                x_c8 = x_c8[..., self.padding : -self.padding]
            else:
                x_c8 = x_c8[..., : -self.padding]

        x_c8 = x_c8.permute(0, 2, 3, 4, 1)

        x_fc1 = self.fc1(x_c8)

        x_fc1 = F.gelu(x_fc1)
        x_out = self.fc2(x_fc1)

        return x_out

    def get_grid(self, shape, device):
        batchsize, size_x, size_y, size_z = shape[0], shape[1], shape[2], shape[3]
        gridx = torch.tensor(np.linspace(0, 2 * np.pi, size_x), dtype=torch.float)
        gridx = gridx.reshape(1, size_x, 1, 1, 1).repeat(
            [batchsize, 1, size_y, size_z, 1]
        )
        gridy = torch.tensor(np.linspace(0, 2 * np.pi, size_y), dtype=torch.float)
        gridy = gridy.reshape(1, 1, size_y, 1, 1).repeat(
            [batchsize, size_x, 1, size_z, 1]
        )
        gridz = torch.tensor(np.linspace(0, 1, size_z), dtype=torch.float)
        gridz = gridz.reshape(1, 1, 1, size_z, 1).repeat(
            [batchsize, size_x, size_y, 1, 1]
        )
        return torch.cat(
            (
                torch.sin(gridx),
                torch.sin(gridy),
                torch.cos(gridx),
                torch.cos(gridy),
                gridz,
            ),
            dim=-1,
        ).to(device)


class Uno3D_T9_256(nn.Module):
    def __init__(self, in_width, width, pad=2, factor=1, pad_both=False):
        super(Uno3D_T10_256, self).__init__()

        self.in_width = in_width  # input channel
        self.width = width

        self.pad = pad
        self.pad_both = pad_both

        self.fc = nn.Linear(self.in_width, self.width // 2)

        self.fc0 = nn.Linear(self.width // 2, self.width)

        self.conv0 = OperatorBlock_3D(
            self.width, 2 * factor * self.width, 64, 64, 6, 32, 32, 3, Normalize=True
        )

        self.conv1 = OperatorBlock_3D(
            2 * factor * self.width, 4 * factor * self.width, 16, 16, 6, 8, 8, 3
        )

        self.conv2 = OperatorBlock_3D(
            4 * factor * self.width, 8 * factor * self.width, 8, 8, 6, 4, 4, 3
        )

        self.conv3 = OperatorBlock_3D(
            8 * factor * self.width,
            16 * factor * self.width,
            8,
            8,
            8,
            4,
            4,
            3,
            Normalize=True,
        )

        self.conv4 = OperatorBlock_3D(
            16 * factor * self.width, 16 * factor * self.width, 8, 8, 8, 4, 4, 4
        )

        self.conv5 = OperatorBlock_3D(
            16 * factor * self.width,
            8 * factor * self.width,
            8,
            8,
            8,
            4,
            4,
            4,
            Normalize=True,
        )

        self.conv6 = OperatorBlock_3D(
            8 * factor * self.width, 4 * factor * self.width, 16, 16, 8, 4, 4, 4
        )

        self.conv7 = OperatorBlock_3D(
            8 * factor * self.width,
            2 * factor * self.width,
            64,
            64,
            9,
            4,
            4,
            4,
            Normalize=True,
        )

        self.conv8 = OperatorBlock_3D(
            4 * factor * self.width, 2 * self.width, 256, 256, 9, 32, 32, 4
        )  # will be reshaped

        self.fc1 = nn.Linear(3 * self.width, 4 * self.width)
        self.fc2 = nn.Linear(4 * self.width, 1)

    def forward(self, x):
        grid = self.get_grid(x.shape, x.device)
        x = torch.cat((x, grid), dim=-1)
        x_fc_1 = self.fc(x)
        x_fc_1 = F.gelu(x_fc_1)
        x_fc0 = self.fc0(x_fc_1)

        x_fc0 = F.gelu(x_fc0)

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

        self.padding = int(self.pad * 0.1 * x_fc0.shape[-1])

        if self.pad_both:
            x_fc0 = F.pad(
                x_fc0, [self.padding, self.padding, 0, 0, 0, 0], mode="constant"
            )
        else:
            x_fc0 = F.pad(x_fc0, [0, self.padding, 0, 0, 0, 0], mode="constant")

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

        x_c0 = self.conv0(x_fc0, D1 // 4, D2 // 4, D3)
        x_c1 = self.conv1(x_c0, D1 // 16, D2 // 16, D3)
        x_c2 = self.conv2(x_c1, D1 // 32, D2 // 32, D3)

        x_c3 = self.conv3(x_c2, D1 // 32, D2 // 32, int(8 * D3 / 6))
        x_c4 = self.conv4(x_c3, D1 // 32, D2 // 32, int(8 * D3 / 6))
        x_c5 = self.conv5(x_c4, D1 // 32, D2 // 32, int(8 * D3 / 6))

        x_c6 = self.conv6(x_c5, D1 // 16, D2 // 16, int(8 * D3 / 6))
        x_c6 = torch.cat(
            [
                x_c6,
                torch.nn.functional.interpolate(
                    x_c1,
                    size=(x_c6.shape[2], x_c6.shape[3], x_c6.shape[4]),
                    mode="trilinear",
                    align_corners=True,
                ),
            ],
            dim=1,
        )

        x_c7 = self.conv7(x_c6, D1 // 4, D2 // 4, int(9 * D3 / 6))
        x_c7 = torch.cat(
            [
                x_c7,
                torch.nn.functional.interpolate(
                    x_c0,
                    size=(x_c7.shape[2], x_c7.shape[3], x_c7.shape[4]),
                    mode="trilinear",
                    align_corners=True,
                ),
            ],
            dim=1,
        )

        x_c8 = self.conv8(x_c7, D1, D2, int(9 * D3 / 6))

        x_c8 = torch.cat(
            [
                x_c8,
                torch.nn.functional.interpolate(
                    x_fc0,
                    size=(x_c8.shape[2], x_c8.shape[3], x_c8.shape[4]),
                    mode="trilinear",
                    align_corners=True,
                ),
            ],
            dim=1,
        )

        if self.padding != 0:
            if self.pad_both:
                x_c8 = x_c8[..., int(9 * self.padding / 6) : -int(9 * self.padding / 6)]
            else:
                x_c8 = x_c8[..., : -int(9 * self.padding / 6)]

        x_c8 = x_c8.permute(0, 2, 3, 4, 1)

        x_fc1 = self.fc1(x_c8)

        x_fc1 = F.gelu(x_fc1)
        x_out = self.fc2(x_fc1)

        return x_out

    def get_grid(self, shape, device):
        batchsize, size_x, size_y, size_z = shape[0], shape[1], shape[2], shape[3]
        gridx = torch.tensor(np.linspace(0, 2 * np.pi, size_x), dtype=torch.float)
        gridx = gridx.reshape(1, size_x, 1, 1, 1).repeat(
            [batchsize, 1, size_y, size_z, 1]
        )
        gridy = torch.tensor(np.linspace(0, 2 * np.pi, size_y), dtype=torch.float)
        gridy = gridy.reshape(1, 1, size_y, 1, 1).repeat(
            [batchsize, size_x, 1, size_z, 1]
        )
        gridz = torch.tensor(np.linspace(0, 1, size_z), dtype=torch.float)
        gridz = gridz.reshape(1, 1, 1, size_z, 1).repeat(
            [batchsize, size_x, size_y, 1, 1]
        )
        return torch.cat(
            (
                torch.sin(gridx),
                torch.sin(gridy),
                torch.cos(gridx),
                torch.cos(gridy),
                gridz,
            ),
            dim=-1,
        ).to(device)