some cleanup

src/scgenerator/physics/pulse.py

@@ -366,27 +366,27 @@ def energy_to_mean_power(energy: float, repetition_rate: float) -> float:
     return energy * repetition_rate
 
 
-def soliton_num_to_peak_power(soliton_num, beta2, gamma, t0):
+def soliton_num_to_peak_power(soliton_num: float, beta2: float, gamma: float, t0):
     return soliton_num**2 * abs(beta2) / (gamma * t0**2)
 
 
-def soliton_num_to_t0(soliton_num, beta2, gamma, peak_power):
+def soliton_num_to_t0(soliton_num: float, beta2: float, gamma: float, peak_power: float):
     return np.sqrt(soliton_num**2 * abs(beta2) / (peak_power * gamma))
 
 
-def soliton_num(dispersion_length, nonlinear_length):
+def soliton_num(dispersion_length: float, nonlinear_length: float):
     return np.sqrt(dispersion_length / nonlinear_length)
 
 
-def dispersion_length(t0, beta2):
+def dispersion_length(t0: float, beta2: float):
     return t0**2 / abs(beta2)
 
 
-def nonlinear_length(peak_power, gamma):
+def nonlinear_length(peak_power: float, gamma: float):
     return 1 / (gamma * peak_power)
 
 
-def soliton_length(dispersion_length):
+def soliton_length(dispersion_length: float):
     return pi / 2 * dispersion_length
 
 
