### Example Usage of Data Utilities

Source: https://context7.com/lu-group/deeponet-fno/llms.txt

Demonstrates loading initial conditions and solutions from a .mat file, normalizing the initial conditions, and calculating relative error using LpLoss.

```python
# Usage example
dataloader = MatReader('burgers_data_R10.mat')
x_data = dataloader.read_field('a')  # Initial conditions
y_data = dataloader.read_field('u')  # Solutions

normalizer = UnitGaussianNormalizer(x_data[:1000])
x_normalized = normalizer.encode(x_data[:1000])

loss_fn = LpLoss(size_average=False)
relative_error = loss_fn(predictions, targets)
```

--------------------------------

### Train 1D FNO for Burgers' Equation

Source: https://context7.com/lu-group/deeponet-fno/llms.txt

Example script for training the 1D FNO model on Burgers' equation data. It includes data loading, model initialization, optimizer setup, and the training loop.

```python
# Training example for Burgers' equation
train_data_res = 128
sub = 2**13 // train_data_res
train, ntest = 1000, 200

dataloader = MatReader('burgers_data_R10.mat')
x_data = dataloader.read_field('a')[:,::sub]
y_data = dataloader.read_field('u')[:,::sub]

x_train, y_train = x_data[:ntrain,:], y_data[:ntrain,:]
x_test, y_test = x_data[-ntest:,:], y_data[-ntest:,:]

# Add spatial grid as input channel
grid = torch.linspace(0, 1, train_data_res).reshape(train_data_res, 1)
x_train = torch.cat([x_train.reshape(ntrain, train_data_res, 1),
                     grid.repeat(ntrain, 1, 1)], dim=2)
x_test = torch.cat([x_test.reshape(ntest, train_data_res, 1),
                    grid.repeat(ntest, 1, 1)], dim=2)

train_loader = torch.utils.data.DataLoader(
    torch.utils.data.TensorDataset(x_train, y_train),
    batch_size=20, shuffle=True
)

model = FNO1d(modes=16, width=64).cuda()
print(f"Parameters: {count_params(model)}")

optimizer = torch.optim.Adam(model.parameters(), lr=0.001, weight_decay=1e-4)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=100, gamma=0.5)
myloss = LpLoss(size_average=False)

for ep in range(500):
    model.train()
    for x, y in train_loader:
        x, y = x.cuda(), y.cuda()
        optimizer.zero_grad()
        out = model(x)
        loss = F.mse_loss(out.view(20, -1), y.view(20, -1))
        loss.backward()
        optimizer.step()
    scheduler.step()
```

--------------------------------

### Train POD-DeepONet Model

Source: https://context7.com/lu-group/deeponet-fno/llms.txt

Example usage for training a POD-DeepONet model for Burgers' equation. This snippet demonstrates data preparation, model instantiation, compilation, and training.

```python
# Usage for Burgers' equation
x_train, y_train, x_test, y_test = get_data(1000, 200)
data = dde.data.TripleCartesianProd(x_train, y_train, x_test, y_test)
y_mean, v = pod(y_test)

modes = 32
m = 2 ** 7
net = PODDeepONet(
    v[:, :modes],                  # POD basis (truncated to modes)
    [m, 128, 128, modes],          # Branch: outputs POD coefficients
    None,                          # No trunk (pure POD)
    "tanh",
    "Glorot normal"
)

def output_transform(inputs, outputs):
    return outputs / modes + y_mean

net.apply_output_transform(output_transform)

model = dde.Model(data, net)
model.compile("adam", lr=3e-4, metrics=["mean l2 relative error"])
losshistory, train_state = model.train(epochs=500000, batch_size=None)
```

--------------------------------

### 2D Darcy Flow Training Setup

Source: https://context7.com/lu-group/deeponet-fno/llms.txt

Sets up the data loading, normalization, and model initialization for training a 2D FNO on Darcy flow problems. Requires `utilities3` for data handling and normalization. Ensure the `.mat` file is accessible.

