Add new get_spectrum and get_duration utils

lasy/utils/laser_utils.py

@@ -40,16 +40,11 @@
 
     envelope = grid.field
 
-    dz = grid.dx[-1] * c
+    dV = get_grid_cell_volume(grid, dim)
 
     if dim == "xyt":
-        dV = grid.dx[0] * grid.dx[1] * dz
         energy = ((dV * epsilon_0 * 0.5) * abs(envelope) ** 2).sum()
-    elif dim == "rt":
-        r = grid.axes[0]
-        dr = grid.dx[0]
-        # 1D array that computes the volume of radial cells
-        dV = np.pi * ((r + 0.5 * dr) ** 2 - (r - 0.5 * dr) ** 2) * dz
+    else:  # dim == "rt":
         energy = (
             dV[np.newaxis, :, np.newaxis]
             * epsilon_0
@@ -211,6 +206,127 @@
     return env, ext
 
 
+def get_spectrum(
+    grid,
+    dim,
+    bins=20,
+    range=None,
+    omega0=None,
+    phase_unwrap_1d=None,
+    is_envelope=True,
+    method="fft",
+):
+    """
+    Get the the frequency spectrum of an envelope.
+
+    Parameters
+    ----------
+    grid : a Grid object.
+        It contains a ndarrays with the field data from which the
+        spectrum is computed, and the associated metadata. The last axis must
+        be the longitudinal dimension.
+
+    dim : string (optional)
+        Dimensionality of the array. Only used if is_envelope is False.
+        Options are:
+
+        - 'xyt': The laser pulse is represented on a 3D grid:
+                 Cartesian (x,y) transversely, and temporal (t) longitudinally.
+        - 'rt' : The laser pulse is represented on a 2D grid:
+                 Cylindrical (r) transversely, and temporal (t) longitudinally.
+
+    bins : int (optional)
+        Number of bins of the spectrum.
+
+    range : list of float (optional)
+        List of two values indicating the minimum and maximum frequency of the
+        spectrum.
+
+    omega0 : scalar
+        Angular frequency at which the envelope is defined.
+
+    phase_unwrap_1d : boolean (optional)
+        Whether the phase unwrapping is done in 1D.
+        This is not recommended, as the unwrapping will not be accurate,
+        but it might be the only practical solution when dim is 'xyt'.
+
+    Returns
+    -------
+    spectrum : ndarray of doubles
+        Array with the spectrum amplitudes.
+
+    omega_spectrum : scalar
+        Array with the frequencies of the spectrum.
+    """
+    if method == "fft":
+        # spectrum = np.fft.fft(grid.field[0, 0]) * grid.dx[-1]
+        spectrum = np.fft.fft(grid.field) * grid.dx[-1]
+        dV = get_grid_cell_volume(grid, dim)
+        if dim == "xyt":
+            spectrum = np.sum(spectrum, axis=(0, 1))
+        else:
+            spectrum = np.sum(spectrum * dV[np.newaxis, :, np.newaxis] / dV[0], axis=1)[
+                0
+            ]
+        # spectrum = np.abs(spectrum[:int(len(spectrum) / 2)])
+        freq = np.fft.fftfreq(
+            spectrum.shape[-1], d=(grid.axes[-1][1] - grid.axes[-1][0])
+        )
+        omega_spectrum = 2 * np.pi * freq
+
+        if is_envelope:
+            omega_spectrum = omega0 - omega_spectrum
+
+        i_sort = np.argsort(omega_spectrum)
+        omega_spectrum = omega_spectrum[i_sort]
+        spectrum = np.abs(spectrum[i_sort])
+
+        i_keep = omega_spectrum >= 0.0
+        omega_spectrum = omega_spectrum[i_keep]
+        spectrum = spectrum[i_keep]
+
+        omega_interp = np.linspace(*range, bins)
+        spectrum = np.interp(omega_interp, omega_spectrum, spectrum)
+        omega_spectrum = omega_interp
+
+        # spectrum = np.fft.fft(grid.field) * grid.dx[-1]
+        # spectrum = np.abs(spectrum[..., :int(spectrum.shape[-1] / 2)])
+        # omega_spectrum = 2 * np.pi / (grid.axes[-1][-1] - grid.axes[-1][0]) * np.arange(spectrum.shape[-1])
+        # if is_envelope:
+        #     omega_spectrum = omega_spectrum
+        # dV = get_grid_cell_volume(grid, dim)
+        # if dim == "xyt":
+        #     spectrum = np.sum(spectrum * dV, axis=(0, 1))
+        # else:
+        #     spectrum = np.sum(spectrum * dV[np.newaxis, :, np.newaxis], axis=1)[0]
+    else:
+        # Get the array of angular frequency.
+        omega, central_omega = get_frequency(
+            grid=grid,
+            dim=dim,
+            is_envelope=True,
+            omega0=omega0,
+            phase_unwrap_1d=phase_unwrap_1d,
+        )
+        # Calculate weights of each frequency (amplitude of the field).
+        dV = get_grid_cell_volume(grid, dim)
+        if dim == "xyt":
+            weights = np.abs(grid.field) * dV
+        else:  # dim == "rt":
+            weights = np.abs(grid.field) * dV[np.newaxis, :, np.newaxis]
+        # Get weighted spectrum.
+        # Neglects the 2 first and last time slices, whose values seems to be
+        # slightly off (maybe due to lower-order derivative at the edges).
+        spectrum, edges = np.histogram(
+            a=np.squeeze(omega)[..., 2:-2],
+            weights=np.squeeze(weights[..., 2:-2]),
+            bins=bins,
+            range=range,
+        )
+        omega_spectrum = edges[1:] - (edges[1] - edges[0]) / 2
+    return spectrum, omega_spectrum
+
+
 def get_frequency(
     grid,
     dim=None,
@@ -313,6 +429,32 @@
     return omega, central_omega
 
 
+def get_duration(grid, dim):
+    """Get envelope duration, measured as RMS.
+
+    Parameters
+    ----------
+    grid : Grid
+        The grid with the envelope to analyze.
+    dim : str
+        Dimensionality of the grid.
+
+    Returns
+    -------
+    float
+        RMS duration of the envelope in seconds.
+    """
+    # Calculate weights of each frequency (amplitude of the field).
+    dV = get_grid_cell_volume(grid, dim)
+    if dim == "xyt":
+        weights = np.abs(grid.field) * dV
+    else:  # dim == "rt":
+        weights = np.abs(grid.field) * dV[np.newaxis, :, np.newaxis]
+    # project weights to longitudinal axes
+    weights = np.sum(weights, axis=(0, 1))
+    return weighted_std(grid.axes[-1], weights)
+
+
 def field_to_vector_potential(grid, omega0):
     """
     Convert envelope from electric field (V/m) to normalized vector potential.
@@ -421,6 +563,53 @@
     return hilbert(grid.field[:, :, ::-1])[:, :, ::-1]
 
 
+def get_grid_cell_volume(grid, dim):
+    """Get the volume of the grid cells.
+
+    Parameters
+    ----------
+    grid : Grid
+        The grid form which to compute the cell volume
+    dim : str
+        Dimensionality of the grid.
+
+    Returns
+    -------
+    float or ndarray
+        A float with the cell volume (if dim=='xyt') or a numpy array with the
+        radial distribution of cell volumes (if dim=='rt').
+    """
+    dz = grid.dx[-1] * c
+    if dim == "xyt":
+        dV = grid.dx[0] * grid.dx[1] * dz
+    else:  # dim == "rt":
+        r = grid.axes[0]
+        dr = grid.dx[0]
+        # 1D array that computes the volume of radial cells
+        dV = np.pi * ((r + 0.5 * dr) ** 2 - (r - 0.5 * dr) ** 2) * dz
+    return dV
+
+
+def weighted_std(values, weights=None):
+    """Calculate the weighted standard deviation of the given values.
+
+    Parameters
+    ----------
+    values: array
+        Contains the values to be analyzed
+
+    weights : array
+        Contains the weights of the values to analyze
+
+    Returns
+    -------
+    A float with the value of the standard deviation
+    """
+    mean_val = np.average(values, weights=weights)
+    std = np.sqrt(np.average((values - mean_val) ** 2, weights=weights))
+    return std
+
+
 def create_grid(array, axes, dim):
     """Create a lasy grid from a numpy array.