Project: s41467-025-59847-3 | Function: DiffractiveLayerRaw.forward (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 Fresnel diffraction propagation for optical neural network layers. Performs forward wave propagation using FFT-based convolution with Fresnel transfer function, applies phase modulation, and simulates optical wave propagation through diffractive elements.

**KEY TERMS:** {
  "h": "Fresnel transfer function",
  "waves": "input optical field",
  "phase": "learnable phase modulation",
  "exp_j_phase": "complex phase mask",
  "distance": "propagation distance"
}

**CONTEXT CODE:** import torch
import numpy as np
from util_code.util_comm import *

class DiffractiveLayerRaw(torch.nn.Module):

    def __init__(self, wavelength=5.32e-7, pixel_size=0.000036, size=60, pad = 0, distance=0.1, name="diffractive_layer_raw",
				amplitude_factor = 6, mesh_size = 4, approx = "Fresnel", phase_mod=True):
        super(DiffractiveLayerRaw, self).__init__()
        self.size = size   # 2D data size
        self.distance = distance   # distance bewteen diffractive layers
        self.pad = pad
        self.ll = pixel_size * (self.size + self.pad*2)   # the length of diffractive layers
        self.wl = wavelength    # wavelength
        self.fi = 1 / self.ll   # frequency interval
        self.wn = 2 * 3.1415926 / self.wl  # wave number
        self.approx = approx
        self.mesh_size = mesh_size
        assert (self.pad > 0), "padding in forward diffraction has to be greater than 0"
        self.pixel_size = pixel_size
        self.ddi = 1 / self.pixel_size
        self.phi = np.fromfunction(
            lambda x, y: np.square((x - ((self.size + self.pad * 2) // 2)) * self.fi) + np.square((y - ((self.size + self.pad * 2) // 2)) * self.fi),
            shape=((self.size + self.pad * 2), (self.size + self.pad * 2)), dtype=np.complex64)

        self.phi2 = np.fromfunction(
            lambda x, y: np.square((x - ((self.size * self.mesh_size + self.pad * 2) // 2)) * self.fi) + np.square(
                (y - ((self.size * self.mesh_size + self.pad * 2)  // 2)) * self.fi),
            shape=((self.size * self.mesh_size + self.pad * 2), (self.size  * self.mesh_size + self.pad * 2)), dtype=np.complex64)

        if self.approx == "Fresnel":
            print("Network is constructed using Fresnel approximation")
            h = np.fft.fftshift(np.exp(1.0j * self.wn * self.distance) * np.exp(-1.0j * self.wl * np.pi * self.distance * self.phi))
            h2 = np.fft.fftshift(np.exp(1.0j * self.wn * self.distance) * np.exp(-1.0j * self.wl * np.pi * self.distance * self.phi2))
            self.h = torch.nn.Parameter(torch.view_as_complex(torch.stack((torch.from_numpy(h.real), torch.from_numpy(h.imag)), dim=-1)), requires_grad=False)
            self.h2 = torch.nn.Parameter(torch.view_as_complex(torch.stack((torch.from_numpy(h2.real), torch.from_numpy(h2.imag)), dim=-1)), requires_grad=False)
        else:
            assert(0), "approximation function %s is not implemented!"

        self.phase = torch.nn.Parameter(torch.from_numpy( 2 * np.pi * torch.nn.init.xavier_uniform_(torch.empty(self.size,self.size)).numpy() ), requires_grad=True)
        self.register_parameter(name, self.phase)
        self.phase_model = phase_mod
        self.amplitude_factor = amplitude_factor

    def forward(self, waves):
        # MASKED: Implements Fresnel diffraction propagation for optical neural network layers using FFT-based convolution with phase modulation implementation
        raise NotImplementedError("This function is masked for testing purposes.")

class DiffractiveLayerRaw_mag(torch.nn.Module):

    def __init__(self, wavelength=5.32e-7, pixel_size=0.000036, size=60, pad = 0, distance=0.1, name="diffractive_layer_raw",
				amplitude_factor = 6, mesh_size = 4, approx = "Fresnel", phase_mod=True):
        super(DiffractiveLayerRaw_mag, self).__init__()
        self.size = size  
        self.height = size # 2D data height
        self.width = 42 # 2D data width
        self.distance = distance  # distance bewteen diffractive layers
        self.pad = pad
        self.ll = pixel_size * (self.size + self.pad*2)  # the length of diffractive layers
        self.wl = wavelength    # wavelength
        self.fi = 1 / self.ll   # frequency interval
        self.wn = 2 * 3.1415926 / self.wl  # wave number
        self.approx = approx
        self.mesh_size = mesh_size
        assert (self.pad > 0), "padding in forward diffraction has to be greater than 0"
        self.pixel_size = pixel_size
        self.ddi = 1 / self.pixel_size

        self.phi = np.fromfunction(
            lambda x, y: np.square((x - ((self.height + self.pad * 2) // 2)) * self.fi) +
                         np.square((y - ((self.width + self.pad * 2) // 2)) * self.fi),
            shape=(self.height + self.pad * 2, self.width + self.pad * 2), dtype=np.complex64)

        self.phi2 = np.fromfunction(
            lambda x, y: np.square((x - ((self.size * self.mesh_size + self.pad * 2) // 2)) * self.fi) + np.square(
                (y - ((self.size * self.mesh_size + self.pad * 2)  // 2)) * self.fi),
            shape=((self.size * self.mesh_size + self.pad * 2), (self.size  * self.mesh_size + self.pad * 2)), dtype=np.complex64)

        if self.approx == "Fresnel":
            print("Network is constructed using Fresnel approximation")
            h = np.fft.fftshift(np.exp(1.0j * self.wn * self.distance) * np.exp(-1.0j * self.wl * np.pi * self.distance * self.phi))
            h2 = np.fft.fftshift(np.exp(1.0j * self.wn * self.distance) * np.exp(-1.0j * self.wl * np.pi * self.distance * self.phi2))
            self.h = torch.nn.Parameter(torch.view_as_complex(torch.stack((torch.from_numpy(h.real), torch.from_numpy(h.imag)), dim=-1)), requires_grad=False)
            self.h2 = torch.nn.Parameter(torch.view_as_complex(torch.stack((torch.from_numpy(h2.real), torch.from_numpy(h2.imag)), dim=-1)), requires_grad=False)

        else:
            assert(0), "approximation function %s is not implemented!"

        self.phase = torch.nn.Parameter(torch.from_numpy( 2 * np.pi * torch.nn.init.xavier_uniform_(torch.empty(self.height,self.width)).numpy() ), requires_grad=True)
        self.register_parameter(name, self.phase)
        self.phase_model = phase_mod
        self.amplitude_factor = amplitude_factor

    def forward(self, waves):
        # MASKED: Implements Fresnel diffraction propagation for optical neural network layers using FFT-based convolution with phase modulation implementation
        raise NotImplementedError("This function is masked for testing purposes.")






Complete the function implementation replacing "raise NotImplementedError".