Project: s41467-024-54696-y | Function: ScalarLight.RS_propagation (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:** Driven by human ingenuity and creativity, the discovery of super-resolution techniques, which circumvent the classical diffraction limit of light, represent a leap in optical microscopy. However, the vast space encompassing all possible experimental configurations suggests that some powerful concepts and techniques might have not been discovered yet, and might never be with a human-driven direct design approach. Thus, AI-based exploration techniques could provide enormous benefit, by exploring this space in a fast, unbiased way. We introduce XLuminA, an open-source computational framework developed using JAX, a high-performance computing library in Python. XLuminA offers enhanced computational speed enabled by JAX’s accelerated linear algebra compiler (XLA), just-in-time compilation, and its seamlessly integrated automatic vectorization, automatic differentiation capabilities and GPU compatibility. XLuminA demonstrates a speed-up of 4 orders of magnitude compared to well-established numerical optimization methods. We showcase XLuminA’s potential by re-discovering three foundational experiments in advanced microscopy, and identifying an unseen experimental blueprint featuring sub-diffraction imaging capabilities. This work constitutes an important step in AI-driven scientific discovery of new concepts in optics and advanced microscopy.

**KEY TERMS:** {
  "Leeuwenhoek": "Optical microscopes, in today’s sense, were invented 300 years ago by Antonj van Leeuwenhoek ^3",
  "Interestingly": "Interestingly, the adjoint method can be seen as a special case of auto-differentiation (which we use) ^44",
  "First": "First, using a data-driven learning methodology, we rediscover an optical configuration commonly used to adjust beam and image sizes",
  "Therefore": "Therefore, all the results hereby presented are computed considering a wide variety of experimental errors"
}

**CONTEXT CODE:**         Returns ScalarLight object after applying the triangular mask.

        >> Diffractio-adapted function (https://pypi.org/project/diffractio/) << 
        """
        x0, y0 = r
        Xrot, Yrot = rotate_mask(self.X, self.Y, angle, origin=(x0, y0))
        Y = -m * jnp.abs(Xrot - x0) + y0
        mask = jnp.where((Yrot < Y) & (Yrot > (y0 - height)), 1, 0)
        output = ScalarLight(self.x, self.y, self.wavelength)
        output.field = self.field * mask
        return output
    
    def apply_rectangular_mask(self, center, width, height, angle):
        """
        Apply a square mask of variable size. Can generate rectangles, squares and rotate them to create diamond shapes.

        Parameters:
            center (float, float): Coordinates of the center (in microns).
            width (float): Width of the rectangle (in microns).
            height (float): Height of the rectangle (in microns).
            angle (float): Angle of rotation of the rectangle (in degrees).

        Returns ScalarLight object after applying the rectangular mask.
        """
        x0, y0 = center
        angle = angle * (jnp.pi/180)
        Xrot, Yrot = rotate_mask(self.X, self.Y, angle, center)
        mask = jnp.where((Xrot < (width/2)) & (Xrot > (-width/2)) & (Yrot < (height/2)) & (Yrot > (-height/2)), 1, 0)
        output = ScalarLight(self.x, self.y, self.wavelength)
        output.field = self.field * mask
        return output
    
    def apply_annular_aperture(self, di, do):
        """
        Apply an annular aperture of variable size.

        Parameters:
            di (float): Radius of the inner circle (in microns).
            do (float): Radius of the outer circle (in microns).

        Returns ScalarLight object after applying the annular mask.
        """
        di = di/2
        do = do/2
        stop = jnp.where(((self.X**2 + self.Y**2) / di**2) < 1, 0, 1)
        ring = jnp.where(((self.X**2 + self.Y**2) / do**2) < 1, 1, 0)
        output = ScalarLight(self.x, self.y, self.wavelength)
        output.field = self.field * stop*ring
        return output
    
    def RS_propagation(self, z):
        """
        Rayleigh-Sommerfeld diffraction integral in both, z>0 and z<0, for ScalarLight.
        [Ref 1: F. Shen and A. Wang, Appl. Opt. 45, 1102-1110 (2006)].
        [Ref 2: J. Li, Z. Fan, Y. Fu, Proc. SPIE 4915, (2002)].

        Parameters:
            z (float): Distance to propagate (in microns).

        Returns ScalarLight object after propagation and the quality factor of the algorithm.
        """
        raise NotImplementedError("Function implementation has been masked")
    
    def get_RS_minimum_z(self, n=1, quality_factor=1):
        """
        Given a quality factor, determines the minimum available (trustworthy) distance for RS_propagation(). 
        [Ref 1: Laser Phys. Lett., 10(6), 065004 (2013)].

        Parameters:
            n (float): refraction index of the medium.
            quality_factor (int): Defaults to 1.

        Returns the minimum distance z (in microns) necessary to achieve qualities larger than quality_factor.
        
        >> Diffractio-adapted function (https://pypi.org/project/diffractio/) << 
        """
        # Check sampling 
        range_x = self.x[-1] - self.x[0]
        range_y = self.y[-1] - self.y[0]
        num_x = jnp.size(self.x)
        num_y = jnp.size(self.y)

        dx = range_x / num_x
        dy = range_y / num_y
        # Delta rho 
        dr_real = jnp.sqrt(dx**2 + dy**2)
        # Rho
        rmax = jnp.sqrt(range_x**2 + range_y**2)

        factor = (((quality_factor * dr_real + rmax)**2 - (self.wavelength / n)**2 - rmax**2) / (2 * self.wavelength / n))**2 - rmax**2

        if factor > 0:
            z_min = jnp.sqrt(factor)
        else:
            z_min = 0

        return print("Minimum distance to propagate (in microns):", z_min)
    
    def CZT(self, z, xout=None, yout=None):
        """
        Chirped z-transform propagation - efficient diffraction using the Bluestein method.
        Useful for imaging light in the focal plane: allows high resolution zoom in z-plane. 
        [Ref] Hu, Y., et al. Light Sci Appl 9, 119 (2020).

        Parameters:
            z (float): Propagation distance (in microns).
            xout (jnp.array): Array with the x-positions for the output plane.
            yout (jnp.array): Array with the y-positions for the output plane.

        Returns ScalarLight object after propagation.
        """
        tic = time.perf_counter()
        if xout is None:
            xout = self.x

        if yout is None:
            yout = self.y

        # Define main set of parameters
        nx, ny, dx, dy, Xout, Yout, Dm, fy_1, fy_2, fx_1, fx_2 = build_CZT_grid(z, self.wavelength, self.x, self.y, xout, yout)
        
        # Compute the diffraction integral using Bluestein method
        field_at_z = CZT_jit(self.field, z, self.wavelength, self.k, nx, ny, dx, dy, Xout, Yout, self.X, self.Y, Dm, fy_1, fy_2, fx_1, fx_2)

        # Build ScalarLight object with output field.
        field_out = ScalarLight(xout, yout, self.wavelength)
        field_out.field = field_at_z
        print(f"Time taken to perform one CZT propagation (in seconds): {(time.perf_counter() - tic):.4f}")
        return field_out

def build_grid(x, y):
    """[From RS_propagation]: Returns the grid where the transfer function is defined."""
    nx = len(x)
    ny = len(y)
    dx = x[1] - x[0]
    dy = y[1] - y[0]
    # Build 2N-1 x 2N-1 (X, Y) space:
    x_padded = jnp.pad((x[0] - x[::-1]), (0, jnp.size(x) - 1), 'reflect')
    y_padded = jnp.pad((y[0] - y[::-1]), (0, jnp.size(y) - 1), 'reflect')
    # Convert the right half into positive values:
    I = jnp.ones((1, int(len(x_padded) / 2) + 1))
    II = -jnp.ones((1, int(len(x_padded) / 2)))
    III = jnp.ravel(jnp.concatenate((I, II), 1))
    Xext, Yext = jnp.meshgrid(x_padded * III, y_padded * III)
    return nx, ny, dx, dy, Xext, Yext

@partial(jit, static_argnums=(2, 3, 4, 5, 8))
def RS_propagation_jit(input_field, z, nx, ny, dx, dy, Xext, Yext, k):
    """[From RS_propagation]: JIT function for Equation (10) in [Ref 1]."""
    # input_field is jnp.array of (N, N)


Complete the function implementation replacing "raise NotImplementedError".