```python
# Training for 2D Darcy flow
s = 29  # Grid resolution
r = (421-1) // (s-1)
ntrain, ntest = 1000, 200

reader = MatReader('piececonst_r421_N1024_smooth1.mat')
x_train = reader.read_field('coeff')[:ntrain,::r,::r][:,:s,:s]
y_train = reader.read_field('sol')[:ntrain,::r,::r][:,:s,:s]

# Normalize inputs and outputs
x_normalizer = UnitGaussianNormalizer(x_train)
x_train = x_normalizer.encode(x_train)
y_normalizer = UnitGaussianNormalizer(y_train)
y_train = y_normalizer.encode(y_train)

# Create 2D grid coordinates
grids = [np.linspace(0, 1, s).reshape(s, 1).astype(np.float32) for _ in range(2)]
grid = np.vstack([xx.ravel() for xx in np.meshgrid(*grids)]).T.reshape(1, s, s, 2)
grid = torch.tensor(grid, dtype=torch.float)
x_train = torch.cat([x_train.reshape(ntrain, s, s, 1), grid.repeat(ntrain, 1, 1, 1)], dim=3)

model = FNO2d(modes1=12, modes2=12, width=32).cuda()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001, weight_decay=1e-4)
```

--------------------------------

### Create DeepONet with CNN branch

Source: https://context7.com/lu-group/deeponet-fno/llms.txt

Initializes a DeepONet model with a CNN branch for the input encoding and a fully connected trunk for the mapping. Output standardization is applied using a StandardScaler.

```python
net = dde.maps.DeepONetCartesianProd(
    [m, branch],                  # CNN branch
    [2, 128, 128, 128, 128],      # Fully connected trunk
    activation,
    "Glorot normal"
)

# Output standardization
scaler = StandardScaler().fit(y_train)
std = np.sqrt(scaler.var_.astype(np.float32))

def output_transform(inputs, outputs):
    return outputs * std + scaler.mean_.astype(np.float32)

net.apply_output_transform(output_transform)

model = dde.Model(data, net)
model.compile(
    tfa.optimizers.AdamW(1e-4, learning_rate=3e-4),
    decay=("inverse time", 1, 1e-4),
    metrics=["mean l2 relative error"],
)
losshistory, train_state = model.train(epochs=100000, batch_size=None)
```

--------------------------------

### Build DeepONet Model with Output Standardization

Source: https://context7.com/lu-group/deeponet-fno/llms.txt

Constructs a DeepONet model using DeepXDE's `DeepONetCartesianProd`. It applies output standardization using `StandardScaler` to normalize the training data's output. The model is compiled with Adam optimizer and inverse time learning rate decay.

```python
m = 2 ** 7  # Branch input dimension

net = dde.maps.DeepONetCartesianProd(
    [m, 128, 128, 128, 128],     # Branch network: input -> hidden layers
    [1, 128, 128, 128],          # Trunk network: location -> hidden layers
    "tanh",                       # Activation function
    "Glorot normal"               # Weight initialization
)

# Apply output standardization transform
scaler = StandardScaler().fit(y_train)
std = np.sqrt(scaler.var_.astype(np.float32))

def output_transform(inputs, outputs):
    return outputs * std + scaler.mean_.astype(np.float32)

net.apply_output_transform(output_transform)

# Create data object and model
data = dde.data.TripleCartesianProd(x_train, y_train, x_test, y_test)
model = dde.Model(data, net)

# Compile with learning rate decay
model.compile(
    "adam",
    lr=0.001,
    metrics=["mean l2 relative error"],
    decay=("inverse time", 100000, 0.5)
)

# Train the model
losshistory, train_state = model.train(epochs=500000, batch_size=None)
```

--------------------------------

### Load Darcy Flow Data for DeepONet

Source: https://context7.com/lu-group/deeponet-fno/llms.txt

Loads and preprocesses Darcy flow data from a .mat file for use with DeepONet. It applies a subsampling rate and enforces zero Dirichlet boundary conditions.

```python
import deepxde as dde
import numpy as np
from deepxde.backend import tf
from sklearn.preprocessing import StandardScaler
import tensorflow_addons as tfa

def get_darcy_data(filename, ndata):
    r = 15  # Subsampling rate: 15 -> 29x29 grid
    s = int(((421 - 1) / r) + 1)

    from scipy import io
    data = io.loadmat(filename)
    x_branch = data["coeff"][:ndata, ::r, ::r].astype(np.float32) * 0.1 - 0.75
    y = data["sol"][:ndata, ::r, ::r].astype(np.float32) * 100

    # Enforce zero Dirichlet BC
    y[:, 0, :] = y[:, -1, :] = y[:, :, 0] = y[:, :, -1] = 0

    grids = [np.linspace(0, 1, s, dtype=np.float32) for _ in range(2)]
    grid = np.vstack([xx.ravel() for xx in np.meshgrid(*grids)]).T

    x_branch = x_branch.reshape(ndata, s * s)
    y = y.reshape(ndata, s * s)
    return (x_branch, grid), y
```

