Requested changes implemented

lasy/utils/denoise.py

@@ -5,14 +5,9 @@
 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,
+def hg_reconstruction(
+    transverse_profile, wavelength, resolution, lo, hi,
+    n_modes_x=10, n_modes_y=10,
 ):
     """
     Denoise the transverse profile by decomposing it into a set of Hermite-Gaussian modes.
@@ -21,61 +16,37 @@ def denoise_transverse_hg(
 
     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.
+    transverse_profile : TransverseProfile object
+        An instance of a class or sub-class of TransverseProfile.
         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.
+        The number of grid points in x and y 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.
+        decomposition is 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.
+    n_modes_x, n_modes_y : ints (optional)
+        The maximum values of `n_x` and `n_y` up to which the expansion is performed.
 
     Returns
     -------
     transverse_profile_cleaned : class instance
         Denoised transverse profile after decomposition and recombination.
 
-    waist : float (meter)
+    waist : array of floats (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
+    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
-        from PIL import Image
-
-        from lasy.profiles.transverse.transverse_profile_from_data import (
-            TransverseProfileFromData,
-        )
-
-        img = Image.open(transverse_profile)
-        intensity_data = np.array(img)
-        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]]
-        )
+    energy_new = 0
+    assert isinstance(transverse_profile, TransverseProfile)
 
     # Calculate the decomposition and waist of the laser pulse
     modeCoeffs, waist = hermite_gauss_decomposition(
@@ -85,19 +56,18 @@ def denoise_transverse_hg(
     # 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
+            waist[0], waist[1], 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
+        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
+    energy_loss = 1 - energy_new
     print(f"Energy loss: {energy_loss * 100:.2f}%")
 
-    return transverse_profile_cleaned, waist, laser_energy_new
+    return transverse_profile_cleaned, waist, energy_new

tests/test_denoise.py

@@ -1,37 +1,37 @@
-"""Test the implementation of the denoise.
+"""Test the implementation of the HG recostruction.
 
-Test checks the implementation of the denoise
-by initializing a super-Gaussian pulse and denoise it. It then
+Test checks the implementation of the HG reconstruction. It does so
+by initializing a super-Gaussian pulse and denoising 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,
+    GaussianTransverseProfile,
 )
-from lasy.utils.denoise import denoise_transverse_hg
+from lasy.utils.denoise import hg_reconstruction
 
 
-def test_denoise_transverse_hg():
+def test_denoise_hg_reconstruction():
     # Parameters
     waist = 20e-6
     shape_parameter = 3
     wavelength = 8e-7
 
     # Define the transverse profile
-    transverse_profile = SuperGaussianTransverseProfile(
+    transverse_profile = GaussianTransverseProfile(
         waist, shape_parameter
     )  # Super-Gaussian profile
-    transverse_profile_cleaned, waist, laser_energy_new = denoise_transverse_hg(
+    transverse_profile_cleaned, waist, laser_energy_new = hg_reconstruction(
         transverse_profile, wavelength
     )  # Denoised profile
 
     # Calculate the error
-    x = np.linspace(-5 * waist, 5 * waist, 500)
+    x = np.linspace(-5 * waist[0], 5 * waist[0], 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
+    assert 0 < np.max(error) < 0.2