@@ -472,45 +472,47 @@ def correct_wavelength(init_wavelength: float, w_c: np.ndarray, field_0: np.ndar
     return units.m.inv(units.m(init_wavelength) - delta_w)
 
 
-def E0_to_P0(E0, t0, shape):
+def E0_to_P0(E0: float, t0: float, shape: str):
     """convert an initial total pulse energy to a pulse peak peak_power"""
     return E0 / (t0 * P0T0_to_E0_fac[shape])
 
 
-def P0_to_E0(P0, t0, shape):
+def P0_to_E0(P0: float, t0: float, shape: str):
     """converts initial peak peak_power to pulse energy"""
     return P0 * t0 * P0T0_to_E0_fac[shape]
 
 
-def sech_pulse(t, t0, P0, offset=0):
+def sech_pulse(t: np.ndarray, t0: float, P0: float):
-    arg = (t - offset) / t0
+    arg = t / t0
     ind = (arg < 700) & (arg > -700)
     out = np.zeros_like(t)
     out[ind] = np.sqrt(P0) / np.cosh(arg[ind])
     return out
 
 
-def gaussian_pulse(t, t0, P0, offset=0):
+def gaussian_pulse(t: np.ndarray, t0: float, P0: float):
-    return np.sqrt(P0) * np.exp(-(((t - offset) / t0) ** 2))
+    return np.sqrt(P0) * np.exp(-((t / t0) ** 2))
 
 
-def photon_number(spec2, w, dw, gamma) -> float:
+def photon_number(spec2: np.ndarray, w: np.ndarray, dw: float, gamma: float) -> float:
     return np.sum(1 / gamma * spec2 / w * dw)
 
 
-def photon_number_with_loss(spec2, w, dw, gamma, alpha, h) -> float:
+def photon_number_with_loss(
+    spec2: np.ndarray, w: np.ndarray, dw: float, gamma: float, alpha: float, h: float
+) -> float:
     return np.sum(1 / gamma * spec2 / w * dw) - h * np.sum(alpha / gamma * spec2 / w * dw)
 
 
-def pulse_energy(spec2, dw) -> float:
+def pulse_energy(spec2: np.ndarray, dw: float) -> float:
     return np.sum(spec2 * dw)
 
 
-def pulse_energy_with_loss(spec2, dw, alpha, h) -> float:
+def pulse_energy_with_loss(spec2: np.ndarray, dw: float, alpha: float, h: float) -> float:
     return np.sum(spec2 * dw) - h * np.sum(alpha * spec2 * dw)
 
 
-def technical_noise(rms_noise, noise_correlation=-0.4):
+def technical_noise(rms_noise: float, noise_correlation: float = -0.4):
     """
     To implement technical noise as described in Grenier2019, we need to know the
     noise properties of the laser, summarized into the RMS amplitude noise
@@ -584,15 +586,16 @@ def C_to_A(C: np.ndarray, effective_area_arr: np.ndarray) -> np.ndarray:
     return (effective_area_arr / effective_area_arr[0]) ** (1 / 4) * C
 
 
-def mean_phase(spectra):
+def mean_phase(spectra: np.ndarray, axis: int = 0):
     """
     computes the mean phase of spectra
 
     Parameters
     ----------
-    spectra : 2D array
+    spectra : np.ndarray, shape (..., n,..., nt)
-        The mean is taken on the 0th axis. This means the array has to be of shape (n, nt)
+        complex spectra from which to compute the mean phase
-
+    axis : int, optional
+        axis on which to take the mean, by default 0
     Returns
     ----------
     mean_phase : 1D array of shape (len(spectra[0]))
@@ -609,33 +612,34 @@ def mean_phase(spectra):
 
     total_phase = np.sum(
         spectra / np.abs(spectra),
-        axis=0,
+        axis=axis,
         where=spectra != 0,
         out=np.zeros(len(spectra[0]), dtype="complex"),
     )
     return (total_phase) / np.abs(total_phase)
 
 
-def flatten_phase(spectra):
+def flatten_phase(spectra: np.ndarray):
     """
     takes the mean phase out of an array of complex numbers
 
     Parameters
     ----------
-    spectra : 2D array of shape (n, nt)
+    spectra : np.ndarray, shape (..., nt)
-        spectra arranged in the same fashion as in `scgenerator.physics.pulse.mean_phase`
+        complex spectra from which to remove the mean phase
 
     Returns
     ----------
-    output : array of same dimensions and amplitude, but with a flattened phase
+    np.ndarray, shape (..., nt)
+        array of same dimensions and amplitude, but with a flattened phase
     """
-    mean_theta = mean_phase(spectra)
+    mean_theta = mean_phase(np.atleast_2d(spectra), axis=0)
     tiled = np.tile(mean_theta, (len(spectra), 1))
     output = spectra * np.conj(tiled)
     return output
 
 
-def compress_pulse(spectra):
+def compress_pulse(spectra: np.ndarray):
     """
     given some complex spectrum, returns the compressed pulse in the time domain
     Parameters
@@ -657,7 +661,7 @@ def compress_pulse(spectra):
         return fftshift(ifft(flatten_phase(spectra)), axes=1)
 
 
-def ideal_compressed_pulse(spectra):
+def ideal_compressed_pulse(spectra: np.ndarray):
     """
     returns the ideal compressed pulse assuming flat phase
     Parameters
@@ -711,7 +715,7 @@ def spectrogram(
     return spec
 
 
-def g12(values):
+def g12(values: np.ndarray):
     """
     computes the first order coherence function of a ensemble of values
 
@@ -739,7 +743,7 @@ def g12(values):
     return np.abs(corr)
 
 
-def avg_g12(values):
+def avg_g12(values: np.ndarray):
     """
     computes the average of the coherence function weighted by amplitude of spectrum
 
@@ -760,7 +764,7 @@ def avg_g12(values):
     return np.sum(coherence * avg_values) / np.sum(avg_values)
 
 
-def fwhm_ind(values, mam=None):
+def fwhm_ind(values: np.ndarray, mam=None):
     """
     returns the indices where values is bigger than half its maximum
 
@@ -788,7 +792,7 @@ def fwhm_ind(values, mam=None):
     return left_ind - 1, right_ind + 1
 
 
-def peak_ind(values, mam=None):
+def peak_ind(values: np.ndarray, mam: tuple[int, int] | None = None):
     """
     returns the indices that encapsulate the entire peak
     Parameters