--------------------------------

### Load and Preprocess Data for Burgers' Equation (DeepONet)

Source: https://context7.com/lu-group/deeponet-fno/llms.txt

Loads and preprocesses data for Burgers' equation using DeepXDE. Ensure 'burgers_data_R10.mat' is available in the current directory. This function prepares data for both branch and trunk networks.

```python
import deepxde as dde
import numpy as np
from scipy import io
from sklearn.preprocessing import StandardScaler

# Load and preprocess data for Burgers' equation
def get_data(ntrain, ntest):
    sub_x = 2 ** 6
    sub_y = 2 ** 6
    data = io.loadmat("burgers_data_R10.mat")
    x_data = data["a"][:, ::sub_x].astype(np.float32)
    y_data = data["u"][:, ::sub_y].astype(np.float32)

    x_branch_train = x_data[:ntrain, :]
    y_train = y_data[:ntrain, :]
    x_branch_test = x_data[-ntest:, :]
    y_test = y_data[-ntest:, :]

    grid = np.linspace(0, 1, num=2 ** 13)[::sub_y, None]
    x_train = (x_branch_train, grid)
    x_test = (x_branch_test, grid)
    return x_train, y_train, x_test, y_test
```

--------------------------------

### Advection Equation Solver with DeepONet

Source: https://context7.com/lu-group/deeponet-fno/llms.txt

Solves the advection equation by treating the spatiotemporal domain as the output space. Loads training and test data, defines a DeepONet model with an initial condition encoder branch and a spatiotemporal location encoder trunk, compiles the model, and trains it.

```python
import numpy as np
import deepxde as dde

def get_advection_data(filename):
    nx, nt = 40, 40
    data = np.load(filename)
    x = data["x"].astype(np.float32)
    t = data["t"].astype(np.float32)
    u = data["u"].astype(np.float32)  # N x nt x nx

    u0 = u[:, 0, :]  # Initial conditions: N x nx
    xt = np.vstack((np.ravel(x), np.ravel(t))).T  # Spatiotemporal grid
    u = u.reshape(-1, nt * nx)  # Flatten solution
    return (u0, xt), u

# Load training and test data
x_train, y_train = get_advection_data("train_IC2.npz")
x_test, y_test = get_advection_data("test_IC2.npz")
data = dde.data.TripleCartesianProd(x_train, y_train, x_test, y_test)

nx, nt = 40, 40
net = dde.maps.DeepONetCartesianProd(
    [nx, 512, 512],           # Branch: initial condition encoder
    [2, 512, 512, 512],       # Trunk: (x, t) location encoder
    "relu",
    "Glorot normal"
)

model = dde.Model(data, net)
model.compile(
    "adam",
    lr=1e-3,
    decay=("inverse time", 1, 1e-4),
    metrics=["mean l2 relative error"],
)
losshistory, train_state = model.train(epochs=250000, batch_size=None)

# Predict and save results
y_pred = model.predict(data.test_x)
np.savetxt("y_pred_deeponet.dat", y_pred[0].reshape(nt, nx))
np.savetxt("y_true_deeponet.dat", data.test_y[0].reshape(nt, nx))
np.savetxt("y_error_deeponet.dat", (y_pred[0] - data.test_y[0]).reshape(nt, nx))
```

--------------------------------

### Define Complete 1D FNO Model

Source: https://context7.com/lu-group/deeponet-fno/llms.txt

Constructs the full 1D FNO model, comprising lifting, four Fourier layers, and projection stages. It takes the number of modes and model width as initialization parameters.

