Utility to denoise laser profile and adding a test for the same

lasy/utils/denoise.py

@@ -0,0 +1,102 @@
+from lasy.profiles.transverse.hermite_gaussian_profile import (
+    HermiteGaussianTransverseProfile,
+)
+from lasy.profiles.transverse.transverse_profile import TransverseProfile
+from lasy.utils.mode_decomposition import hermite_gauss_decomposition
+
+
+def denoise_transverse_hg(
+    transverse_profile,
+    wavelength,
+    resolution=0.2e-6,
+    lo=[-2e-4, -2e-4],
+    hi=[2e-4, 2e-4],
+    n_modes_x=2,
+    n_modes_y=2,
+):
+    """
+    Denoise the transverse profile by decomposing it into a set of Hermite-Gaussian modes.
+
+    The profiles are weighted according to mode coefficients and then added.
+
+    Parameters
+    ----------
+    transverse_profile : class instance
+        An instance of a class or sub-class of TransverseProfile. Or it could be
+        a path to an image file from which intensity data will be read.
+        Defines the transverse envelope of the laser.
+
+    wavelength : float (in meter)
+        Central wavelength at which the Hermite-Gauss beams are to be defined.
+
+    resolution : float
+        The resolution of grid points in x and y that will be used
+        during the decomposition calculation.
+
+    lo, hi : array of floats
+        The lower and upper bounds of the spatial grid on which the
+        decomposition will be performed.
+
+    n_modes_x, n_modes_y : ints
+        The maximum values of `n_x` and `n_y` up to which the
+        expansion will be performed.
+
+    Returns
+    -------
+    transverse_profile_cleaned : class instance
+        Denoised transverse profile after decomposition and recombination.
+
+    waist : float (meter)
+        Beam waist for which the decomposition is calculated.
+        It is computed as the waist for which the weight of order 0 is maximum.
+
+    laser_energy_new : float
+        The total energy of the laser pulse after decomposition.
+    """
+    laser_energy_new = 0
+    if isinstance(transverse_profile, TransverseProfile):
+        pass
+    else:
+        import numpy as np
+        import skimage
+
+        from lasy.profiles.transverse.transverse_profile_from_data import (
+            TransverseProfileFromData,
+        )
+
+        intensity_data = skimage.io.imread(transverse_profile)
+        intensity_scale = np.max(intensity_data)  # Maximum value of the intensity
+        intensity_data[intensity_data < intensity_scale / 100] = 0
+        nx, ny = intensity_data.shape
+        lb = (0, 0)  # Lower bounds in x and y
+        ub = (ny * resolution, nx * resolution)  # Upper bounds in x and y
+
+        # Create the transverse profile. This also centers the data by default
+        transverse_profile = TransverseProfileFromData(
+            intensity_data, [lb[0], lb[1]], [ub[0], ub[1]]
+        )
+
+    # Calculate the decomposition and waist of the laser pulse
+    modeCoeffs, waist = hermite_gauss_decomposition(
+        transverse_profile, wavelength, resolution, lo, hi, n_modes_x, n_modes_y
+    )
+
+    # Denosing the laser profile
+    for i, mode_key in enumerate(list(modeCoeffs)):
+        transverse_profile_temp = HermiteGaussianTransverseProfile(
+            waist, waist, mode_key[0], mode_key[1], wavelength
+        )  # Create a new profile for each mode
+
+        print(f"Mode {i}: {mode_key} with coefficient {modeCoeffs[mode_key]}")
+        laser_energy_new += modeCoeffs[mode_key] ** 2  # Energy fraction of the mode
+
+        if i == 0:  # First mode (0,0)
+            transverse_profile_cleaned = modeCoeffs[mode_key] * transverse_profile_temp
+        else:  # All other modes
+            transverse_profile_cleaned += modeCoeffs[mode_key] * transverse_profile_temp
+
+    # Energy loss due to decomposition
+    energy_loss = 1 - laser_energy_new
+    print(f"Energy loss: {energy_loss * 100:.2f}%")
+
+    return transverse_profile_cleaned, waist, laser_energy_new

tests/test_denoise.py

@@ -0,0 +1,37 @@
+"""Test the implementation of the denoise.
+
+Test checks the implementation of the denoise
+by initializing a super-Gaussian pulse and denoise it. It then
+checks that the error remains positive and less than the predefined
+value.
+"""
+
+import numpy as np
+
+from lasy.profiles.transverse.super_gaussian_profile import (
+    SuperGaussianTransverseProfile,
+)
+from lasy.utils.denoise import denoise_transverse_hg
+
+
+def test_denoise_transverse_hg():
+    # Parameters
+    waist = 20e-6
+    shape_parameter = 3
+    wavelength = 8e-7
+
+    # Define the transverse profile
+    transverse_profile = SuperGaussianTransverseProfile(
+        waist, shape_parameter
+    )  # Super-Gaussian profile
+    transverse_profile_cleaned, waist, laser_energy_new = denoise_transverse_hg(
+        transverse_profile, wavelength
+    )  # Denoised profile
+
+    # Calculate the error
+    x = np.linspace(-5 * waist, 5 * waist, 500)
+    X, Y = np.meshgrid(x, x)
+    prof1 = np.abs(transverse_profile.evaluate(X, Y)) ** 2
+    prof2 = np.abs(transverse_profile_cleaned.evaluate(X, Y)) ** 2
+    error = (prof1 - prof2) / np.max(prof1)
+    assert 0 < np.max(error) < 0.2