Project: s41467-025-59847-3 | Function: compl_mul2d (python)

system: You are an expert Python programmer implementing algorithms from scientific papers.
Complete the masked function based on the algorithm description.
Follow Python best practices and maintain consistent code style.
Return **only** the complete Python function code.

user: Complete this Python function:

**DESCRIPTION:** Implements efficient 2D complex multiplication operation for spectral convolution in Fourier Neural Operators. Uses Einstein summation for tensor contraction to perform channel mixing in Fourier domain for PDE solving.

**KEY TERMS:** {
  "input": "Fourier domain features",
  "weights": "learnable spectral weights",
  "bixy": "batch-input-spatial dimensions",
  "boxy": "batch-output-spatial dimensions"
}

**CONTEXT CODE:** from util_code.util_comm import *
from DONN_Model import layers as layers

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

################################################################
# Fourier layer
################################################################
class SpectralConv2d(nn.Module):
    def __init__(self, in_channels, out_channels, modes1, modes2, config):
        super(SpectralConv2d, self).__init__()

        """
        2D Fourier layer. It does FFT, linear transform, and Inverse FFT.    
        """
        self.config = config
        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, dtype=torch.cfloat))

        ### DONN
        self.wavelength = config['wavelength']
        self.pixel_size = config['pixel_size']
        self.img_size = config['img_size']
        self.distance = config['distance'] * config['distance_factor']
        self.pad = config['pad']  # for phase mask
        self.padsize = config['padsize']  # for input
        self.sys_size = self.img_size + 2 * self.padsize
        self.amp_factor = config['amp_factor']
        self.approx = config['approx']
        self.device = torch.device(config['device'])

        self.fourier_layer1 = DiffractiveClassifier_CoDesign(device=self.device,  det_size=6,
                                       wavelength= self.wavelength, pixel_size= self.pixel_size, sys_size= self.sys_size,
                                       pad=self.pad, padsize=self.padsize, distance=self.distance,
                                       tau=10, amp_factor=self.amp_factor).to(self.device)

        self.fourier_layer2 = DiffractiveClassifier_CoDesign(device=self.device,  det_size=6,
                                       wavelength=self.wavelength, pixel_size=self.pixel_size, sys_size=self.sys_size,
                                       pad=self.pad, padsize=self.padsize, distance=self.distance,
                                       tau=10, amp_factor=self.amp_factor).to(self.device)


    # # Complex multiplication
    def compl_mul2d(self, input, weights):
        # MASKED: Implements efficient 2D complex multiplication operation for spectral convolution in Fourier Neural Operators using Einstein summation implementation
        raise NotImplementedError("This function is masked for testing purposes.")

    def forward(self, x):
        # DONN layer 1
        x_ft = self.fourier_layer1(x)

        # Channel mixing
        out_ft = self.compl_mul2d(x_ft, self.weights1)

        # DONN layer 2
        out_ft = self.fourier_layer2(out_ft)
        return torch.abs(out_ft)**2

################################################################
# DONN layer
################################################################
class DiffractiveClassifier_CoDesign(torch.nn.Module):
    def __init__(self, det_size,
                 device, tau=10, wavelength=5.32e-7, pixel_size=0.000036,
                 batch_norm=False,
                 sys_size=200, pad=1, padsize=150, distance=0.1, precision=256, amp_factor=6,
                 approx="Fresnel"):
        super(DiffractiveClassifier_CoDesign, self).__init__()
        self.amp_factor = amp_factor
        self.tau = tau
        self.size = sys_size
        self.det_size = det_size
        self.distance = distance
        self.wavelength = wavelength
        self.pixel_size = pixel_size
        self.pad = pad
        self.padsize = padsize
        self.approx = approx
        
        self.diffractive_layers_1 = layers.DiffractiveLayerRaw(
                                                            wavelength=self.wavelength, pixel_size=self.pixel_size,
                                                            size=self.size, pad=self.pad, distance=self.distance,
                                                            amplitude_factor=self.amp_factor, mesh_size = 4,
                                                            phase_mod=True)
        self.last_diffraction = layers.DiffractiveLayerRaw(wavelength=self.wavelength,
                                                        pixel_size=self.pixel_size,
                                                        size=self.size, pad=self.pad, distance=self.distance,
                                                        amplitude_factor=self.amp_factor, mesh_size = 4,
                                                        phase_mod=False)

    def forward(self, x, evaluate=False):

        bs, c, height, width = x.shape
        output = torch.zeros_like(x, dtype=torch.complex64)

        # Process each channel individually
        for i in range(c):
            # Select the current channel
            current_channel = x[:, i:i+1, :, :]  # Keeps dimensions (bs, 1, x, y)

            # Process the current channel
            processed_channel = self.diffractive_layers_1(current_channel.squeeze())
            processed_channel = self.last_diffraction(processed_channel)

            # Store the processed channel's result in the output tensor
            output[:, i:i+1, :, :] = processed_channel.unsqueeze(dim = 1)

        return output

###################################'
###  ONE model
class FNO2d(nn.Module):
    def __init__(self, config
                 ):
        super(FNO2d, self).__init__()

        self.config = config
        self.modes1 = self.config['modes']
        self.modes2 = self.config['modes']
        self.width = self.config['width']  # width is 1
        self.ratio = self.config['DONN_ratio']
        self.padding = 9  # pad the domain if input is non-periodic
        self.fc0 = nn.Linear(5, self.width)  # input channel is 3: (a(x, y), x, y)

        self.conv0 = SpectralConv2d(self.width, self.width, self.modes1, self.modes2, self.config)
        self.conv1 = SpectralConv2d(self.width, self.width, self.modes1, self.modes2, self.config)
        self.conv2 = SpectralConv2d(self.width, self.width, self.modes1, self.modes2, self.config)
        self.conv3 = SpectralConv2d(self.width, self.width, self.modes1, self.modes2, self.config)

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

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

    def forward(self, x):
        grid = self.get_grid(x.shape, x.device)
        x = torch.cat((x, grid), dim=-1)
        x = self.fc0(x)
        x = x.permute(0, 3, 1, 2)
        x = F.pad(x, [0, self.padding, 0, self.padding])


Complete the function implementation replacing "raise NotImplementedError".