```python
class FNO1d(nn.Module):
    """Complete 1D FNO with 4 Fourier layers."""
    def __init__(self, modes, width):
        super(FNO1d, self).__init__()
        self.modes1 = modes
        self.width = width
        self.fc0 = nn.Linear(2, self.width)  # Input: (a(x), x)

        self.conv0 = SpectralConv1d(self.width, self.width, self.modes1)
        self.conv1 = SpectralConv1d(self.width, self.width, self.modes1)
        self.conv2 = SpectralConv1d(self.width, self.width, self.modes1)
        self.conv3 = SpectralConv1d(self.width, self.width, self.modes1)
        self.w0 = nn.Conv1d(self.width, self.width, 1)
        self.w1 = nn.Conv1d(self.width, self.width, 1)
        self.w2 = nn.Conv1d(self.width, self.width, 1)
        self.w3 = nn.Conv1d(self.width, self.width, 1)
        self.fc1 = nn.Linear(self.width, 128)
        self.fc2 = nn.Linear(128, 1)

    def forward(self, x):
        x = self.fc0(x)
        x = x.permute(0, 2, 1)

        # 4 Fourier layers: u' = (W + K)(u)
        for conv, w in [(self.conv0, self.w0), (self.conv1, self.w1),
                        (self.conv2, self.w2), (self.conv3, self.w3)]:
            x1 = conv(x)
            x2 = w(x)
            x = F.relu(x1 + x2) if conv != self.conv3 else x1 + x2

        x = x.permute(0, 2, 1)
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return x
```

--------------------------------

### Load MATLAB Data with MatReader

Source: https://context7.com/lu-group/deeponet-fno/llms.txt

Use MatReader to load data from MATLAB .mat files, supporting both old and HDF5 formats. It can convert data to PyTorch tensors and move them to CUDA if needed.

```python
import torch
import numpy as np
import scipy.io
import h5py

class MatReader(object):
    """Load data from MATLAB .mat files (supports both old and HDF5 formats)."""
    def __init__(self, file_path, to_torch=True, to_cuda=False, to_float=True):
        self.to_torch = to_torch
        self.to_cuda = to_cuda
        self.to_float = to_float
        self.file_path = file_path
        self._load_file()

    def _load_file(self):
        try:
            self.data = scipy.io.loadmat(self.file_path)
            self.old_mat = True
        except:
            self.data = h5py.File(self.file_path)
            self.old_mat = False

    def read_field(self, field):
        x = self.data[field]
        if not self.old_mat:
            x = x[()]
            x = np.transpose(x, axes=range(len(x.shape) - 1, -1, -1))
        if self.to_float:
            x = x.astype(np.float32)
        if self.to_torch:
            x = torch.from_numpy(x)
            if self.to_cuda:
                x = x.cuda()
        return x
```

--------------------------------

### Build CNN Branch Network for DeepONet

Source: https://context7.com/lu-group/deeponet-fno/llms.txt

Constructs a Convolutional Neural Network (CNN) as the branch network for DeepONet. This network is designed to capture spatial patterns in 2D data.

```python
# Build CNN branch network
m = 29 ** 2
activation = "relu"
branch = tf.keras.Sequential([
    tf.keras.layers.InputLayer(input_shape=(m,)),
    tf.keras.layers.Reshape((29, 29, 1)),
    tf.keras.layers.Conv2D(64, (5, 5), strides=2, activation=activation),
    tf.keras.layers.Conv2D(128, (5, 5), strides=2, activation=activation),
    tf.keras.layers.Flatten(),
    tf.keras.layers.Dense(128, activation=activation),
    tf.keras.layers.Dense(128),
])
branch.summary()
```

--------------------------------

### Normalize Data with UnitGaussianNormalizer

Source: https://context7.com/lu-group/deeponet-fno/llms.txt

Normalizes data to have zero mean and unit variance per point. Supports encoding data and decoding normalized data, with an option to move normalization parameters to CUDA.

```python
class UnitGaussianNormalizer(object):
    """Pointwise Gaussian normalization (zero mean, unit variance per point)."""
    def __init__(self, x, eps=0.00001):
        self.mean = torch.mean(x, 0)
        self.std = torch.std(x, 0)
        self.eps = eps

    def encode(self, x):
        return (x - self.mean) / (self.std + self.eps)

    def decode(self, x, sample_idx=None):
        if sample_idx is None:
            return x * (self.std + self.eps) + self.mean
        std = self.std[sample_idx] + self.eps
        mean = self.mean[sample_idx]
        return x * std + mean

    def cuda(self):
        self.mean = self.mean.cuda()
        self.std = self.std.cuda()
```

--------------------------------

### 2D Fourier Neural Operator (FNO) Model

Source: https://context7.com/lu-group/deeponet-fno/llms.txt

The complete 2D FNO model architecture. It uses SpectralConv2d layers for spectral processing and includes linear layers for input/output transformations. Ensure input tensor includes spatial coordinates.

```python
class FNO2d(nn.Module):
    """Complete 2D FNO for Darcy flow problems."""
    def __init__(self, modes1, modes2, width):
        super(FNO2d, self).__init__()
        self.modes1 = modes1
        self.modes2 = modes2
        self.width = width
        self.fc0 = nn.Linear(3, self.width)  # Input: (a(x,y), x, y)

        self.conv0 = SpectralConv2d(self.width, self.width, modes1, modes2)
        self.conv1 = SpectralConv2d(self.width, self.width, modes1, modes2)
        self.conv2 = SpectralConv2d(self.width, self.width, modes1, modes2)
        self.conv3 = SpectralConv2d(self.width, self.width, modes1, modes2)
        self.w0 = nn.Conv1d(self.width, self.width, 1)
        self.w1 = nn.Conv1d(self.width, self.width, 1)
        self.w2 = nn.Conv1d(self.width, self.width, 1)
        self.w3 = nn.Conv1d(self.width, self.width, 1)
        self.fc1 = nn.Linear(self.width, 128)
        self.fc2 = nn.Linear(128, 1)

    def forward(self, x):
        batchsize = x.shape[0]
        size_x, size_y = x.shape[1], x.shape[2]
        x = self.fc0(x)
        x = x.permute(0, 3, 1, 2)

        for conv, w in [(self.conv0, self.w0), (self.conv1, self.w1),
                        (self.conv2, self.w2), (self.conv3, self.w3)]:
            x1 = conv(x)
            x2 = w(x.view(batchsize, self.width, -1)).view(batchsize, self.width, size_x, size_y)
            x = F.relu(x1 + x2) if conv != self.conv3 else x1 + x2

        x = x.permute(0, 2, 3, 1)
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return x
```

--------------------------------

### Implement 1D Spectral Convolution Layer

Source: https://context7.com/lu-group/deeponet-fno/llms.txt

Defines the SpectralConv1d module, which performs a linear transformation in Fourier space. It requires the input tensor, weights, and the number of modes to consider.

```python
import torch
import torch.nn as nn
import torch.nn.functional as F
from utilities3 import MatReader, LpLoss, count_params

class SpectralConv1d(nn.Module):
    """1D Fourier layer: FFT -> linear transform -> inverse FFT."""
    def __init__(self, in_channels, out_channels, modes1):
        super(SpectralConv1d, self).__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.modes1 = modes1
        self.scale = (1 / (in_channels * out_channels))
        self.weights1 = nn.Parameter(
            self.scale * torch.rand(in_channels, out_channels, self.modes1, dtype=torch.cfloat)
        )

    def compl_mul1d(self, input, weights):
        return torch.einsum("bix,iox->box", input, weights)

    def forward(self, x):
        batchsize = x.shape[0]
        x_ft = torch.fft.rfft(x)
        out_ft = torch.zeros(batchsize, self.out_channels, x.size(-1)//2 + 1,
                            device=x.device, dtype=torch.cfloat)
        out_ft[:, :, :self.modes1] = self.compl_mul1d(x_ft[:, :, :self.modes1], self.weights1)
        x = torch.fft.irfft(out_ft, n=x.size(-1))
        return x
```

--------------------------------

### PODDeepONet Class Definition

Source: https://context7.com/lu-group/deeponet-fno/llms.txt

Defines the PODDeepONet neural network architecture, integrating POD basis functions with optional trunk network augmentation. Use this class when implementing POD-enhanced DeepONet models.

```python
class PODDeepONet(dde.maps.NN):
    def __init__(self, pod_basis, layer_sizes_branch, layer_sizes_trunk,
                 activation, kernel_initializer):
        super().__init__()
        self.activation_trunk = dde.maps.activations.get(activation)
        self.pod_basis = tf.convert_to_tensor(pod_basis, dtype=tf.float32)
        self.branch = dde.maps.FNN(
            layer_sizes_branch, activation, kernel_initializer
        )
        self.trunk = None
        if layer_sizes_trunk is not None:
            self.trunk = dde.maps.FNN(
                layer_sizes_trunk, self.activation_trunk, kernel_initializer
            )
            self.b = tf.Variable(tf.zeros(1))

    def call(self, inputs, training=False):
        x_func = inputs[0]
        x_loc = inputs[1]
        x_func = self.branch(x_func)
        if self.trunk is None:
            # Pure POD mode
            x = tf.einsum("bi,ni->bn", x_func, self.pod_basis)
        else:
            # POD + Trunk augmentation
            x_loc = self.activation_trunk(self.trunk(x_loc))
            x = tf.einsum("bi,ni->bn", x_func, tf.concat((self.pod_basis, x_loc), 1))
            x += self.b
        if self._output_transform is not None:
            x = self._output_transform(inputs, x)
        return x
```

--------------------------------

### Calculate Relative Lp Loss

Source: https://context7.com/lu-group/deeponet-fno/llms.txt

Computes the relative Lp loss between two tensors, useful for measuring approximation errors during training and evaluation. Supports averaging and summation of losses.

```python
class LpLoss(object):
    """Relative Lp loss for measuring approximation error."""
    def __init__(self, d=2, p=2, size_average=True, reduction=True):
        self.d = d
        self.p = p
        self.reduction = reduction
        self.size_average = size_average

    def rel(self, x, y):
        num_examples = x.size()[0]
        diff_norms = torch.norm(x.reshape(num_examples, -1) - y.reshape(num_examples, -1), self.p, 1)
        y_norms = torch.norm(y.reshape(num_examples, -1), self.p, 1)
        if self.reduction:
            if self.size_average:
                return torch.mean(diff_norms / y_norms)
            return torch.sum(diff_norms / y_norms)
        return diff_norms / y_norms

    def __call__(self, x, y):
        return self.rel(x, y)
```

--------------------------------

### Compute POD Basis

Source: https://context7.com/lu-group/deeponet-fno/llms.txt

Computes the Proper Orthogonal Decomposition basis from training data. This function is essential for dimensionality reduction and feature extraction in POD-based methods.

```python
import deepxde as dde
import numpy as np
from deepxde.backend import tf

def pod(y):
    """Compute POD basis from training data."""
    n = len(y)
    y_mean = np.mean(y, axis=0)
    y = y - y_mean
    C = 1 / (n - 1) * y.T @ y
    w, v = np.linalg.eigh(C)
    w = np.flip(w)
    v = np.fliplr(v)
    v *= len(y_mean) ** 0.5
    return y_mean, v
```

--------------------------------

### 2D Spectral Convolution Layer

Source: https://context7.com/lu-group/deeponet-fno/llms.txt

Implements a 2D spectral convolution using FFT. It applies a linear transform in the spectral domain. Ensure input data is in the correct format for FFT operations.

```python
import torch
import torch.nn as nn
import torch.nn.functional as F
from utilities3 import MatReader, UnitGaussianNormalizer, LpLoss

class SpectralConv2d(nn.Module):
    """2D Fourier layer: 2D FFT -> linear transform -> inverse 2D FFT."""
    def __init__(self, in_channels, out_channels, modes1, modes2):
        super(SpectralConv2d, self).__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.modes1 = modes1
        self.modes2 = modes2
        self.scale = (1 / (in_channels * out_channels))
        self.weights1 = nn.Parameter(
            self.scale * torch.rand(in_channels, out_channels, modes1, modes2, dtype=torch.cfloat)
        )
        self.weights2 = nn.Parameter(
            self.scale * torch.rand(in_channels, out_channels, modes1, modes2, dtype=torch.cfloat)
        )

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

    def forward(self, x):
        batchsize = x.shape[0]
        x_ft = torch.fft.rfft2(x)
        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] = \
            self.compl_mul2d(x_ft[:, :, :self.modes1, :self.modes2], self.weights1)
        out_ft[:, :, -self.modes1:, :self.modes2] = \
            self.compl_mul2d(x_ft[:, :, -self.modes1:, :self.modes2], self.weights2)
        x = torch.fft.irfft2(out_ft, s=(x.size(-2), x.size(-1)))
        return x
```

=== COMPLETE CONTENT === This response contains all available snippets from this library. No additional content exists. Do not make further requests.